# -*- coding: utf-8 -*-
# @Author  : xuelun

import cv2
import torch
import argparse
import warnings
import numpy as np
import matplotlib.pyplot as plt
import torchvision.transforms.functional as F

from os.path import join

from dkm.models.model_zoo.DKMv3 import DKMv3
from gluefactory.superpoint import SuperPoint
from gluefactory.models.matchers.lightglue import LightGlue

DEFAULT_MIN_NUM_MATCHES = 4
DEFAULT_RANSAC_MAX_ITER = 10000
DEFAULT_RANSAC_CONFIDENCE = 0.999
DEFAULT_RANSAC_REPROJ_THRESHOLD = 8
DEFAULT_RANSAC_METHOD = "USAC_MAGSAC"

RANSAC_ZOO = {
    "RANSAC": cv2.RANSAC,
    "USAC_FAST": cv2.USAC_FAST,
    "USAC_MAGSAC": cv2.USAC_MAGSAC,
    "USAC_PROSAC": cv2.USAC_PROSAC,
    "USAC_DEFAULT": cv2.USAC_DEFAULT,
    "USAC_FM_8PTS": cv2.USAC_FM_8PTS,
    "USAC_ACCURATE": cv2.USAC_ACCURATE,
    "USAC_PARALLEL": cv2.USAC_PARALLEL,
}


def read_image(path, grayscale=False):
    if grayscale:
        mode = cv2.IMREAD_GRAYSCALE
    else:
        mode = cv2.IMREAD_COLOR
    image = cv2.imread(str(path), mode)
    if image is None:
        raise ValueError(f"Cannot read image {path}.")
    if not grayscale and len(image.shape) == 3:
        image = image[:, :, ::-1]  # BGR to RGB
    return image


def resize_image(image, size, interp):
    assert interp.startswith("cv2_")
    if interp.startswith("cv2_"):
        interp = getattr(cv2, "INTER_" + interp[len("cv2_") :].upper())
        h, w = image.shape[:2]
        if interp == cv2.INTER_AREA and (w < size[0] or h < size[1]):
            interp = cv2.INTER_LINEAR
        resized = cv2.resize(image, size, interpolation=interp)
    # elif interp.startswith('pil_'):
    #     interp = getattr(PIL.Image, interp[len('pil_'):].upper())
    #     resized = PIL.Image.fromarray(image.astype(np.uint8))
    #     resized = resized.resize(size, resample=interp)
    #     resized = np.asarray(resized, dtype=image.dtype)
    else:
        raise ValueError(f"Unknown interpolation {interp}.")
    return resized


