"""
Stitching sample (advanced)
===========================
Show how to use Stitcher API from python.
Adapted from https://raw.githubusercontent.com/opencv/opencv/master/samples/python/stitching_detailed.py
"""
import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import json
import pathlib
import copy


def resize_image(img, scale):
    return cv.resize(src=img, dsize=None, fx=scale, fy=scale, interpolation=cv.INTER_LINEAR_EXACT)


def resize_mask(msk, target_msk):
    return cv.resize(src=msk, dsize=(target_msk.shape[1], target_msk.shape[0]), fx=0, fy=0,
                     interpolation=cv.INTER_LINEAR_EXACT)


def find_features(imgs):
    features = []
    finder = cv.ORB.create()
    for img in imgs:
        img_feat = cv.detail.computeImageFeatures2(finder, img)  # insert optional parameter 'mask' here
        features.append(img_feat)
    return features


def find_features_masked(imgs, masks):
    features = []
    finder = cv.ORB.create()
    for i, img in enumerate(imgs):
        img_feat = cv.detail.computeImageFeatures2(finder, img, masks[i])
        features.append(img_feat)
    return features


def match_features(features, match_conf=0.3):
    matcher = cv.detail_AffineBestOf2NearestMatcher(False, False, match_conf, num_matches_thresh1=6)
    matches = matcher.apply2(features)
    matcher.collectGarbage()
    return matches


def get_matches_graph(img_names, matches, conf_thresh=1.0):
    return cv.detail.matchesGraphAsString(img_names, matches, conf_thresh)


def get_biggest_subset(img_names, imgs, features, matches, conf_thresh=1.0, match_conf=0.3):
    indices = cv.detail.leaveBiggestComponent(features, matches, conf_thresh)
    img_subset = []
    img_names_subset = []
    features_subset = []
    num = len(indices)
    if num < 2:
        print('Need more images, no matches found!')
        raise SystemExit
    for i in range(num):
        img_subset.append(imgs[indices[i, 0]])
        img_names_subset.append(img_names[indices[i, 0]])
        features_subset.append(features[indices[i, 0]])
    matches_subset = match_features(features_subset, match_conf)
    return indices, img_names_subset, img_subset, features_subset, matches_subset


def estimate_cameras(features, matches):
    estimator = cv.detail_AffineBasedEstimator()
    b, cameras = estimator.apply(features, matches, None)
    if not b:
        print('Homography estimation failed')
        raise SystemExit
    for cam in cameras:
        cam.R = cam.R.astype(np.float32)
    return cameras


def refine_cameras(features, matches, cameras, refine_mask):
    adjuster = cv.detail_BundleAdjusterAffinePartial()
    adjuster.setConfThresh(1)
    mask = np.zeros((3, 3), np.uint8)
    if refine_mask[0] == 'x':
        mask[0, 0] = 1
    if refine_mask[1] == 'x':
        mask[0, 1] = 1
    if refine_mask[2] == 'x':
        mask[0, 2] = 1
    if refine_mask[3] == 'x':
        mask[1, 1] = 1
    if refine_mask[4] == 'x':
        mask[1, 2] = 1
    adjuster.setRefinementMask(mask)
    b, refined_cams = adjuster.apply(features, matches, cameras)
    if not b:
        print('Camera parameters adjusting failed')
        raise SystemExit
    return refined_cams


