from scan import *
import cvxpy as cp
from itertools import product


class EmptyTileError(Exception):
    pass


class OutOfBoundsError(Exception):
    pass


class Tile:
    def __init__(self, side_length=100, top_left=(0, 0), roi=(0, 0, 100, 100), img_names=None):
        self.side_length = side_length
        self.top_left = top_left
        self.roi = roi
        self.img_names = img_names

    def is_out_of_roi(self):
        x, y = self.top_left[0], self.top_left[1]
        check_x = (x >= self.roi[0]) & (x <= self.roi[0] + self.roi[2])
        check_y = (y >= self.roi[1]) & (y <= self.roi[1] + self.roi[3])
        return not (check_x & check_y)

    def check_intersections(self, data):
        tl = np.array(self.top_left)
        sl = self.side_length
        pts = [tl, tl + np.array([sl, 0]), tl + np.array([sl, sl]), tl + np.array([0, sl])]
        img_lis = []
        if have_overlap(pts, 0, np.array(get_vertices(data[0]['dst_shape'], 0, 0, 0)).reshape((4, 2)), 0):
            img_lis.append(data[0]['dst_name'])
        for dat in data:
            if have_overlap(pts, 0, get_rectangle(dat), dat['angle']):
                img_lis.append(dat['src_name'])
        if len(img_lis) > 0:
            self.img_names = img_lis


def get_linear_ineq(rect, angle):
    rad = angle * np.pi / 180
    if np.absolute(rad) <= np.pi / 4:
        slope = np.tan(rad)
        inter_1 = rect[0][1] - slope * rect[0][0]
        a1, a2, b1 = slope, - 1, - inter_1
        inter_2 = rect[3][1] - slope * rect[3][0]
        a3, a4, b2 = - slope, 1, inter_2
        inter_3 = rect[0][0] + slope * rect[0][1]
        a5, a6, b3 = - 1, - slope, - inter_3
        inter_4 = rect[1][0] + slope * rect[1][1]
        a7, a8, b4 = 1, slope, inter_4
    else:
        # exchange x and y here? right now only copied from above
        slope = np.tan(rad)
        inter_1 = rect[0][1] - slope * rect[0][0]
        a1, a2, b1 = slope, - 1, - inter_1
        inter_2 = rect[3][1] - slope * rect[3][0]
        a3, a4, b2 = - slope, 1, inter_2
        inter_3 = rect[0][0] + slope * rect[0][1]
        a5, a6, b3 = - 1, - slope, - inter_3
        inter_4 = rect[1][0] + slope * rect[1][1]
        a7, a8, b4 = 1, slope, inter_4
    A1 = np.array([[a1, a2], [a3, a4]])
    B1 = np.array([b1, b2])
    A2 = np.array([[a5, a6], [a7, a8]])
    B2 = np.array([b3, b4])
    # print(A1, B1)
    # print(A2, B2)
    return A1, B1, A2, B2


def have_overlap(rect_1, angle_1, rect_2, angle_2):
    p0 = rect_1[0]
    shifted_1 = copy.copy(rect_1)
    shifted_2 = copy.copy(rect_2)
    for i in range(4):
        shifted_1[i] = rect_1[i] - np.array(p0)
        shifted_2[i] = rect_2[i] - np.array(p0)
    A1, B1, A2, B2 = get_linear_ineq(shifted_1, - angle_1)
    A3, B3, A4, B4 = get_linear_ineq(shifted_2, - angle_2)
    x = cp.Variable(2)
    prob = cp.Problem(cp.Minimize(2), [A1 @ x <= B1, A2 @ x <= B2, A3 @ x <= B3, A4 @ x <= B4])
    prob.solve()
    solution = x.value
    return solution is not None


def get_rectangle(dat):
    return get_vertices(np.array(dat['src_shape']) * dat['scale'], dat['angle'], dat['x'], dat['y']).reshape((4, 2))


def get_warp_scale(data, name):
    for dat in data:
        if dat['src_name'] == name:
            return dat['scale']
        elif dat['dst_name'] == name:
            return 1


def get_angle(data, name):
    for dat in data:
        if dat['src_name'] == name:
            return dat['angle']
        elif dat['dst_name'] == name:
            return 0


def get_translation_vec(data, name):
    for dat in data:
        if dat['src_name'] == name:
            return dat['x'], dat['y']
        elif dat['dst_name'] == name:
            return 0, 0


def get_data_from_names(data, names):
    ret_data = []
    for name in names:
        for dat in data:
            if dat['src_name'] == name:
                ret_data.append(dat)
            else:
                ret_data.append(data[0])
    return ret_data


def smart_slice(output, input, x1, x2, y1, y2):
    xmin = max(0, x1)
    xmax = min(x2, input.shape[1])
    ymin = max(0, y1)
    ymax = min(y2, input.shape[0])
    tile_xmin = max(0, -x1)
    tile_xmax = min(input.shape[1] - x1, output.shape[1])
    tile_ymin = max(0, -y1)
    tile_ymax = min(input.shape[0] - y1, output.shape[0])
    output[tile_ymin: tile_ymax, tile_xmin: tile_xmax] = input[ymin: ymax, xmin: xmax]
    return output