def fast_make_matching_figure(data, b_id):
    color0 = (
        (data["color0"][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8)
    )  # (rH, rW, 3)
    color1 = (
        (data["color1"][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8)
    )  # (rH, rW, 3)
    gray0 = cv2.cvtColor(color0, cv2.COLOR_RGB2GRAY)
    gray1 = cv2.cvtColor(color1, cv2.COLOR_RGB2GRAY)
    kpts0 = data["mkpts0_f"].cpu().detach().numpy()
    kpts1 = data["mkpts1_f"].cpu().detach().numpy()
    mconf = data["mconf"].cpu().detach().numpy()
    inliers = data["inliers"]

    rows = 2
    margin = 2
    (h0, w0), (h1, w1) = data["hw0_i"], data["hw1_i"]
    h = max(h0, h1)
    H, W = margin * (rows + 1) + h * rows, margin * 3 + w0 + w1

    # canvas
    out = 255 * np.ones((H, W), np.uint8)

    wx = [margin, margin + w0, margin + w0 + margin, margin + w0 + margin + w1]
    hx = lambda row: margin * row + h * (row - 1)
    out = np.stack([out] * 3, -1)

    sh = hx(row=1)
    out[sh : sh + h0, wx[0] : wx[1]] = color0
    out[sh : sh + h1, wx[2] : wx[3]] = color1

    sh = hx(row=2)
    out[sh : sh + h0, wx[0] : wx[1]] = color0
    out[sh : sh + h1, wx[2] : wx[3]] = color1
    mkpts0, mkpts1 = np.round(kpts0).astype(int), np.round(kpts1).astype(int)
    for (x0, y0), (x1, y1) in zip(mkpts0[inliers], mkpts1[inliers]):
        c = (0, 255, 0)
        cv2.circle(out, (x0, y0 + sh), 3, c, -1, lineType=cv2.LINE_AA)
        cv2.circle(out, (x1 + margin + w0, y1 + sh), 3, c, -1, lineType=cv2.LINE_AA)

    return out


def fast_make_matching_overlay(data, b_id):
    color0 = (
        (data["color0"][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8)
    )  # (rH, rW, 3)
    color1 = (
        (data["color1"][b_id].permute(1, 2, 0).cpu().detach().numpy() * 255).round().astype(np.uint8)
    )  # (rH, rW, 3)
    gray0 = cv2.cvtColor(color0, cv2.COLOR_RGB2GRAY)
    gray1 = cv2.cvtColor(color1, cv2.COLOR_RGB2GRAY)
    kpts0 = data["mkpts0_f"].cpu().detach().numpy()
    kpts1 = data["mkpts1_f"].cpu().detach().numpy()
    mconf = data["mconf"].cpu().detach().numpy()
    inliers = data["inliers"]

    rows = 2
    margin = 2
    (h0, w0), (h1, w1) = data["hw0_i"], data["hw1_i"]
    h = max(h0, h1)
    H, W = margin * (rows + 1) + h * rows, margin * 3 + w0 + w1

    # canvas
    out = 255 * np.ones((H, W), np.uint8)

    wx = [margin, margin + w0, margin + w0 + margin, margin + w0 + margin + w1]
    hx = lambda row: margin * row + h * (row - 1)
    out = np.stack([out] * 3, -1)

    sh = hx(row=1)
    out[sh : sh + h0, wx[0] : wx[1]] = color0
    out[sh : sh + h1, wx[2] : wx[3]] = color1

    sh = hx(row=2)
    out[sh : sh + h0, wx[0] : wx[1]] = color0
    out[sh : sh + h1, wx[2] : wx[3]] = color1
    mkpts0, mkpts1 = np.round(kpts0).astype(int), np.round(kpts1).astype(int)
    for (x0, y0), (x1, y1) in zip(mkpts0[inliers], mkpts1[inliers]):
        c = (0, 255, 0)
        cv2.line(out, (x0, y0 + sh), (x1 + margin + w0, y1 + sh), color=c, thickness=1, lineType=cv2.LINE_AA)
        cv2.circle(out, (x0, y0 + sh), 3, c, -1, lineType=cv2.LINE_AA)
        cv2.circle(out, (x1 + margin + w0, y1 + sh), 3, c, -1, lineType=cv2.LINE_AA)

    return out


def preprocess(image: np.ndarray, grayscale: bool = False, resize_max: int = None, dfactor: int = 8):
    image = image.astype(np.float32, copy=False)
    size = image.shape[:2][::-1]
    scale = np.array([1.0, 1.0])

    if resize_max:
        scale = resize_max / max(size)
        if scale < 1.0:
            size_new = tuple(int(round(x * scale)) for x in size)
            image = resize_image(image, size_new, "cv2_area")
            scale = np.array(size) / np.array(size_new)

    if grayscale:
        assert image.ndim == 2, image.shape
        image = image[None]
    else:
        image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
    image = torch.from_numpy(image / 255.0).float()

    # assure that the size is divisible by dfactor
    size_new = tuple(map(lambda x: int(x // dfactor * dfactor), image.shape[-2:]))
    image = F.resize(image, size=size_new)
    scale = np.array(size) / np.array(size_new)[::-1]
    return image, scale


def compute_geom(
    data,
    ransac_method=DEFAULT_RANSAC_METHOD,
    ransac_reproj_threshold=DEFAULT_RANSAC_REPROJ_THRESHOLD,
    ransac_confidence=DEFAULT_RANSAC_CONFIDENCE,
    ransac_max_iter=DEFAULT_RANSAC_MAX_ITER,
) -> dict:

    mkpts0 = data["mkpts0_f"].cpu().detach().numpy()
    mkpts1 = data["mkpts1_f"].cpu().detach().numpy()

    if len(mkpts0) < 2 * DEFAULT_MIN_NUM_MATCHES:
        return {}

    h1, w1 = data["hw0_i"]

    geo_info = {}

    F, inliers = cv2.findFundamentalMat(
        mkpts0,
        mkpts1,
        method=RANSAC_ZOO[ransac_method],
        ransacReprojThreshold=ransac_reproj_threshold,
        confidence=ransac_confidence,
        maxIters=ransac_max_iter,
    )
    if F is not None:
        geo_info["Fundamental"] = F.tolist()

    H, _ = cv2.findHomography(
        mkpts1,
        mkpts0,
        method=RANSAC_ZOO[ransac_method],
        ransacReprojThreshold=ransac_reproj_threshold,
        confidence=ransac_confidence,
        maxIters=ransac_max_iter,
    )
    if H is not None:
        geo_info["Homography"] = H.tolist()
        _, H1, H2 = cv2.stereoRectifyUncalibrated(
            mkpts0.reshape(-1, 2),
            mkpts1.reshape(-1, 2),
            F,
            imgSize=(w1, h1),
        )
        geo_info["H1"] = H1.tolist()
        geo_info["H2"] = H2.tolist()

    return geo_info


def wrap_images(img0, img1, geo_info, geom_type):
    img0 = img0[0].permute((1, 2, 0)).cpu().detach().numpy()[..., ::-1]
    img1 = img1[0].permute((1, 2, 0)).cpu().detach().numpy()[..., ::-1]

    h1, w1, _ = img0.shape
    h2, w2, _ = img1.shape

    rectified_image0 = img0
    rectified_image1 = None
    H = np.array(geo_info["Homography"])
    F = np.array(geo_info["Fundamental"])

    title = []
    if geom_type == "Homography":
        rectified_image1 = cv2.warpPerspective(img1, H, (img0.shape[1], img0.shape[0]))
        title = ["Image 0", "Image 1 - warped"]
    elif geom_type == "Fundamental":
        H1, H2 = np.array(geo_info["H1"]), np.array(geo_info["H2"])
        rectified_image0 = cv2.warpPerspective(img0, H1, (w1, h1))
        rectified_image1 = cv2.warpPerspective(img1, H2, (w2, h2))
        title = ["Image 0 - warped", "Image 1 - warped"]
    else:
        print("Error: Unknown geometry type")

    fig = plot_images(
        [rectified_image0.squeeze(), rectified_image1.squeeze()],
        title,
        dpi=300,
    )

    img = fig2im(fig)

    plt.close(fig)

    return img


def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
    """Plot a set of images horizontally.
    Args:
        imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).
        titles: a list of strings, as titles for each image.
        cmaps: colormaps for monochrome images.
        dpi:
        size:
        pad:
    """
    n = len(imgs)
    if not isinstance(cmaps, (list, tuple)):
        cmaps = [cmaps] * n

    figsize = (size * n, size * 6 / 5) if size is not None else None
    fig, ax = plt.subplots(1, n, figsize=figsize, dpi=dpi)

    if n == 1:
        ax = [ax]
    for i in range(n):
        ax[i].imshow(imgs[i], cmap=plt.get_cmap(cmaps[i]))
        ax[i].get_yaxis().set_ticks([])
        ax[i].get_xaxis().set_ticks([])
        ax[i].set_axis_off()
        for spine in ax[i].spines.values():  # remove frame
            spine.set_visible(False)
        if titles:
            ax[i].set_title(titles[i])

    fig.tight_layout(pad=pad)

    return fig


def fig2im(fig):
    fig.canvas.draw()
    w, h = fig.canvas.get_width_height()
    buf_ndarray = np.frombuffer(fig.canvas.tostring_rgb(), dtype="u1")
    im = buf_ndarray.reshape(h, w, 3)
    return im


if __name__ == "__main__":
    model_zoo = ["gim_dkm", "gim_lightglue"]

    # model
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="gim_dkm", choices=model_zoo)
    args = parser.parse_args()

    # device
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # load model
    ckpt = None
    model = None
    detector = None
    if args.model == "gim_dkm":
        ckpt = "gim_dkm_100h.ckpt"
        model = DKMv3(weights=None, h=672, w=896)
    elif args.model == "gim_lightglue":
        ckpt = "gim_lightglue_100h.ckpt"
        detector = SuperPoint(
            {
                "max_num_keypoints": 2048,
                "force_num_keypoints": True,
                "detection_threshold": 0.0,
                "nms_radius": 3,
                "trainable": False,
            }
        )
        model = LightGlue(
            {
                "filter_threshold": 0.1,
                "flash": False,
                "checkpointed": True,
            }
        )

    # weights path
    checkpoints_path = join("weights", ckpt)

    # load state dict
    if args.model == "gim_dkm":
        state_dict = torch.load(checkpoints_path, map_location="cpu")
        if "state_dict" in state_dict.keys():
            state_dict = state_dict["state_dict"]
        for k in list(state_dict.keys()):
            if k.startswith("model."):
                state_dict[k.replace("model.", "", 1)] = state_dict.pop(k)
            if "encoder.net.fc" in k:
                state_dict.pop(k)
        model.load_state_dict(state_dict)

    elif args.model == "gim_lightglue":
        state_dict = torch.load(checkpoints_path, map_location="cpu")
        if "state_dict" in state_dict.keys():
            state_dict = state_dict["state_dict"]
        for k in list(state_dict.keys()):
            if k.startswith("model."):
                state_dict.pop(k)
            if k.startswith("superpoint."):
                state_dict[k.replace("superpoint.", "", 1)] = state_dict.pop(k)
        detector.load_state_dict(state_dict)

        state_dict = torch.load(checkpoints_path, map_location="cpu")
        if "state_dict" in state_dict.keys():
            state_dict = state_dict["state_dict"]
        for k in list(state_dict.keys()):
            if k.startswith("superpoint."):
                state_dict.pop(k)
            if k.startswith("model."):
                state_dict[k.replace("model.", "", 1)] = state_dict.pop(k)
        model.load_state_dict(state_dict)

    # eval mode
    if detector is not None:
        detector = detector.eval().to(device)
    model = model.eval().to(device)

    name0 = "a1"
    name1 = "a2"
    postfix = ".png"
    image_dir = join("assets", "demo")
    img_path0 = join(image_dir, name0 + postfix)
    img_path1 = join(image_dir, name1 + postfix)

    image0 = read_image(img_path0)
    image1 = read_image(img_path1)
    image0, scale0 = preprocess(image0)
    image1, scale1 = preprocess(image1)

    image0 = image0.to(device)[None]
    image1 = image1.to(device)[None]

    data = dict(color0=image0, color1=image1, image0=image0, image1=image1)

    if args.model == "gim_dkm":
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            dense_matches, dense_certainty = model.match(image0, image1)
            sparse_matches, mconf = model.sample(dense_matches, dense_certainty, 5000)

        height0, width0 = image0.shape[-2:]
        height1, width1 = image1.shape[-2:]

        kpts0 = sparse_matches[:, :2]
        kpts0 = torch.stack(
            (width0 * (kpts0[:, 0] + 1) / 2, height0 * (kpts0[:, 1] + 1) / 2),
            dim=-1,
        )
        kpts1 = sparse_matches[:, 2:]
        kpts1 = torch.stack(
            (width1 * (kpts1[:, 0] + 1) / 2, height1 * (kpts1[:, 1] + 1) / 2),
            dim=-1,
        )
        b_ids = torch.where(mconf[None])[0]
    elif args.model == "gim_lightglue":
        gray0 = read_image(img_path0, grayscale=True)
        gray1 = read_image(img_path1, grayscale=True)
        gray0 = preprocess(gray0, grayscale=True)[0]
        gray1 = preprocess(gray1, grayscale=True)[0]

        gray0 = gray0.to(device)[None]
        gray1 = gray1.to(device)[None]
        scale0 = torch.tensor(scale0).to(device)[None]
        scale1 = torch.tensor(scale1).to(device)[None]

        data.update(dict(gray0=gray0, gray1=gray1))

        size0 = torch.tensor(data["gray0"].shape[-2:][::-1])[None]
        size1 = torch.tensor(data["gray1"].shape[-2:][::-1])[None]

        data.update(dict(size0=size0, size1=size1))
        data.update(dict(scale0=scale0, scale1=scale1))

        pred = {}
        pred.update(
            {
                k + "0": v
                for k, v in detector(
                    {
                        "image": data["gray0"],
                        "image_size": data["size0"],
                    }
                ).items()
            }
        )
        pred.update(
            {
                k + "1": v
                for k, v in detector(
                    {
                        "image": data["gray1"],
                        "image_size": data["size1"],
                    }
                ).items()
            }
        )
        pred.update(model({**pred, **data, **{"resize0": data["size0"], "resize1": data["size1"]}}))

        kpts0 = torch.cat([kp * s for kp, s in zip(pred["keypoints0"], data["scale0"][:, None])])
        kpts1 = torch.cat([kp * s for kp, s in zip(pred["keypoints1"], data["scale1"][:, None])])
        m_bids = torch.nonzero(pred["keypoints0"].sum(dim=2) > -1)[:, 0]
        matches = pred["matches"]
        bs = data["image0"].size(0)
        kpts0 = torch.cat([kpts0[m_bids == b_id][matches[b_id][..., 0]] for b_id in range(bs)])
        kpts1 = torch.cat([kpts1[m_bids == b_id][matches[b_id][..., 1]] for b_id in range(bs)])
        b_ids = torch.cat([m_bids[m_bids == b_id][matches[b_id][..., 0]] for b_id in range(bs)])
        mconf = torch.cat(pred["scores"])

    # robust fitting
    _, mask = cv2.findFundamentalMat(
        kpts0.cpu().detach().numpy(),
        kpts1.cpu().detach().numpy(),
        cv2.USAC_MAGSAC,
        ransacReprojThreshold=1.0,
        confidence=0.999999,
        maxIters=10000,
    )
    mask = mask.ravel() > 0

    data.update(
        {
            "hw0_i": image0.shape[-2:],
            "hw1_i": image1.shape[-2:],
            "mkpts0_f": kpts0,
            "mkpts1_f": kpts1,
            "m_bids": b_ids,
            "mconf": mconf,
            "inliers": mask,
        }
    )

    # save visualization
    alpha = 0.5
    out = fast_make_matching_figure(data, b_id=0)
    overlay = fast_make_matching_overlay(data, b_id=0)
    out = cv2.addWeighted(out, 1 - alpha, overlay, alpha, 0)
    cv2.imwrite(join(image_dir, f"{name0}_{name1}_{args.model}_match.png"), out[..., ::-1])

    geom_info = compute_geom(data)
    wrapped_images = wrap_images(image0, image1, geom_info, "Homography")
    cv2.imwrite(join(image_dir, f"{name0}_{name1}_{args.model}_warp.png"), wrapped_images)
