"""
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


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 truncate_path(s):
    return s.replace('imgs/danube_10_pct/', '')


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_img_shape(img_name):
    img = cv.imread(img_name)
    return img.shape[1], img.shape[0]


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:
            # test_scale = ['imgs/danube_10_pct/13.jpg', 'imgs/danube_10_pct/14_scaled_down.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)]
            for comp in components:
                print(comp)
            with open('stitch_data.json', 'w') as fp:
                json.dump(components, fp)
            # with open('stitch_data.json', 'r') as fp:
            #     components = json.load(fp)
            # H_2 = feature_testing(comp_2, 2, match_conf=0.4)
            # H_3 = feature_testing(comp_3, 3, match_conf=0.4)
            # H_4 = feature_testing(comp_4, 4, match_conf=0.4)
            # conf_1 = feature_testing(comp_1, 1, match_conf=0.8)
            # conf_2 = feature_testing(comp_2, 2, match_conf=0.8)
            # conf_3 = feature_testing(comp_3, 3, match_conf=0.8)
            # conf_4 = feature_testing(comp_4, 4, match_conf=0.8)
            # with open('confs_match_conf_0.8.csv', 'w') as file:
            #     json.dump([conf_1, conf_2, conf_3, conf_4], file)
            # confs = []
            # for i in range(2, 9):
            #     with open(f'confs_match_conf_0.{i}.csv', 'r') as file:
            #         lis = json.load(file)
            #     flat_lis = [item for sublist in lis for item in sublist]
            #     avg_conf = np.mean(flat_lis, axis=0)
            #     confs.append(avg_conf)
            # plt.plot(np.arange(0.2, 0.8, 0.1), confs)
            # plt.legend(['AKAZE', 'BRISK', 'KAZE', 'ORB', 'SIFT'])
            # plt.title('Average confidence for different match_conf')
            # plt.savefig('avg_conf.png')
            # orb = cv.ORB.create()
            # crop1 = imgs[0][0:2765, 1920:3246]
            # crop2 = imgs[1][0:3089, 0:1170]
            # imgs = [crop1, crop2]
            # debug_stitching_pipeline(imgs, img_names)
            # kp1, des1 = orb.detectAndCompute(imgs[0], None)
            # kp2, des2 = orb.detectAndCompute(imgs[1], None)
            # feat1 = cv.detail.computeImageFeatures2(orb, imgs[0])
            # feat2 = cv.detail.computeImageFeatures2(orb, imgs[1])
            # afm = cv.detail_AffineBestOf2NearestMatcher.create()
            # matches_info = afm.apply(feat1, feat2)
            # matches = matches_info.getMatches()
            # matches = sorted(matches, key=lambda x: x.distance)
            # img = cv.drawMatches(imgs[0], kp1, imgs[1], kp2, matches[:20], None)
            # cv.imwrite('imgs_out/matches3.jpg', img)
        else:
            imgs = [cv.imread(name) for name in img_names]
            result = full_stitching_pipeline(imgs, img_names)
            res_name = 'area.jpg'
            cv.imwrite(res_name, result)
    except SystemExit:
        print('program execution stopped!')