def draw_tiles(data, side_length=200, incl_first=True):
    roi = get_roi(data, incl_first)
    print(roi)
    range_x = range(roi[0], roi[0] + roi[2], side_length)
    range_y = range(roi[1], roi[1] + roi[3], side_length)
    corners = product(range_x, range_y)
    tiles = []
    for top_left in corners:
        tile = Tile(side_length, top_left, roi)
        tile.check_intersections(data)
        if tile.img_names is not None:
            tiles.append(tile)
    tile_dict = {}
    for tile in tiles:
        if tuple(tile.img_names) in tile_dict.keys():
            tile_dict[tuple(tile.img_names)] += [tile.top_left]
        else:
            tile_dict[tuple(tile.img_names)] = [tile.top_left]
    for names, tile_corners in tile_dict.items():
        blender = cv.detail_MultiBandBlender()
        blend_width = 10
        num_bands = np.log(blend_width) / np.log(2.) - 1.
        blender.setNumBands(num_bands.astype(np.int32))
        local_roi = get_roi(get_data_from_names(data, names), incl_first)
        blender.prepare(local_roi)
        print(names, tile_corners)
        for name in names:
            warper = cv.PyRotationWarper('affine', get_warp_scale(data, name))
            img = cv.imread(f'imgs/danube_10_pct/{name}').astype(np.int16)
            corner, img = warper.warp(img, get_3x3_id(), compose_R_mat(get_angle(data, name), 1, 0, 0),
                                      cv.INTER_LINEAR, cv.BORDER_REFLECT)
            mask = cv.imread(f'imgs/danube_10_pct/masks/mask_{name}'.replace('jpg', 'png'), flags=cv.IMREAD_GRAYSCALE)
            mask = cv.UMat(mask)
            corner, mask = warper.warp(mask, get_3x3_id(), compose_R_mat(get_angle(data, name), 1, 0, 0),
                                       cv.INTER_NEAREST, cv.BORDER_CONSTANT)
            corner = (corner[0] + round(get_translation_vec(data, name)[0]),
                      corner[1] + round(get_translation_vec(data, name)[1]))
            blender.feed(cv.UMat(img), mask, corner)
        result, result_mask = blender.blend(None, None)
        for tl in tile_corners:
            tile_img = 255 * np.ones((side_length, side_length, 3), np.int16)
            slice_x = tl[0] - local_roi[0]
            slice_y = tl[1] - local_roi[1]
            img_slice = 255 * np.ones((side_length, side_length, 3), np.int16)
            img_slice = smart_slice(img_slice, result, slice_x, slice_x + side_length, slice_y, slice_y + side_length)
            mask_slice = 255 * np.ones((side_length, side_length), np.uint8)
            mask_slice = smart_slice(mask_slice, result_mask, slice_x, slice_x + side_length, slice_y, slice_y + side_length)
            mask_slice_inv = cv.bitwise_not(mask_slice)
            tile_img = cv.bitwise_and(tile_img, tile_img, mask=mask_slice_inv)
            img_slice = cv.bitwise_and(img_slice, img_slice, mask=mask_slice)
            tile_img = cv.add(tile_img, img_slice)
            cv.imwrite(f'imgs_out/tiles/{tl[0] - roi[0]}_{tl[1] - roi[1]}.png', tile_img)


if __name__ == '__main__':
    with open('10pct_stitch_data.json', 'r') as fp:
        all_comp = json.load(fp)
    # rect5 = get_rectangle(all_comp[21])
    # rect6 = get_rectangle(all_comp[22])
    # res = have_overlap(rect5, all_comp[21]['angle'], rect6, all_comp[22]['angle'])
    # if res is not None:
    #     roi = get_roi(all_comp[21:23], incl_first=False)
    #     circ_x = round(res[0] - roi[0])
    #     circ_y = round(res[1] - roi[1])
    #     print(circ_x)
    #     print(circ_y)
    #     img = cv.circle(img, (circ_x, circ_y), radius=5, color=0, thickness=-1)
    # roi = get_roi(all_comp[0:4], incl_first=True)
    # tile = Tile(200, (600, 40), roi)
    # print(tile.is_out_of_roi())
    # tile.check_intersections(all_comp[0:4])
    # print(tile.img_names)
    # img = draw_rectangles_from_data(all_comp[0:4], draw_first=True)
    # img = cv.rectangle(img, (tile.top_left[0] - roi[0], tile.top_left[1] - roi[1]),
    #                    (tile.top_left[0] + tile.side_length - roi[0], tile.top_left[1] + tile.side_length - roi[1]),
    #                    0, thickness=10, lineType=cv.LINE_8)
    # cv.imwrite('test.png', img)
    # draw_tiles(all_comp[14:19] + all_comp[-6:], incl_first=False)
    draw_tiles(all_comp, side_length=400, incl_first=True)
    #
    # debug in the case angle > pi/4:
    # rect1 = get_rectangle(all_comp[5])
    # rect2 = get_rectangle(all_comp[6])
    # print(have_overlap(rect1, all_comp[5]['angle'], rect2, all_comp[6]['angle']))
