from collections import defaultdict

from skimage.transform import resize
from scipy.ndimage import gaussian_filter
from matplotlib import pyplot as plt

from src.processing.imageprocessing import rgb2gray_value
from scipy import signal
import numpy as np
from PIL import Image

class Line:
    def __init__(self, intercept=0, slope=0, points=None):
        self.intercept = intercept
        self.slope = slope
        self.points = [] if points is None else points
        self.splits = []

    def __str__(self):
        return "m="+str(self.slope)+";n="+str(self.intercept)+";split="+str(self.splits)

def cut_receipt(image):
    image = image[:, :, :3]
    gray = rgb2gray_value(image)
    gray = resize(gray, (500, int(gray.shape[1] / gray.shape[0] * 500)), mode="reflect")
    scale = image.shape[0] / gray.shape[0]
    print(scale)
    gray = gaussian_filter(gray, sigma=max(3, scale*2-8))
    #gauss = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype='float')/16
    #gray = signal.convolve2d(gray, gauss, boundary='symm', mode='same')
    sobely = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype='float')
    sobely = np.array([[0, 2, 0], [0, 0, 0], [0, -2, 0]], dtype='float')
    gray_y = signal.convolve2d(gray, sobely, boundary='symm', mode='same')
    sobelx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype='float')
    sobelx = np.array([[0, 0, 0], [-2, 0, 2], [0, 0, 0]], dtype='float')
    gray_x = signal.convolve2d(gray, sobelx, boundary='symm', mode='same')
    grad_strength = np.sqrt(np.square(gray_y)+np.square(gray_x))
    grad_angle = np.arctan(np.true_divide(gray_y, gray_x, where=gray_x != 0))
    angles = np.copy(grad_angle)
    grad_angle = np.round(grad_angle / np.pi * 4) + 2
    grad_angle = np.array(grad_angle, dtype="uint8")
    grad_angle[grad_angle == 4] = 0
    h, w = grad_strength.shape
    #plt.imshow(grad_strength, cmap="gray")
    #plt.show()

    # Canny
    CANNY_MIN = 0.01
    CANNY_MAX = 0.002
    strenghts = np.zeros_like(grad_strength)
    for y in range(1, h-1):
        for x in range(1, w-1):
            if grad_strength[y, x] < CANNY_MIN:
                continue
            if grad_angle[y, x] == 0:
                if grad_strength[y-1, x] < grad_strength[y, x] and grad_strength[y+1, x] < grad_strength[y, x]:
                    strenghts[y, x] = grad_strength[y, x]
            elif grad_angle[y, x] == 1:
                if grad_strength[y-1, x-1] < grad_strength[y, x] and grad_strength[y+1, x+1] < grad_strength[y, x]:
                    strenghts[y, x] = grad_strength[y, x]
            elif grad_angle[y, x] == 2:
                if grad_strength[y, x-1] < grad_strength[y, x] and grad_strength[y, x+1] < grad_strength[y, x]:
                    strenghts[y, x] = grad_strength[y, x]
            elif grad_angle[y, x] == 3:
                if grad_strength[y-1, x+1] < grad_strength[y, x] and grad_strength[y+1, x-1] < grad_strength[y, x]:
                    strenghts[y, x] = grad_strength[y, x]

    # Hough-Lines
    theta_res = 5
    width_res = 5
    theta_angle = 180 // theta_res
    mid = int(np.round(np.sqrt((h//width_res)**2+(w//width_res)**2)))
    hough = np.zeros((theta_angle, mid*2+2))
    references = defaultdict(list)
    for y in range(1, h-1):
        for x in range(1, w-1):
            if strenghts[y, x] > 0:
                for theta in range(theta_angle):
                    t = theta * np.pi / theta_angle
                    q = (x * np.cos(t) + y * np.sin(t))/width_res + mid
                    #print(theta*theta_res, y//width_res, x//width_res, q - mid, t)
                    hough[theta, int(q)] += 1#strenghts[y, x]
                    #print(theta, int(q), 175//theta_res, -120//width_res + mid + 1)
                    references[int(q) * theta_angle + theta].append([y, x])
    lines = np.unravel_index(np.argsort(hough.ravel())[-100:][::-1], hough.shape)
    lines = np.array(lines).T
    results = []
    COVER_Y = 10
    COVER_X = 15
    for line in lines:
        if hough[line[0], line[1]] > 0:
            results.append([line[0]*theta_res, (line[1]-mid)*width_res])
            hough[max(0, line[0]-COVER_Y):line[0]+COVER_Y, max(0, line[1]-COVER_X):line[1]+COVER_X] = 0
            y1, x1, y2, x2 = None, None, None, None
            if line[0]-COVER_Y < 0:
                y1 = hough.shape[0]+line[0]-COVER_Y
            if line[1]-COVER_X < 0:
                x1 = hough.shape[1]+line[1]-COVER_X
            if line[0]+COVER_Y > hough.shape[0]:
                y2 = line[0]+COVER_Y-hough.shape[0]
            if line[1]+COVER_X > hough.shape[1]:
                x2 = line[1]+COVER_X-hough.shape[1]
            if any(x is not None for x in [y1, x1, y2, x2]):
                ty1 = y1 if y1 is not None else (0 if y2 is not None else line[0]-COVER_Y)
                tx1 = x1 if x1 is not None else (0 if x2 is not None else line[1]-COVER_X)
                ty2 = y2 if y2 is not None else (hough.shape[0] if y1 is not None else line[0]+COVER_Y)
                tx2 = x2 if x2 is not None else (hough.shape[1] if x1 is not None else line[1]+COVER_X)
                hough[ty1:ty2, tx1:tx2] = 0
            if len(results) > 5:
                break
    results = np.array(results)
    #print(results)

    # draw image (removable)
    draw_image = np.copy(image)
    RED_WIDTH = 10
    GREEN_WIDTH = 5
    for result in results:
        refs = references[(result[1]/width_res+mid) * theta_angle + result[0]//theta_res]
        for ref in refs:
            xa = int(scale * ref[1])
            ya = int(scale * ref[0])
            draw_image[max(0, ya - RED_WIDTH):ya + RED_WIDTH, max(0, xa - RED_WIDTH):xa + RED_WIDTH] = np.array([0, 255, 0])
        if result[0] == 0:
            x = int(result[1] * scale)
            for y in range(image.shape[0]):
                draw_image[max(0, y-GREEN_WIDTH):y+GREEN_WIDTH, max(0, x-GREEN_WIDTH):x+GREEN_WIDTH] = np.array([255, 0, 0])
        else:
            angle = (90 - result[0]) / 180 * np.pi
            m = np.sin(angle) / np.cos(angle)
            n = result[1] / np.cos(angle) * scale
            for x in range(image.shape[1]):
                y = int(n - x * m)
                if 0 < y < image.shape[0]:
                    draw_image[max(0, y - GREEN_WIDTH):y + GREEN_WIDTH, max(0, x - GREEN_WIDTH):x + GREEN_WIDTH] = np.array([255, 0, 0])
    #plt.imshow(draw_image)
    #plt.show()

    # convert to original image pixel
    scale = image.shape[0] / gray.shape[0]
    lines = []
    for result in results:
        points = []
        refs = references[(result[1]/width_res+mid) * theta_angle + result[0]//theta_res]
        for ref in refs:
            xa = int(scale * ref[1])
            ya = int(scale * ref[0])
            points.append([ya, xa])
        if result[0] == 0:
            x = int(result[1] * scale)
            lines.append(Line(intercept=x, slope=None, points=points))
        else:
            angle = (90 - result[0]) / 180 * np.pi
            m = np.sin(angle) / -np.cos(angle)
            n = result[1] / np.cos(angle) * scale
            lines.append(Line(intercept=n, slope=m, points=points))

    # split segments
    for idx1, line1 in enumerate(lines):
        for idx2, line2 in enumerate(lines[idx1+1:], idx1+1):
            if line1.slope == line2.slope:
                continue
            elif line1.slope is None:
                x = line1.intercept
                y = line2.intercept + line2.slope * line1.intercept
                if 0 < y < image.shape[0]:
                    line1.splits.append((int(y), int(x), idx2))
                    line2.splits.append((int(y), int(x), idx1))
            elif line2.slope is None:
                x = line2.intercept
                y = line1.intercept + line1.slope * line2.intercept
                if 0 < y < image.shape[0]:
                    line1.splits.append((int(y), int(x), idx2))
                    line2.splits.append((int(y), int(x), idx1))
            else:
                x = (line2.intercept - line1.intercept) / (line1.slope - line2.slope)
                y = line1.intercept + line1.slope * x
                #print(x, y)
                if 0 < x < image.shape[1] and 0 < y < image.shape[0]:
                    line1.splits.append((int(y), int(x), idx2))
                    line2.splits.append((int(y), int(x), idx1))
        if line1.slope is None:
            line1.splits.append((0, int(line1.intercept), None))
            line1.splits.append((image.shape[0], int(line1.intercept), None))
        elif line1.slope == 0:
            line1.splits.append((int(line1.intercept), 0, None))
            line1.splits.append((int(line1.intercept), image.shape[1], None))
        else:
            y = min(max(0, line1.intercept), image.shape[0])
            x = (y - line1.intercept) / line1.slope
            line1.splits.append((int(y), int(x), None))
            y = min(max(0, line1.intercept+line1.slope*image.shape[1]), image.shape[0])
            x = (y - line1.intercept) / line1.slope
            line1.splits.append((int(y), int(x), None))
        if line1.slope is not None:
            line1.splits = sorted(line1.splits, key=lambda x: (x[1], x[0]))
        else:
            line1.splits = sorted(line1.splits, key=lambda x: (x[0], x[1]))
        #print(line1)

    # find important segments
    for line in lines:
        counts = np.zeros((len(line.splits)-1, ))
        for point in line.points:
            if line.slope is None:
                x = point[0]
            else:
                a, b = point[1], point[0]
                m, n = line.slope, line.intercept
                x = (a+b*m-m*n) / (m**2 + 1)
            for i in range(len(line.splits)-1):
                if line.slope is not None:
                    lower = line.splits[i][1]
                    upper = line.splits[i+1][1]
                else:
                    lower = line.splits[i][0]
                    upper = line.splits[i+1][0]
                if lower <= x < upper:
                    counts[i] += 1
                    break
        #print(counts)
        #print("before", line)
        counts = counts / np.sum(counts)
        start = None
        end = None
        for idx, count in enumerate(counts):
            if count > 0.10:
                if start is None:
                    start = idx
                end = idx + 1
        if start is None:
            line.splits = []
        else:
            while start > 0 and np.sqrt((line.splits[start][0]-line.splits[start-1][0])**2 + (line.splits[start][1]-line.splits[start-1][1])**2) / scale < 30:
                #print("start", np.sqrt((line.splits[start][0]-line.splits[start-1][0])**2 + (line.splits[start][1]-line.splits[start-1][1])**2) / scale)
                start -= 1
            while end < len(line.splits)-1 and np.sqrt((line.splits[end][0]-line.splits[end+1][0])**2 + (line.splits[end][1]-line.splits[end+1][1])**2) / scale < 30:
                #print("end", np.sqrt((line.splits[end][0]-line.splits[end+1][0])**2 + (line.splits[end][1]-line.splits[end+1][1])**2) / scale)
                end += 1
            line.splits = line.splits[start:end+1]
        #print("after", line)

    for idx, line in enumerate(lines):
        new_splits = []
        for split in line.splits:
            if split[2] is None:
                new_splits.append(split)
                continue
            for split2 in lines[split[2]].splits:
                if split2[2] == idx:
                    new_splits.append(split)
                    break
        line.splits = new_splits
        #print("after2", line)

    print()

    # find largest rectangle
    def find_polygon(number, used, next, target, edges):
        #print(number, used, next, target, edges)

        max_score, corners = 0, None
        if number < 2:
            return max_score, corners
        if next is None:
            return max_score, corners
        current = lines[next]
        for neighbor in current.splits:
            #print(number, "--", neighbor)
            if next == target and number == 2:
                if neighbor[2] == used[0]:
                    e = np.array(edges[:] + [tuple(neighbor[:2])])
                    #print(used[:] + [next], target, edges[:] + [tuple(neighbor[:2])])
                    a = abs((e[0, 1] - e[2, 1])*(e[3, 0] - e[1, 0]) + (e[1, 1] - e[3, 1])*(e[0, 0] - e[2, 0]))/2
                    a = a/image.shape[0]/image.shape[1]
                    #print(a)
                    return a, edges[:] + [tuple(neighbor[:2])]
                else:
                    continue
            if neighbor[2] not in used:
                res_score, res_corners = find_polygon(number - 1, used[:] + [next], neighbor[2], target, edges[:] + [tuple(neighbor[:2])])
                if res_score > max_score:
                    max_score, corners = res_score, res_corners
        return max_score, corners

    max_score, corners = 0, None
    for lidx, line in enumerate(lines):
        for idx, s in enumerate(line.splits[:-1]):
            for e in line.splits[idx + 1:]:
                # print(line.name, (s[0], s[1], None if s[2] is None else lines[s[2]].name), (e[0], e[1], None if e[2] is None else lines[e[2]].name))
                res_score, res_corners = find_polygon(4, [lidx], e[2], s[2], [tuple(e[:2])])
                if res_score > max_score and res_score > 0.15:
                    max_score, corners = res_score, res_corners
    #print(max_score, corners)
    #print(image.shape)


    if corners is not None:
        corners = np.array(corners)
        # check rectangle validity
        l = corners.shape[0]
        for i in range(corners.shape[0]):
            a, b, c = i%l, (i+1)%l, (i+2)%l
            v1 = corners[b, :] - corners[a, :]
            v2 = corners[b, :] - corners[c, :]
            #print(v1, v2, np.linalg.norm(v1))
            phi = np.arccos(v1.dot(v2) / np.linalg.norm(v1) / np.linalg.norm(v2))
            #print(phi, corners[b])
            if phi < 0.8:
                corners = None
                break

    if corners is not None:
        # draw image (removable)
        draw_image = np.copy(image)
        for i in range(corners.shape[0]):
            a, b = corners[i%l], corners[(i+1)%l]
            for i in range(5000):
                x = int(a[1] + (b[1] - a[1]) * i / 5000)
                y = int(a[0] + (b[0] - a[0]) * i / 5000)
                draw_image[max(0, y - 15):y + 10, max(0, x - 10):x + 10] = np.array([255, 0, 0])
        #plt.imshow(draw_image)
        #plt.show()

        # crop image
        topleft = np.argmin(np.linalg.norm(corners, axis=1))
        corners = np.roll(corners, -topleft-1, axis=0)
        print(corners)
        h = int((corners[0, 0] + corners[1, 0] - corners[2, 0] - corners[3, 0]) / 2)
        w = int((corners[1, 1] + corners[2, 1] - corners[0, 1] - corners[3, 1]) / 2)
        pb = np.copy(corners)[:, ::-1].reshape((8,))
        img = Image.fromarray(image)
        convert = np.rot90(np.asarray(img.transform((h, w), Image.QUAD, pb, Image.BICUBIC)))
        plt.imshow(convert)
        plt.show()
        return convert
    else:
        return image
