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


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)
        features.append(img_feat)
    return features


def match_features(features, match_conf=0.3):
    matcher = cv.detail_AffineBestOf2NearestMatcher(False, False, match_conf)
    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')
        exit()
    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')
        exit()
    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')
        exit()
    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)
        corners.append(corner)
        images_warped.append(image_wp)
        p, mask_wp = warper.warp(masks[idx], K, cameras[idx].R, cv.INTER_LINEAR, cv.BORDER_REFLECT)
        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):
    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)
    indices, img_names_subset, img_subset, features_subset, matches_subset = get_biggest_subset(
        img_names, imgs_med_res, features, matches, conf_thresh=0.3)
    graph = get_matches_graph(img_names_subset, matches_subset, conf_thresh=0.3)
    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]
    reg_scale = get_scale(imgs[0], megapix=reg_megapix)
    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


if __name__ == '__main__':
    img_names = [f'pano_{i}.jpg' for i in range(1, 4)]
    imgs = [cv.imread(f'imgs/{name}') for name in img_names]
    # print(imgs[0].shape[0] * imgs[0].shape[1] / 1e6)
    result = full_stitching_pipeline(imgs, img_names)
    print(result.shape)
    out_name = 'imgs_out/pano.jpg'
    cv.imwrite(out_name, result)
    print('Stitched image saved as ' + out_name)
