"""
"LiftFeat: 3D Geometry-Aware Local Feature Matching"
"""

import numpy as np
import os
import torch
from torch import nn
import torch.nn.functional as F

import tqdm
import math
import cv2

import sys

sys.path.append("/home/yepeng_liu/code_python/laiwenpeng/LiftFeat")
from utils.featurebooster import FeatureBooster
from utils.config import featureboost_config

# from models.model_dfb import LiftFeatModel
# from models.interpolator import InterpolateSparse2d
# from third_party.config import featureboost_config

"""
foundational functions
"""


def simple_nms(scores, radius):
    """Perform non maximum suppression on the heatmap using max-pooling.
    This method does not suppress contiguous points that have the same score.
    Args:
        scores: the score heatmap of size `(B, H, W)`.
        radius: an integer scalar, the radius of the NMS window.
    """

    def max_pool(x):
        return torch.nn.functional.max_pool2d(x, kernel_size=radius * 2 + 1, stride=1, padding=radius)

    zeros = torch.zeros_like(scores)
    max_mask = scores == max_pool(scores)
    for _ in range(2):
        supp_mask = max_pool(max_mask.float()) > 0
        supp_scores = torch.where(supp_mask, zeros, scores)
        new_max_mask = supp_scores == max_pool(supp_scores)
        max_mask = max_mask | (new_max_mask & (~supp_mask))
    return torch.where(max_mask, scores, zeros)


def top_k_keypoints(keypoints, scores, k):
    if k >= len(keypoints):
        return keypoints, scores
    scores, indices = torch.topk(scores, k, dim=0, sorted=True)
    return keypoints[indices], scores


def sample_k_keypoints(keypoints, scores, k):
    if k >= len(keypoints):
        return keypoints, scores
    indices = torch.multinomial(scores, k, replacement=False)
    return keypoints[indices], scores[indices]


def soft_argmax_refinement(keypoints, scores, radius: int):
    width = 2 * radius + 1
    sum_ = torch.nn.functional.avg_pool2d(scores[:, None], width, 1, radius, divisor_override=1)
    ar = torch.arange(-radius, radius + 1).to(scores)
    kernel_x = ar[None].expand(width, -1)[None, None]
    dx = torch.nn.functional.conv2d(scores[:, None], kernel_x, padding=radius)
    dy = torch.nn.functional.conv2d(scores[:, None], kernel_x.transpose(2, 3), padding=radius)
    dydx = torch.stack([dy[:, 0], dx[:, 0]], -1) / sum_[:, 0, :, :, None]
    refined_keypoints = []
    for i, kpts in enumerate(keypoints):
        delta = dydx[i][tuple(kpts.t())]
        refined_keypoints.append(kpts.float() + delta)
    return refined_keypoints