def get_warped_image_scale(cameras):
    focals = [cam.focal for cam in cameras]
    focals.sort()
    if len(focals) % 2 == 1:
        warped_image_scale = focals[len(focals) // 2]
    else:
        warped_image_scale = (focals[len(focals) // 2] + focals[len(focals) // 2 - 1]) / 2
    return warped_image_scale


def get_scale(img, megapix):
    return min(1.0, np.sqrt(megapix * 1e6 / (img.shape[0] * img.shape[1])))


def get_warped_masks_and_images(imgs, cameras, warp_scale, seam_aspect):
    corners = []
    masks = []
    masks_warped = []
    images_warped = []
    num_imgs = len(imgs)
    sizes = get_sizes(imgs)
    for i in range(0, num_imgs):
        um = cv.UMat(255 * np.ones((sizes[i][1], sizes[i][0]), np.uint8))
        masks.append(um)
    warper = cv.PyRotationWarper('affine', warp_scale * seam_aspect)
    for idx in range(num_imgs):
        K = cameras[idx].K().astype(np.float32)
        swa = seam_aspect
        K[0, 0] *= swa
        K[0, 2] *= swa
        K[1, 1] *= swa
        K[1, 2] *= swa
        corner, image_wp = warper.warp(imgs[idx], K, cameras[idx].R, cv.INTER_LINEAR, cv.BORDER_REFLECT)
        # corner, image_wp = warper.warp(imgs[idx], K, cameras[idx].R, cv.INTER_NEAREST, cv.BORDER_CONSTANT)
        corners.append(corner)
        images_warped.append(image_wp)
        p, mask_wp = warper.warp(masks[idx], K, cameras[idx].R, cv.INTER_NEAREST, cv.BORDER_CONSTANT)
        masks_warped.append(mask_wp.get())
    return corners, images_warped, masks_warped


def estimate_exposure_compensation(corners, images_warped, masks_warped):
    comp = cv.detail.ExposureCompensator_createDefault(cv.detail.ExposureCompensator_GAIN_BLOCKS)
    comp.feed(corners=corners, images=images_warped, masks=masks_warped)
    return comp


def find_seams(corners, imgs_w, msks_w):
    imgs_w_f = [img.astype(np.float32) for img in imgs_w]
    seam_finder = cv.detail.SeamFinder_createDefault(cv.detail.SeamFinder_VORONOI_SEAM)
    # seam_finder = cv.detail_DpSeamFinder('COLOR')
    # seam_finder = cv.detail.SeamFinder_createDefault(cv.detail.SeamFinder_NO)
    updated_masks = seam_finder.find(imgs_w_f, corners, msks_w)
    return updated_masks


def resize_seam_mask(msk_w, orig_msk):
    dilated_mask = cv.dilate(msk_w, None)
    seam_mask = resize_mask(dilated_mask, orig_msk)
    mask_warped = cv.bitwise_and(seam_mask, orig_msk)
    return mask_warped


def get_sizes(imgs):
    sizes = []
    for img in imgs:
        x, y = img.shape[1], img.shape[0]
        sizes.append((x, y))
    return sizes


def blend_images(corners, imgs, seam_msks, msks_w, comp):
    sizes = get_sizes(imgs)
    blender = cv.detail_MultiBandBlender()
    dst_sz = cv.detail.resultRoi(corners=corners, sizes=sizes)
    blend_strength = 5
    blend_width = np.sqrt(dst_sz[2] * dst_sz[3]) * blend_strength / 100
    blender.setNumBands((np.log(blend_width) / np.log(2.) - 1.).astype(np.int32))
    blender.prepare(dst_sz)
    for idx in range(len(imgs)):
        comp.apply(idx, corners[idx], imgs[idx], msks_w[idx])
        img_s = imgs[idx].astype(np.int16)
        msk_w = resize_seam_mask(seam_msks[idx], msks_w[idx])
        blender.feed(cv.UMat(img_s), msk_w, corners[idx])
    result = None
    result_mask = None
    return blender.blend(result, result_mask)


def get_reg_data(imgs, img_names, reg_megapix=0.6):
    if reg_megapix == -1:
        imgs_med_res = imgs
    else:
        reg_scale = get_scale(imgs[0], reg_megapix)
        imgs_med_res = [resize_image(img, reg_scale) for img in imgs]
    features = find_features(imgs_med_res)
    matches = match_features(features, match_conf=0.35)
    indices, img_names_subset, img_subset, features_subset, matches_subset = get_biggest_subset(
        img_names, imgs_med_res, features, matches, conf_thresh=0.6, match_conf=0.35)

    graph = get_matches_graph(img_names_subset, matches_subset, conf_thresh=0.6)
    cameras = estimate_cameras(features_subset, matches_subset)
    cameras = refine_cameras(features_subset, matches_subset, cameras, 'xxxxx')
    return indices, graph, cameras


def prepare_reg_data(imgs, indices, reg_megapix=0.6, seam_megapix=0.1):
    imgs_subset = [imgs[i[0]] for i in indices]
    if reg_megapix == -1:
        reg_scale = 1
    else:
        reg_scale = get_scale(imgs[0], megapix=reg_megapix)
    if seam_megapix == -1:
        seam_scale = 1
    else:
        seam_scale = get_scale(imgs[0], megapix=seam_megapix)
    return imgs_subset, reg_scale, seam_scale


def compose_from_reg_data(imgs, cameras, reg_scale, seam_scale):
    warp_scale_full = get_warped_image_scale(cameras)
    corners_full, images_w_full, masks_w_full = get_warped_masks_and_images(imgs, cameras, warp_scale_full, 1/reg_scale)
    imgs_low_res = [resize_image(img, seam_scale) for img in imgs]
    warp_scale_low = get_warped_image_scale(cameras)
    corners_low, images_w_low, masks_w_low = get_warped_masks_and_images(imgs_low_res, cameras, warp_scale_low,
                                                                         seam_scale/reg_scale)
    comp = estimate_exposure_compensation(corners_low, images_w_low, masks_w_low)
    seam_masks = find_seams(corners_low, images_w_low, masks_w_low)
    result, result_mask = blend_images(corners_full, images_w_full, seam_masks, masks_w_full, comp)
    return result, result_mask


def full_stitching_pipeline(imgs, img_names):
    indices, graph, cameras = get_reg_data(imgs, img_names, reg_megapix=0.6)
    imgs_subset, reg_scale, seam_scale = prepare_reg_data(imgs, indices, reg_megapix=0.6, seam_megapix=0.1)
    result, result_mask = compose_from_reg_data(imgs_subset, cameras, reg_scale, seam_scale)
    return result


def debug_stitching_pipeline(imgs, img_names):
    features = find_features(imgs)
    matches = match_features(features, match_conf=0.3)
    # for match in matches:
    #     print(match.getInliers())
    #     print(match.getMatches())
    #     print(match.confidence)
    #     print(match.src_img_idx, match.dst_img_idx)
    #     print(match.num_inliers)
    #     print(match.H)
    graph = get_matches_graph(img_names, matches, conf_thresh=0.3)
    print(graph)
    cameras = estimate_cameras(features, matches)
    # cameras = refine_cameras(features, matches, cameras, 'xxxxx')
    for cam in cameras:
        print(cam.R)
    corners, images_w, masks_w = get_warped_masks_and_images(imgs, cameras, 1, 1)
    print(corners)
    corners = [(-3788, -362), (-1872, -178), (0, 0), (1889, -1135), (3450, -2492)]
    # corners = [(0, 0), (1947, -556)] #- (2914, 2736) + (2879, 2733) = (-35, -3), (2305, 823) - (2335, 826) = (-30, -3)
    # corners = [(0, 0), (1912, -559)]
    # corners = [(0, 0), (1583, -88)]
    # cv.imwrite('32_w.jpg', images_w[0])
    # cv.imwrite('32_m_w.jpg', masks_w[0])
    # cv.imwrite('33_w.jpg', images_w[1])
    # cv.imwrite('33_m_w.jpg', masks_w[1])
    seam_masks = find_seams(corners, images_w, masks_w)
    cv.imwrite('24_sm.jpg', seam_masks[0])
    cv.imwrite('25_sm.jpg', seam_masks[1])
    sizes = get_sizes(images_w)
    blender = cv.detail_MultiBandBlender()
    dst_sz = cv.detail.resultRoi(corners=corners, sizes=sizes)
    blend_strength = 5
    blend_width = np.sqrt(dst_sz[2] * dst_sz[3]) * blend_strength / 100
    blender.setNumBands((np.log(blend_width) / np.log(2.) - 1.).astype(np.int32))
    blender.prepare(dst_sz)
    # comp = estimate_exposure_compensation(corners, images_w, masks_w)
    for idx in range(len(images_w)):
        # comp.apply(idx, corners[idx], images_w[idx], masks_w[idx])
        img_s = images_w[idx].astype(np.int16)
        blender.feed(cv.UMat(img_s), seam_masks[idx], corners[idx])
    result, result_mask = blender.blend(None, None)
    cv.imwrite('test.jpg', result)
    cv.imwrite('test_m.jpg', result_mask)


def feature_testing(img_names, comp, match_conf=0.3):
    features = []
    feat_finders = [cv.AKAZE_create(), cv.BRISK_create(), cv.KAZE_create(), cv.ORB_create(), cv.SIFT_create()]
    for name in img_names:
        img = cv.imread(name)
        feat_lis = []
        for finder in feat_finders:
            feat = cv.detail.computeImageFeatures2(finder, img)
            feat_lis.append(feat)
        features.append(feat_lis)
    features_transposed = list(map(list, zip(*features)))
    matches_lis = [match_features(feat, match_conf=match_conf) for feat in features_transposed]
    H_lis = []
    for m in matches_lis:
        c = []
        for info in m:
            if info.src_img_idx + 1 == info.dst_img_idx:
                c.append(info)
        H_lis.append(c)
    # H_transposed = list(map(list, zip(*H_lis)))
    # plt.plot(list(range(len(img_names) - 1)), confidences_transposed)
    # plt.title(f'Confidences of feature detection algorithms at match_conf={match_conf}')
    # plt.legend(['AKAZE', 'BRISK', 'KAZE', 'ORB', 'SIFT'])
    # plt.savefig(f'comp_{comp}_match_conf_{match_conf}.png')
    # plt.clf()
    return H_lis


def match_sift_features(img_names, match_conf=0.4):
    features = []
    feat_finder = cv.SIFT_create()
    for name in img_names:
        img = cv.imread(name)
        feat = cv.detail.computeImageFeatures2(feat_finder, img)
        features.append(feat)
    matches = match_features(features, match_conf=match_conf)
    overlap_matches = []
    for match in matches:
        if match.src_img_idx + 1 == match.dst_img_idx:
            overlap_matches.append(match)
    return overlap_matches


def extract_from_R_mat(R_mat):
    scale = np.sqrt(R_mat[0, 0] ** 2 + R_mat[0, 1] ** 2)
    angle = np.arcsin(R_mat[1, 0] / scale) * 180 / np.pi
    t_x = R_mat[0, 2]
    t_y = R_mat[1, 2]
    return angle, scale, t_x, t_y


def compose_R_mat(angle, scale, t_x, t_y):
    rad = angle * np.pi / 180
    R = np.zeros((3, 3), np.float32)
    R[0, 0] = np.cos(rad) * scale
    R[0, 1] = - np.sin(rad) * scale
    R[0, 2] = t_x
    R[1, 0] = np.sin(rad) * scale
    R[1, 1] = np.cos(rad) * scale
    R[1, 2] = t_y
    R[2, 2] = 1
    return R


def truncate_path(s):
    return s.replace('imgs/danube_10_pct/', '')


def get_img_shape(img_name):
    img = cv.imread(img_name)
    return img.shape[1], img.shape[0]


def get_match_details(img_names, matches_lis):
    details_lis = []
    for match in matches_lis:
        details = {}
        angle, scale, x, y = extract_from_R_mat(match.H)
        src_idx = match.src_img_idx
        dst_idx = match.dst_img_idx
        details['src_name'] = truncate_path(img_names[src_idx])
        details['src_shape'] = get_img_shape(img_names[src_idx])
        details['dst_name'] = truncate_path(img_names[dst_idx])
        details['dst_shape'] = get_img_shape(img_names[dst_idx])
        details['angle'] = angle
        details['scale'] = scale
        details['x'] = x
        details['y'] = y
        details_lis.append(details)
    return details_lis


def get_3x3_id():
    K = np.zeros((3, 3), np.float32)
    K[0, 0] = 1.0
    K[1, 1] = 1.0
    K[2, 2] = 1.0
    return K


def rotate_vec(pt, angle, center):
    x, y = pt[0] - center[0], pt[1] - center[1]
    rad = angle * np.pi / 180
    rot_x, rot_y = x * np.cos(rad) - y * np.sin(rad), x * np.sin(rad) + y * np.cos(rad)
    return rot_x + center[0], rot_y + center[1]


def get_seam_mask(img_name, data):
    subset_data = []
    nsh_lis = {}
    for data_set in data:
        if data_set['src_name'] == img_name or data_set['dst_name'] == img_name:
            subset_data.append(data_set)
            nsh_lis[data_set['src_name']] = data_set['src_shape']
            nsh_lis[data_set['dst_name']] = data_set['dst_shape']
    masks = {}
    for name, shape in nsh_lis.items():
        masks[name] = cv.UMat(255 * np.ones((shape[1], shape[0]), np.uint8))
    if subset_data[0]['src_name'] == img_name:  # start of the chain
        warper2 = cv.PyRotationWarper('affine', 1 / subset_data[0]['scale'])
        R_2 = compose_R_mat(subset_data[0]['angle'], 1, 0, 0)
        corner2, mask_w_2 = warper2.warp(masks[subset_data[0]['dst_name']], get_3x3_id(), R_2, cv.INTER_LINEAR,
                                         cv.BORDER_CONSTANT)
        rot_x, rot_y = rotate_vec((-subset_data[0]['x'] / subset_data[0]['scale'],
                                   -subset_data[0]['y'] / subset_data[0]['scale']),
                                  -subset_data[0]['angle'], (0, 0))
        corner2 = (corner2[0] + rot_x, corner2[1] + rot_y)
        corners = [(0, 0), corner2]
        masks_w = [masks[subset_data[0]['src_name']], mask_w_2]
        seam_finder = cv.detail.SeamFinder_createDefault(cv.detail.SeamFinder_VORONOI_SEAM)
        updated_masks = seam_finder.find(masks_w, corners, masks_w)
        return updated_masks[0].get()
    elif (len(nsh_lis) == 2) and (subset_data[0]['dst_name'] == img_name):  # end of the chain
        warper1 = cv.PyRotationWarper('affine', subset_data[0]['scale'])
        R_1 = compose_R_mat(-subset_data[0]['angle'], 1, 0, 0)
        corner1, mask_w_1 = warper1.warp(masks[subset_data[0]['src_name']], get_3x3_id(), R_1, cv.INTER_LINEAR,
                                         cv.BORDER_CONSTANT)
        corner1 = (corner1[0] + subset_data[0]['x'], corner1[1] + subset_data[0]['y'])
        corners = [corner1, (0, 0)]
        masks_w = [mask_w_1, masks[subset_data[0]['dst_name']]]
        seam_finder = cv.detail.SeamFinder_createDefault(cv.detail.SeamFinder_VORONOI_SEAM)
        updated_masks = seam_finder.find(masks_w, corners, masks_w)
        return updated_masks[1].get()
    else:  # middle of the chain
        warper1 = cv.PyRotationWarper('affine', subset_data[0]['scale'])
        R_1 = compose_R_mat(-subset_data[0]['angle'], 1, 0, 0)
        corner1, mask_w_1 = warper1.warp(masks[subset_data[0]['src_name']], get_3x3_id(), R_1, cv.INTER_LINEAR,
                                         cv.BORDER_CONSTANT)
        corner1 = (corner1[0] + subset_data[0]['x'], corner1[1] + subset_data[0]['y'])
        warper2 = cv.PyRotationWarper('affine', 1/subset_data[1]['scale'])
        R_2 = compose_R_mat(subset_data[1]['angle'], 1, 0, 0)
        corner2, mask_w_2 = warper2.warp(masks[subset_data[1]['dst_name']], get_3x3_id(), R_2, cv.INTER_LINEAR,
                                         cv.BORDER_CONSTANT)
        rot_x, rot_y = rotate_vec((-subset_data[1]['x']/subset_data[1]['scale'],
                                   -subset_data[1]['y']/subset_data[1]['scale']),
                                  -subset_data[1]['angle'], (0, 0))
        corner2 = (corner2[0] + rot_x, corner2[1] + rot_y)
        corners = [corner1, (0, 0), corner2]
        masks_w = [mask_w_1, masks[subset_data[0]['dst_name']], mask_w_2]
        seam_finder = cv.detail.SeamFinder_createDefault(cv.detail.SeamFinder_VORONOI_SEAM)
        updated_masks = seam_finder.find(masks_w, corners, masks_w)
        return updated_masks[1].get()


def transform_data_relative_to_first(data):
    new_data = []
    angle, scale, x, y = 0, 1, 0, 0
    first_name = data[0]['src_name']
    first_shape = data[0]['src_shape']
    for dat in data:
        scale *= 1 / dat['scale']
        angle += dat['angle']
        xr, yr = rotate_vec((- dat['x'] * scale, - dat['y'] * scale), - angle, (0, 0))
        x += xr
        y += yr
        new_dat = {'src_name': dat['dst_name'],
                   'src_shape': dat['dst_shape'],
                   'dst_name': first_name,
                   'dst_shape': first_shape,
                   'angle': angle,
                   'scale': scale,
                   'x': x,
                   'y': y}
        new_data.append(new_dat)
    return new_data


def get_vertices(shape, angle, x, y):
    center = [shape[0]/2, shape[1]/2]
    pts = [(0, 0), (shape[0], 0), (shape[0], shape[1]), (0, shape[1])]
    vertices = []
    x0, y0 = rotate_vec(pts[0], - angle, center)
    for i in range(0, 4):
        vert = rotate_vec(pts[i], - angle, center)
        vert = [round(vert[0] + x - x0), round(vert[1] + y - y0)]
        vertices.append(vert)
    vertices = np.array(vertices, np.int32).reshape((-1, 1, 2))
    return vertices


def get_roi(data, incl_first=True):
    if incl_first:
        points = [get_vertices(data[0]['dst_shape'], 0, 0, 0)]
    else:
        points = []
    for dat in data:
        rect = get_vertices(np.array(dat['src_shape']) * dat['scale'], dat['angle'], dat['x'], dat['y'])
        points.append(rect)
    points = np.reshape(points, (-1, 2))
    pt_max = np.max(points, axis=0)
    pt_min = np.min(points, axis=0)
    x, y = pt_min[0], pt_min[1]
    w, h = pt_max[0] - pt_min[0] + 1, pt_max[1] - pt_min[1] + 1
    return x, y, w, h


def draw_rectangles_from_data(data, draw_first=True):
    dst_sz = get_roi(data, draw_first)
    new_img = 255 * np.ones((dst_sz[3], dst_sz[2]), np.uint8)
    if draw_first:
        new_img = cv.rectangle(new_img, (- dst_sz[0], - dst_sz[1]),
                               (data[0]['dst_shape'][0] - dst_sz[0], data[0]['dst_shape'][1] - dst_sz[1]),
                               128, thickness=10, lineType=cv.LINE_8)
    for dat in data:
        points = get_vertices(np.array(dat['src_shape']) * dat['scale'], dat['angle'], dat['x'] - dst_sz[0],
                              dat['y'] - dst_sz[1])
        new_img = cv.polylines(img=new_img, pts=[points],
                               isClosed=True, color=128, thickness=10)
        points = points.reshape((4, 2))
        for pt in points:
            # print(pt)
            new_img = cv.putText(new_img, str(int(pt[0])) + ', ' + str(int(pt[1])), pt, 0, 10, 0, thickness=5)
    new_img = cv.resize(new_img, None, fx=0.15, fy=0.15, interpolation=cv.INTER_AREA)
    # cv.imwrite('test.png', new_img)
    return new_img


def draw_images_from_data(data, draw_first=True):
    dst_sz = get_roi(data, draw_first)
    print(dst_sz)
    blender = cv.detail_MultiBandBlender()
    # blend_strength = 5
    # blend_width = np.sqrt(dst_sz[2] * dst_sz[3]) * blend_strength / 100
    blend_width = 5
    num_bands = np.log(blend_width) / np.log(2.) - 1.
    print(num_bands)
    blender.setNumBands(num_bands.astype(np.int32))
    blender.prepare(dst_sz)
    if draw_first:
        img = cv.imread('imgs/danube_10_pct/9.jpg').astype(np.int16)
        mask = cv.imread('imgs/danube_10_pct/masks/mask_9.png', flags=cv.IMREAD_GRAYSCALE)
        blender.feed(cv.UMat(img), cv.UMat(mask), (0, 0))
    for dat in data:
        print(dat['src_name'])
        warper = cv.PyRotationWarper('affine', dat['scale'])
        img = cv.imread(f'imgs/danube_10_pct/{dat["src_name"]}').astype(np.int16)
        corner, img = warper.warp(img, get_3x3_id(), compose_R_mat(dat['angle'], 1, 0, 0),
                                  cv.INTER_LINEAR, cv.BORDER_REFLECT)
                                  # cv.INTER_NEAREST, cv.BORDER_CONSTANT)
        mask_path = pathlib.Path(f'imgs/danube_10_pct/masks/mask_{dat["src_name"]}'.replace('.jpg', '_cropped.jpg'))
        mask_png = pathlib.Path(f'imgs/danube_10_pct/masks/mask_{dat["src_name"]}'.replace('.jpg', '.png'))
        if mask_png.exists():
            mask = cv.imread(mask_png.as_posix(), flags=cv.IMREAD_GRAYSCALE)
        elif mask_path.exists():
            mask = cv.imread(mask_path.as_posix(), flags=cv.IMREAD_GRAYSCALE)
        else:
            mask = cv.imread(f'imgs/danube_10_pct/masks/mask_{dat["src_name"]}', flags=cv.IMREAD_GRAYSCALE)
        mask = cv.UMat(mask)
        corner, mask = warper.warp(mask, get_3x3_id(), compose_R_mat(dat['angle'], 1, 0, 0),
                                   cv.INTER_NEAREST, cv.BORDER_CONSTANT)
        corner = (corner[0] + round(dat['x']), corner[1] + round(dat['y']))
        blender.feed(cv.UMat(img), mask, corner)
    result, result_mask = blender.blend(None, None)
    result = cv.resize(result, None, fx=0.15, fy=0.15, interpolation=cv.INTER_AREA)
    result_mask = cv.resize(result_mask, None, fx=0.15, fy=0.15, interpolation=cv.INTER_AREA)
    cv.imwrite('test_draw_images.png', result)
    cv.imwrite('test_draw_images_mask.png', result_mask)


def connections_between_components():
    conn_1_2 = {'src_name': '42.jpg', 'src_shape': [1083, 624],
                'dst_name': '35.jpg', 'dst_shape': [891, 627],
                'angle': 14.75, 'scale': 1.0035,
                'x': 713.0, 'y': 46.5}
    # conn_1_2_old = {'src_name': '42.jpg', 'src_shape': [1083, 624],
    #             'dst_name': '35.jpg', 'dst_shape': [891, 627],
    #             'angle': 14.68, 'scale': 1.0035,
    #             'x': 715.0, 'y': 46.0}
    conn_1_3 = {'src_name': '39.jpg', 'src_shape': [1121, 1352],
                'dst_name': '24.jpg', 'dst_shape': [1072, 633],
                'angle': 0.190, 'scale': 1.012,
                'x': -193.0, 'y': -1277.0}
    # conn_1_3_old = {'src_name': '39.jpg', 'src_shape': [1121, 1352],
    #             'dst_name': '24.jpg', 'dst_shape': [1072, 633],
    #             'angle': 0.0, 'scale': 1.00,
    #             'x': -183.0, 'y': -1260.0}
    conn_1_4 = {'src_name': '25.jpg', 'src_shape': [977, 627],
                'dst_name': '36.jpg', 'dst_shape': [1045, 825],
                'angle': 0.0, 'scale': 0.997,
                'x': -95.0, 'y': -607.0}
    # conn_1_4_old = {'src_name': '25.jpg', 'src_shape': [977, 627],
    #             'dst_name': '36.jpg', 'dst_shape': [1045, 825],
    #             'angle': 0.0, 'scale': 1.0,
    #             'x': -96.0, 'y': -601.0}
    conn_1_5 = {'src_name': '38.jpg', 'src_shape': [1715, 1083],
                'dst_name': '28.jpg', 'dst_shape': [1020, 627],
                'angle': 0.0, 'scale': 1.0065,
                'x': -665.0, 'y': -1059.0}
    # conn_1_5_old = {'src_name': '38.jpg', 'src_shape': [1715, 1083],
    #             'dst_name': '28.jpg', 'dst_shape': [1020, 627],
    #             'angle': 0.0, 'scale': 1.0,
    #             'x': -660.0, 'y': -1051.0}
    return [conn_1_2, conn_1_3, conn_1_4, conn_1_5]


def invert_coord_frame(dat):
    xr, yr = rotate_vec((dat['x'] / dat['scale'], dat['y'] / dat['scale']), dat['angle'], (0, 0))
    new_dat = {'src_name': dat['dst_name'], 'src_shape': dat['dst_shape'],
               'dst_name': dat['src_name'], 'dst_shape': dat['src_shape'],
               'angle': dat['angle'], 'scale': 1 / dat['scale'],
               'x': -xr, 'y': -yr}
    return new_dat


def do_corrections(data):
    data[0]['data'][6]['scale'] = 1.0
    data[0]['data'][6]['x'] += 9
    data[0]['data'][6]['y'] += -5
    data[0]['data'][6]['angle'] += 0.55
    data[0]['data'][7]['scale'] = 1.0
    data[0]['data'][7]['x'] += -2
    data[0]['data'][7]['y'] += -1
    data[0]['data'][22]['angle'] += -0.35
    data[0]['data'][22]['scale'] = 0.995
    data[0]['data'][22]['x'] += 4
    data[0]['data'][22]['y'] += -5
    data[0]['data'][23]['scale'] = 0.997
    data[0]['data'][23]['angle'] += -0.1
    data[0]['data'][23]['x'] += 4
    data[0]['data'][23]['y'] += 4
    data[0]['data'][24]['scale'] *= 1.004
    data[0]['data'][24]['x'] += -3
    data[0]['data'][24]['y'] += -2
    return data


def get_compass_list():
    lis = {'9.jpg': (3797.6 - 3963.0, 201.1 - 526.8),
           '10.jpg': (3886.1 - 4048.0, 201.2 - 498.9),
           '11.jpg': (3995.0 - 4145.5, 207.7 - 490.5),
           '12.jpg': (3996.3 - 4133.0, 224.0 - 457.5),
           '13.jpg': (3731.0 - 3941.2, 533.3 - 828.0),
           '14.jpg': (4595.2 - 4664.9, 477.5 - 808.6),
           '15.jpg': (4018.7 - 4020.3, 646.9 - 996.0),
           '16.jpg': (4204.3 - 4022.3, 236.3 - 548.0),
           '17.jpg': (3850.3 - 3740.0, 153.0 - 409.8),
           '18.jpg': (4225.9 - 4232.4, 260.7 - 523.5),
           '19.jpg': (3982.9 - 4009.5, 254.4 - 518.1),
           '20.jpg': (4285.1 - 4378.9, 201.6 - 360.2),
           '21.jpg': (4183.0 - 4309.0, 358.4 - 586.7),
           '22.jpg': (4049.5 - 4154.0, 722.5 - 987.0),
           '23.jpg': (4467.6 - 4515.1, 424.0 - 618.2),
           '24.jpg': (4240.2 - 4288.0, 105.1 - 326.5),
           '25.jpg': (4174.0 - 4239.2, 183.8 - 364.0),
           '26.jpg': (4380.2 - 4430.0, 559.4 - 726.8),
           '27.jpg': (4531.4 - 4532.0, 232.0 - 447.8),
           '28.jpg': (4096.2 - 4094.2, 231.1 - 420.8),
           '29.jpg': (4281.7 - 4279.4, 248.5 - 521.4),
           '30.jpg': (3948.5 - 3945.3, 115.4 - 268.0),
           '32.jpg': (3885.4 - 3994.5, 431.1 - 433.5),
           '33.jpg': (4377.0 - 4555.2, 137.2 - 97.0),
           '34.jpg': (4025.5 - 4232.5, 285.1 - 237.6),
           '35.jpg': (4184.3 - 4395.7, 255.0 - 206.5),
           '38.jpg': (7396.4 - 7399.5, 2450.6 - 2672.2),
           '41.jpg': (4560.0 - 4638.8, 4783.5 - 4980.0),
           '42.jpg': (4044.4 - 4255.5, 199.4 - 202.9),
           '43.jpg': (3649.2 - 3855.9, 129.0 - 132.1),
           '44.jpg': (3664.5 - 3860.2, 108.1 - 104.6),
           '45.jpg': (4596.0 - 4722.2, 209.6 - 169.8),
           '46.jpg': (5048.7 - 5252.4, 327.8 - 328.5),
           '47.jpg': (5445.3 - 5648.4, 2864.5 - 2843.5),
           '48.jpg': (4672.9 - 4825.3, 1752.7 - 1766.9),
           '49.jpg': (3920.5 - 4109.8, 1685.0 - 1761.2),
           '50.jpg': (3842.5 - 4003.3, 1725.1 - 1788.0)}
    return lis


def create_matches_data():
    path = 'imgs/danube_10_pct'
    comp_1 = [f'{path}/{i}.jpg' for i in range(9, 36)]
    comp_2 = [f'{path}/{i}.jpg' for i in range(42, 51)]
    comp_3 = [f'{path}/39.jpg', f'{path}/40.jpg', f'{path}/41.jpg']
    comp_4 = [f'{path}/36.jpg', f'{path}/37.jpg']
    comp_5 = [f'{path}/38.jpg']
    all_matches = [match_sift_features(eval(f'comp_{i}'), match_conf=0.4) for i in range(1, 5)]
    match_details = [get_match_details(eval(f'comp_{i+1}'), all_matches[i]) for i in range(0, 4)]
    components = [{'component': i+1, 'data': match_details[i]} for i in range(0, 4)]
    with open('stitch_data.json', 'w') as fp:
        json.dump(components, fp)


def collect_all_components(components):
    conn = connections_between_components()
    connect_1_2 = invert_coord_frame(conn[0])
    comp_1_2 = components[0]['data'] + [connect_1_2] + components[1]['data']
    comp_1_2_abs = transform_data_relative_to_first(comp_1_2)
    comp_1_3 = [invert_coord_frame(comp_1_2_abs[14]), invert_coord_frame(conn[1])] + components[2]['data']
    comp_1_3_abs = transform_data_relative_to_first(comp_1_3)
    comp_1_4 = [invert_coord_frame(comp_1_2_abs[15]), conn[2]] + components[3]['data']
    comp_1_4_abs = transform_data_relative_to_first(comp_1_4)
    comp_1_5 = [invert_coord_frame(comp_1_2_abs[18]), invert_coord_frame(conn[3])]
    comp_1_5_abs = transform_data_relative_to_first(comp_1_5)
    all_comp = comp_1_2_abs + comp_1_3_abs[1:] + comp_1_4_abs[1:] + comp_1_5_abs[1:]
    return all_comp


if __name__ == '__main__':
    path = 'imgs/danube_10_pct'
    comp_1 = [f'{path}/{i}.jpg' for i in range(9, 36)]
    comp_2 = [f'{path}/{i}.jpg' for i in range(42, 51)]
    comp_3 = [f'{path}/39.jpg', f'{path}/40.jpg', f'{path}/41.jpg']
    comp_4 = [f'{path}/36.jpg', f'{path}/37.jpg']
    comp_5 = [f'{path}/38.jpg']
    img_names = [f'imgs/danube_10_pct/{i}.jpg' for i in range(9, 51)]
    # imgs = [cv.imread(f'imgs/{name}') for name in img_names]
    debug_mode = True
    try:
        if debug_mode:
            # create_matches_data()
            with open('stitch_data.json', 'r') as fp:
                components = json.load(fp)
            components = do_corrections(components)
            all_comp = collect_all_components(components)
            # with open('10pct_stitch_data.json', 'w') as fp:
            #     json.dump(all_comp, fp)
            # draw_images_from_data(all_comp[14:19] + all_comp[-6:], draw_first=False)
            # draw_images_from_data(all_comp[17:19] + all_comp[-1:], draw_first=False)
            draw_images_from_data(all_comp, draw_first=True)
            # rect_img = draw_rectangles_from_data(all_comp, draw_first=True)
            # cv.imwrite('test_draw_rectangles.png', rect_img)
        else:
            imgs = [cv.imread(name) for name in comp_2]
            result = full_stitching_pipeline(imgs, comp_2)
            res_name = 'area.jpg'
            cv.imwrite(res_name, result)
    except SystemExit:
        print('program execution stopped!')
