"""latest version of SuperpointNet. Use it!

"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import xavier_uniform_, zeros_

from models.unet_parts import *
import numpy as np

# from models.SubpixelNet import SubpixelNet
class SuperPointNet_gauss2_ssmall(torch.nn.Module):
    """Pytorch definition of SuperPoint Network."""

    def __init__(self, n_classes=133, subpixel_channel=1):
        super(SuperPointNet_gauss2_ssmall, self).__init__()
        c1, c2, c3, c4, c5, d1 = 64, 64, 128, 128, 256, 256
        det_h = 65
        self.inc = inconv(1, c1)
        self.down1 = down(c1, c2)
        self.down2 = down(c2, c3)
        self.down3 = down(c3, c4)

        # self.down4 = down(c4, 512)
        # self.up1 = up(c4+c3, c2)
        # self.up2 = up(c2+c2, c1)
        # self.up3 = up(c1+c1, c1)
        # self.outc = outconv(c1, subpixel_channel)
        self.relu = torch.nn.ReLU(inplace=True)
        # self.outc = outconv(64, n_classes)

        # Detector Head.
        self.convPa = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)
        self.bnPa = nn.BatchNorm2d(c5)
        self.convPb = torch.nn.Conv2d(c5, det_h, kernel_size=1, stride=1, padding=0)
        self.bnPb = nn.BatchNorm2d(det_h)

        # Descriptor Head.
        self.convDa = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)
        self.bnDa = nn.BatchNorm2d(c5)
        self.convDb = torch.nn.Conv2d(c5, d1, kernel_size=1, stride=1, padding=0)
        self.bnDb = nn.BatchNorm2d(d1)

        # Seg Head.
        self.convDS = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)
        self.bnS1 = nn.BatchNorm2d(c5)
        self.convSout = torch.nn.Conv2d(c5, n_classes, kernel_size=1, stride=1, padding=0)

        # self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
        # self.convS1 = double_conv(c5, c4)
        # self.convS2 = double_conv(c4, c3)
        # self.convSout = torch.nn.Conv2d(c3, n_classes, kernel_size=(1, 1), padding=0)

        self.output = None

    def forward(self, x, subpixel=False):
        # def forward(self, x): #uncomment this to extract the scripted version
        """Forward pass that jointly computes unprocessed point and descriptor
        tensors.
        Input
          x: Image pytorch tensor shaped N x 1 x patch_size x patch_size.
        Output
          semi: Output point pytorch tensor shaped N x 65 x H/8 x W/8.
          desc: Output descriptor pytorch tensor shaped N x 256 x H/8 x W/8.
          sem: Output semantic segmentation tensor shape N x N_CLASSES x H/2 x H/2.
        """

        x_hw = x.shape[2:]

        # Let's stick to this version: first BN, then relu
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)

        # Detector Head.
        cPa = self.relu(self.bnPa(self.convPa(x4)))
        semi = self.bnPb(self.convPb(cPa))

        # Descriptor Head.
        cDa = self.relu(self.bnDa(self.convDa(x4)))
        desc = self.bnDb(self.convDb(cDa))

        # Sem Head
        sem = self.convDS(x4)
        sem = self.bnS1(sem)
        sem = self.relu(sem)
        sem = self.convSout(sem)
        sem = F.interpolate(sem, x_hw, mode="bilinear", align_corners=False)
        # sem = nn.Softmax(dim=1)(sem)

        dn = torch.norm(desc, p=2, dim=1)  # Compute the norm.
        desc = desc.div(torch.unsqueeze(dn, 1))  # Divide by norm to normalize.
        output = {"semi": semi, "desc": desc, "sem": sem}
        # output = [semi, desc] #uncomment this to extract the scripted version

        return output

    def removeSem(self):
        del self.convDS
        del self.bnS1
        del self.convSout

    def process_output(self, sp_processer):
        """
        input:
          N: number of points
        return: -- type: tensorFloat
          pts: tensor [batch, N, 2] (no grad)  (x, y)
          pts_offset: tensor [batch, N, 2] (grad) (x, y)
          pts_desc: tensor [batch, N, 256] (grad)
        """
        from utils.utils import flattenDetection

        # from models.model_utils import pred_soft_argmax, sample_desc_from_points
        output = self.output
        semi = output["semi"]
        desc = output["desc"]
        sem = output["sem"]
        # flatten
        heatmap = flattenDetection(semi)  # [batch_size, 1, H, W]
        # nms
        heatmap_nms_batch = sp_processer.heatmap_to_nms(heatmap, tensor=True)
        # extract offsets
        outs = sp_processer.pred_soft_argmax(heatmap_nms_batch, heatmap)
        residual = outs["pred"]
        # extract points
        outs = sp_processer.batch_extract_features(desc, heatmap_nms_batch, residual)
        outs["sem"] = sem

        # output.update({'heatmap': heatmap, 'heatmap_nms': heatmap_nms, 'descriptors': descriptors})
        output.update(outs)
        self.output = output
        return output


def get_matches(deses_SP):
    from models.model_wrap import PointTracker

    tracker = PointTracker(max_length=2, nn_thresh=1.2)
    f = lambda x: x.cpu().detach().numpy()
    # tracker = PointTracker(max_length=2, nn_thresh=1.2)
    # print("deses_SP[1]: ", deses_SP[1].shape)
    matching_mask = tracker.nn_match_two_way(f(deses_SP[0]).T, f(deses_SP[1]).T, nn_thresh=1.2)
    return matching_mask

    # print("matching_mask: ", matching_mask.shape)
    # f_mask = lambda pts, maks: pts[]
    # pts_m = []
    # pts_m_res = []
    # for i in range(2):
    #     idx = xs_SP[i][matching_mask[i, :].astype(int), :]
    #     res = reses_SP[i][matching_mask[i, :].astype(int), :]
    #     print("idx: ", idx.shape)
    #     print("res: ", idx.shape)
    #     pts_m.append(idx)
    #     pts_m_res.append(res)
    #     pass

    # pts_m = torch.cat((pts_m[0], pts_m[1]), dim=1)
    # matches_test = toNumpy(pts_m)
    # print("pts_m: ", pts_m.shape)

    # pts_m_res = torch.cat((pts_m_res[0], pts_m_res[1]), dim=1)
    # # pts_m_res = toNumpy(pts_m_res)
    # print("pts_m_res: ", pts_m_res.shape)
    # # print("pts_m_res: ", pts_m_res)

    # pts_idx_res = torch.cat((pts_m, pts_m_res), dim=1)
    # print("pts_idx_res: ", pts_idx_res.shape)


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SuperPointNet_gauss2_unetV1_sem()
    model = model.to(device)

    # check keras-like model summary using torchsummary
    from torchsummary import summary

    summary(model, input_size=(1, 240, 320))

    ## test
    image = torch.zeros((2, 1, 120, 160))
    outs = model(image.to(device))

    from utils.print_tool import print_dict_attr

    print_dict_attr(outs, "shape")

    from models.model_utils import SuperPointNet_process

    params = {"out_num_points": 500, "patch_size": 5, "device": device, "nms_dist": 4, "conf_thresh": 0.015}

    sp_processer = SuperPointNet_process(**params)
    outs = model.process_output(sp_processer)
    print("outs: ", list(outs))
    print_dict_attr(outs, "shape")

    # timer
    import time
    from tqdm import tqdm

    iter_max = 50

    start = time.time()
    print("Start timer!")
    for i in tqdm(range(iter_max)):
        outs = model(image.to(device))
    end = time.time()
    print("forward only: ", iter_max / (end - start), " iter/s")

    start = time.time()
    print("Start timer!")
    xs_SP, deses_SP, reses_SP = [], [], []
    for i in tqdm(range(iter_max)):
        outs = model(image.to(device))
        outs = model.process_output(sp_processer)
        xs_SP.append(outs["pts_int"].squeeze())
        deses_SP.append(outs["pts_desc"].squeeze())
        reses_SP.append(outs["pts_offset"].squeeze())
    end = time.time()
    print("forward + process output: ", iter_max / (end - start), " iter/s")

    start = time.time()
    print("Start timer!")
    for i in tqdm(range(len(xs_SP))):
        get_matches([deses_SP[i][0], deses_SP[i][1]])
    end = time.time()
    print("nn matches: ", iter_max / (end - start), " iters/s")


if __name__ == "__main__":
    main()
