import numpy as np
from scipy.ndimage import measurements

from src.processing.imageprocessing import rgb2gray

from matplotlib import pyplot as plt

from src.utils.cmap_generator import rand_cmap


def find_lines(image, draw_steps=False):
    gray, binary, magnitude = preparation(image)
    #backtrack = load_numpy("result/backtrack.npz")
    #if backtrack is None:
    energy, backtrack = minimum_seam(binary, magnitude)
        #save_numpy("result/backtrack.npz", backtrack)
    seams = calculate_seams(backtrack)
    labeled, ncomponents = group_empty_boxes(seams)
    lines = generate_lines(labeled, ncomponents, gray, draw_steps)
    return filter_lines(lines)


def preparation(image):
    gray = rgb2gray(image)
    cnt, vals = np.histogram(gray, 256)
    threshold = get_threshold(cnt)/256*0.96
    binary = (gray > threshold).astype(np.int_)
    magnitude = np.ones_like(binary)-binary
    imin = 0
    while np.sum(binary[:, imin]) / binary.shape[0] < 0.6:
        imin += 1
    imax = binary.shape[1]
    while np.sum(binary[:, imax-1]) / binary.shape[0] < 0.6:
        imax -= 1
    jmin = 0
    while np.sum(binary[jmin]) / binary.shape[1] < 0.6:
        jmin += 1
    jmax = binary.shape[0]
    while np.sum(binary[jmax-1]) / binary.shape[1] < 0.6:
        jmax -= 1
    return gray[jmin:jmax, imin:imax], binary[jmin:jmax, imin:imax], magnitude[jmin:jmax, imin:imax]


def get_threshold(hist, thresh=None):
    # ISO data algorithm
    # https://felixniklas.com/imageprocessing/binarization
    if thresh is None:
        thresh = hist.shape[0] // 2
    m1 = median(hist, 0, thresh)
    m2 = median(hist, thresh, hist.shape[0])
    tk = int((m1 + m2) / 2)
    if thresh == tk:
        return np.round(tk)
    else:
        return get_threshold(hist, tk)


def median(values, start, stop):
    p, x = 0, 0
    for idx, val in enumerate(values[start:stop], start=start):
        p += val
        x += idx*val
    if p == 0:
        return start if start != 0 else stop+1
    return x/p


def minimum_seam(img, energy_map):
    r, c = img.shape
    M = energy_map.copy()
    backtrack = np.zeros_like(M, dtype=np.int)
    for j in range(1, c):
        for i in range(0, r):
            # Handle the top edge of the image, to ensure we don't index -1
            if i == 0:
                idx = np.argmin(M[i:i+2, j-1])
                backtrack[i, j] = idx + i
                min_energy = M[idx+i, j-1]
                if idx > 0:
                    min_energy += 1
            else:
                m = M[i-1:i+2, j-1]
                idx = (np.argmin(np.roll(m, 2)) - 2) % len(m)
                backtrack[i, j] = idx + i - 1
                min_energy = M[idx+i-1, j-1]
                if idx != 1:
                    min_energy += 1
            M[i, j] += min_energy
    return M, backtrack


def calculate_seams(links):
    h, w = links.shape
    seams = np.zeros_like(links)
    seams[:, w-1] = 1
    for x in range(w-1, 0, -1):
        for y in range(h):
            seams[links[y, x], x-1] += seams[y, x]
    return seams


def group_empty_boxes(seams):
    clouds = 1 - np.minimum(seams, 1)
    structure = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.int)
    labeled, ncomponents = measurements.label(clouds, structure)
    return labeled, ncomponents


def generate_lines(labeled, ncomponents, gray, draw_steps=False):
    if draw_steps:
        plt.imshow(labeled, cmap=rand_cmap(ncomponents+1, "hard"))
        plt.title("Found line groups")
        plt.show()

    groups = np.copy(labeled)
    group_id = 0
    entries = []
    pixelgroup = None
    in_top_mode = True
    indices = np.indices(labeled.shape).T[:, :, [1, 0]]
    indices = np.swapaxes(indices, 0, 1)

    def submit_entry():
        nonlocal pixelgroup, group_id, entries
        group_id += 1
        minps = np.min(pixelgroup, axis=0)
        maxps = np.max(pixelgroup, axis=0)
        entry = gray[minps[1]:maxps[1] + 1, minps[0]:maxps[0] + 1]
        if pixelgroup.shape[0] < 20 or maxps[1] - minps[1] < 5 or maxps[0] - minps[0] < 5:
            return
        if pixelgroup.shape[0] / entry.shape[0] / entry.shape[1] < 0.10:
            return
        pixelgroup = np.subtract(pixelgroup, minps)
        white = np.ones_like(entry) * np.max(entry)
        for w in range(entry.shape[1]):
            ys = pixelgroup[pixelgroup[:, 0] == w][:, 1]
            if ys.size > 0:
                ymin = np.min(ys)
                ymax = np.max(ys)
                white[ymin:ymax, w] = entry[ymin:ymax, w]
        #white[pixelgroup[:, 1], pixelgroup[:, 0]] = entry[pixelgroup[:, 1], pixelgroup[:, 0]]
        entries.append(white)
        pixelgroup = None

    for label in range(1, ncomponents+1):
        pixel = indices[labeled == label]
        minp = np.min(pixel, axis=0)
        maxp = np.max(pixel, axis=0)
        right_pixel = pixel[pixel[:, 0] == maxp[0]][:, 1]
        second_pixel = pixel[pixel[:, 0] == maxp[0]-1][:, 1]
        if second_pixel.size > 0 and min(right_pixel) > min(second_pixel):
            # up
            if not in_top_mode:
                submit_entry()
            in_top_mode = True
        if second_pixel.size > 0 and max(right_pixel) < max(second_pixel):
            # down
            in_top_mode = False
        if pixelgroup is not None and minp[1] - 5 > np.max(pixelgroup, axis=0)[1]:
            submit_entry()
        groups[labeled == label] = group_id+1
        if pixelgroup is None:
            pixelgroup = pixel[:, :]
        else:
            pixelgroup = np.concatenate((pixelgroup, pixel))
    submit_entry()
    if draw_steps:
        plt.imshow(groups, cmap=rand_cmap(group_id+1, "hard"))
        plt.title("Combined and valid groups")
        plt.show()
    return entries

def filter_lines(lines):
    filtered = []
    for line in lines:
        cnt, vals = np.histogram(line, 256)
        threshold = get_threshold(cnt)/256*0.96
        binary = (line > threshold).astype(np.int_)
        labeled, ncomponents = group_empty_boxes(binary)
        count = 0
        for label in range(1, ncomponents+1):
            pixcnt = np.sum(labeled == label)
            if pixcnt > 10:
                count += 1
        if count > 2:
            filtered.append(line)
    return filtered