# Legacy (broken) sampling of the descriptors
def sample_descriptors(keypoints, descriptors, s):
    b, c, h, w = descriptors.shape
    keypoints = keypoints - s / 2 + 0.5
    keypoints /= torch.tensor(
        [(w * s - s / 2 - 0.5), (h * s - s / 2 - 0.5)],
    ).to(
        keypoints
    )[None]
    keypoints = keypoints * 2 - 1  # normalize to (-1, 1)
    args = {"align_corners": True} if torch.__version__ >= "1.3" else {}
    descriptors = torch.nn.functional.grid_sample(descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", **args)
    descriptors = torch.nn.functional.normalize(descriptors.reshape(b, c, -1), p=2, dim=1)
    return descriptors


# The original keypoint sampling is incorrect. We patch it here but
# keep the original one above for legacy.
def sample_descriptors_fix_sampling(keypoints, descriptors, s: int = 8):
    """Interpolate descriptors at keypoint locations"""
    b, c, h, w = descriptors.shape
    keypoints = keypoints / (keypoints.new_tensor([w, h]) * s)
    keypoints = keypoints * 2 - 1  # normalize to (-1, 1)
    descriptors = torch.nn.functional.grid_sample(
        descriptors, keypoints.view(b, 1, -1, 2), mode="bilinear", align_corners=False
    )
    descriptors = torch.nn.functional.normalize(descriptors.reshape(b, c, -1), p=2, dim=1)
    return descriptors


class UpsampleLayer(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        # 定义特征提取层，减少通道数同时增加特征提取能力
        self.conv = nn.Conv2d(in_channels, in_channels // 2, kernel_size=3, stride=1, padding=1)
        # 使用BN层
        self.bn = nn.BatchNorm2d(in_channels // 2)
        # 使用LeakyReLU激活函数
        self.leaky_relu = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
        x = self.leaky_relu(self.bn(self.conv(x)))

        return x


class KeypointHead(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.layer1 = BaseLayer(in_channels, 32)
        self.layer2 = BaseLayer(32, 32)
        self.layer3 = BaseLayer(32, 64)
        self.layer4 = BaseLayer(64, 64)
        self.layer5 = BaseLayer(64, 128)

        self.conv = nn.Conv2d(128, out_channels, kernel_size=3, stride=1, padding=1)
        self.bn = nn.BatchNorm2d(65)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.bn(self.conv(x))
        return x


class DescriptorHead(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.layer = nn.Sequential(
            BaseLayer(in_channels, 32),
            BaseLayer(32, 32, activation=False),
            BaseLayer(32, 64, activation=False),
            BaseLayer(64, out_channels, activation=False),
        )

    def forward(self, x):
        x = self.layer(x)
        # x=nn.functional.softmax(x,dim=1)
        return x


class HeatmapHead(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels):
        super().__init__()
        self.convHa = nn.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, padding=1)
        self.bnHa = nn.BatchNorm2d(mid_channels)
        self.convHb = nn.Conv2d(mid_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.bnHb = nn.BatchNorm2d(out_channels)
        self.leaky_relu = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.leaky_relu(self.bnHa(self.convHa(x)))
        x = self.leaky_relu(self.bnHb(self.convHb(x)))

        x = torch.sigmoid(x)
        return x


class DepthHead(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.upsampleDa = UpsampleLayer(in_channels)
        self.upsampleDb = UpsampleLayer(in_channels // 2)
        self.upsampleDc = UpsampleLayer(in_channels // 4)

        self.convDepa = nn.Conv2d(in_channels // 2 + in_channels, in_channels // 2, kernel_size=3, stride=1, padding=1)
        self.bnDepa = nn.BatchNorm2d(in_channels // 2)
        self.convDepb = nn.Conv2d(
            in_channels // 4 + in_channels // 2, in_channels // 4, kernel_size=3, stride=1, padding=1
        )
        self.bnDepb = nn.BatchNorm2d(in_channels // 4)
        self.convDepc = nn.Conv2d(in_channels // 8 + in_channels // 4, 3, kernel_size=3, stride=1, padding=1)
        self.bnDepc = nn.BatchNorm2d(3)

        self.leaky_relu = nn.LeakyReLU(0.1)

    def forward(self, x):
        x0 = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
        x1 = self.upsampleDa(x)
        x1 = torch.cat([x0, x1], dim=1)
        x1 = self.leaky_relu(self.bnDepa(self.convDepa(x1)))

        x1_0 = F.interpolate(x1, scale_factor=2, mode="bilinear", align_corners=False)
        x2 = self.upsampleDb(x1)
        x2 = torch.cat([x1_0, x2], dim=1)
        x2 = self.leaky_relu(self.bnDepb(self.convDepb(x2)))

        x2_0 = F.interpolate(x2, scale_factor=2, mode="bilinear", align_corners=False)
        x3 = self.upsampleDc(x2)
        x3 = torch.cat([x2_0, x3], dim=1)
        x = self.leaky_relu(self.bnDepc(self.convDepc(x3)))

        x = F.normalize(x, p=2, dim=1)
        return x


class BaseLayer(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False, activation=True):
        super().__init__()
        if activation:
            self.layer = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
                nn.BatchNorm2d(out_channels, affine=False),
                nn.ReLU(inplace=True),
            )
        else:
            self.layer = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
                nn.BatchNorm2d(out_channels, affine=False),
            )

    def forward(self, x):
        return self.layer(x)


class LiftFeatSPModel(nn.Module):
    default_conf = {
        "has_detector": True,
        "has_descriptor": True,
        "descriptor_dim": 64,
        # Inference
        "sparse_outputs": True,
        "dense_outputs": False,
        "nms_radius": 4,
        "refinement_radius": 0,
        "detection_threshold": 0.005,
        "max_num_keypoints": -1,
        "max_num_keypoints_val": None,
        "force_num_keypoints": False,
        "randomize_keypoints_training": False,
        "remove_borders": 4,
        "legacy_sampling": True,  # True to use the old broken sampling
    }

    def __init__(self, featureboost_config, use_kenc=False, use_normal=True, use_cross=True):
        super().__init__()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.descriptor_dim = 64

        self.norm = nn.InstanceNorm2d(1)

        self.relu = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        c1, c2, c3, c4, c5 = 24, 24, 64, 64, 128

        self.conv1a = nn.Conv2d(1, c1, kernel_size=3, stride=1, padding=1)
        self.conv1b = nn.Conv2d(c1, c1, kernel_size=3, stride=1, padding=1)
        self.conv2a = nn.Conv2d(c1, c2, kernel_size=3, stride=1, padding=1)
        self.conv2b = nn.Conv2d(c2, c2, kernel_size=3, stride=1, padding=1)
        self.conv3a = nn.Conv2d(c2, c3, kernel_size=3, stride=1, padding=1)
        self.conv3b = nn.Conv2d(c3, c3, kernel_size=3, stride=1, padding=1)
        self.conv4a = nn.Conv2d(c3, c4, kernel_size=3, stride=1, padding=1)
        self.conv4b = nn.Conv2d(c4, c4, kernel_size=3, stride=1, padding=1)
        self.conv5a = nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)
        self.conv5b = nn.Conv2d(c5, c5, kernel_size=3, stride=1, padding=1)

        self.upsample4 = UpsampleLayer(c4)
        self.upsample5 = UpsampleLayer(c5)
        self.conv_fusion45 = nn.Conv2d(c5 // 2 + c4, c4, kernel_size=3, stride=1, padding=1)
        self.conv_fusion34 = nn.Conv2d(c4 // 2 + c3, c3, kernel_size=3, stride=1, padding=1)

        # detector
        self.keypoint_head = KeypointHead(in_channels=c3, out_channels=65)
        # descriptor
        self.descriptor_head = DescriptorHead(in_channels=c3, out_channels=self.descriptor_dim)
        # # heatmap
        # self.heatmap_head = HeatmapHead(in_channels=c3,mid_channels=c3,out_channels=1)
        # depth
        self.depth_head = DepthHead(c3)

        self.fine_matcher = nn.Sequential(
            nn.Linear(128, 512),
            nn.BatchNorm1d(512, affine=False),
            nn.ReLU(inplace=True),
            nn.Linear(512, 512),
            nn.BatchNorm1d(512, affine=False),
            nn.ReLU(inplace=True),
            nn.Linear(512, 512),
            nn.BatchNorm1d(512, affine=False),
            nn.ReLU(inplace=True),
            nn.Linear(512, 512),
            nn.BatchNorm1d(512, affine=False),
            nn.ReLU(inplace=True),
            nn.Linear(512, 64),
        )

        # feature_booster
        self.feature_boost = FeatureBooster(
            featureboost_config, use_kenc=use_kenc, use_cross=use_cross, use_normal=use_normal
        )

    def feature_extract(self, x):
        x1 = self.relu(self.conv1a(x))
        x1 = self.relu(self.conv1b(x1))
        x1 = self.pool(x1)
        x2 = self.relu(self.conv2a(x1))
        x2 = self.relu(self.conv2b(x2))
        x2 = self.pool(x2)
        x3 = self.relu(self.conv3a(x2))
        x3 = self.relu(self.conv3b(x3))
        x3 = self.pool(x3)
        x4 = self.relu(self.conv4a(x3))
        x4 = self.relu(self.conv4b(x4))
        x4 = self.pool(x4)
        x5 = self.relu(self.conv5a(x4))
        x5 = self.relu(self.conv5b(x5))
        x5 = self.pool(x5)
        return x3, x4, x5

    def fuse_multi_features(self, x3, x4, x5):
        # upsample x5 feature
        x5 = self.upsample5(x5)
        x4 = torch.cat([x4, x5], dim=1)
        x4 = self.conv_fusion45(x4)

        # upsample x4 feature
        x4 = self.upsample4(x4)
        x3 = torch.cat([x3, x4], dim=1)
        x = self.conv_fusion34(x3)
        return x

    def _unfold2d(self, x, ws=2):
        """
        Unfolds tensor in 2D with desired ws (window size) and concat the channels
        """
        B, C, H, W = x.shape
        x = x.unfold(2, ws, ws).unfold(3, ws, ws).reshape(B, C, H // ws, W // ws, ws**2)
        return x.permute(0, 1, 4, 2, 3).reshape(B, -1, H // ws, W // ws)

    def forward1(self, x):
        """
        input:
            x -> torch.Tensor(B, C, H, W) grayscale or rgb images
        return:
            feats     ->  torch.Tensor(B, 64, H/8, W/8) dense local features
            keypoints ->  torch.Tensor(B, 65, H/8, W/8) keypoint logit map
            heatmap   ->  torch.Tensor(B,  1, H/8, W/8) reliability map

        """
        with torch.no_grad():
            x = x.mean(dim=1, keepdim=True)
            x = self.norm(x)

        x3, x4, x5 = self.feature_extract(x)

        # features fusion
        x = self.fuse_multi_features(x3, x4, x5)

        # keypoint
        keypoint_map = self.keypoint_head(x)
        # descriptor
        des_map = self.descriptor_head(x)
        # # heatmap
        # heatmap = self.heatmap_head(x)

        # import pdb;pdb.set_trace()
        # depth
        d_feats = self.depth_head(x)

        return des_map, keypoint_map, d_feats
        # return des_map, keypoint_map, heatmap, d_feats

    def forward2(self, descs, kpts, normals):
        # import pdb;pdb.set_trace()
        normals_feat = self._unfold2d(normals, ws=8)
        normals_v = normals_feat.squeeze(0).permute(1, 2, 0).reshape(-1, normals_feat.shape[1])
        descs_v = descs.squeeze(0).permute(1, 2, 0).reshape(-1, descs.shape[1])
        kpts_v = kpts.squeeze(0).permute(1, 2, 0).reshape(-1, kpts.shape[1])
        descs_refine = self.feature_boost(descs_v, kpts_v, normals_v)
        return descs_refine

    def forward(self, x):
        M1, K1, D1 = self.forward1(x)
        descs_refine = self.forward2(M1, K1, D1)
        return descs_refine, M1, K1, D1


if __name__ == "__main__":
    img_path = os.path.join(os.path.dirname(__file__), "../assert/ref.jpg")
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, (800, 608))
    import pdb

    pdb.set_trace()
    img = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float() / 255.0
    img = img.cuda() if torch.cuda.is_available() else img
    liftfeat_sp = LiftFeatSPModel(featureboost_config).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
    des_map, keypoint_map, d_feats = liftfeat_sp.forward1(img)
    des_fine = liftfeat_sp.forward2(des_map, keypoint_map, d_feats)
    print(des_map.shape)
