# -*- coding: UTF-8 -*-
'''=================================================
@Project -> File   pram -> gm
@IDE    PyCharm
@Author fx221@cam.ac.uk
@Date   07/02/2024 10:47
=================================================='''
import torch
import torch.nn as nn
import torch.nn.functional as F
from nets.layers import KeypointEncoder, AttentionalPropagation
from nets.utils import normalize_keypoints, arange_like

eps = 1e-8


def dual_softmax(M, dustbin):
    M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1)
    M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2)
    score = torch.log_softmax(M, dim=-1) + torch.log_softmax(M, dim=1)
    return torch.exp(score)


def sinkhorn(M, r, c, iteration):
    p = torch.softmax(M, dim=-1)
    u = torch.ones_like(r)
    v = torch.ones_like(c)
    for _ in range(iteration):
        u = r / ((p * v.unsqueeze(-2)).sum(-1) + eps)
        v = c / ((p * u.unsqueeze(-1)).sum(-2) + eps)
    p = p * u.unsqueeze(-1) * v.unsqueeze(-2)
    return p


def sink_algorithm(M, dustbin, iteration):
    M = torch.cat([M, dustbin.expand([M.shape[0], M.shape[1], 1])], dim=-1)
    M = torch.cat([M, dustbin.expand([M.shape[0], 1, M.shape[2]])], dim=-2)
    r = torch.ones([M.shape[0], M.shape[1] - 1], device='cuda')
    r = torch.cat([r, torch.ones([M.shape[0], 1], device='cuda') * M.shape[1]], dim=-1)
    c = torch.ones([M.shape[0], M.shape[2] - 1], device='cuda')
    c = torch.cat([c, torch.ones([M.shape[0], 1], device='cuda') * M.shape[2]], dim=-1)
    p = sinkhorn(M, r, c, iteration)
    return p


class AttentionalGNN(nn.Module):
    def __init__(self, feature_dim: int, layer_names: list, hidden_dim: int = 256, ac_fn: str = 'relu',
                 norm_fn: str = 'bn'):
        super().__init__()
        self.layers = nn.ModuleList([
            AttentionalPropagation(feature_dim=feature_dim, num_heads=4, hidden_dim=hidden_dim, ac_fn=ac_fn,
                                   norm_fn=norm_fn)
            for _ in range(len(layer_names))])
        self.names = layer_names

    def forward(self, desc0, desc1):
        # desc0s = []
        # desc1s = []

        for i, (layer, name) in enumerate(zip(self.layers, self.names)):
            if name == 'cross':
                src0, src1 = desc1, desc0
            else:
                src0, src1 = desc0, desc1
            delta0 = layer(desc0, src0)
            # prob0 = layer.attn.prob
            delta1 = layer(desc1, src1)
            # prob1 = layer.attn.prob
            desc0, desc1 = (desc0 + delta0), (desc1 + delta1)

            # if name == 'cross':
            #     desc0s.append(desc0)
            #     desc1s.append(desc1)
        return [desc0], [desc1]

    def predict(self, desc0, desc1, n_it=-1):
        for i, (layer, name) in enumerate(zip(self.layers, self.names)):
            if name == 'cross':
                src0, src1 = desc1, desc0
            else:
                src0, src1 = desc0, desc1
            delta0 = layer(desc0, src0)
            # prob0 = layer.attn.prob
            delta1 = layer(desc1, src1)
            # prob1 = layer.attn.prob
            desc0, desc1 = (desc0 + delta0), (desc1 + delta1)

            if name == 'cross' and i == n_it:
                break
        return [desc0], [desc1]


class GM(nn.Module):
    default_config = {
        'descriptor_dim': 128,
        'hidden_dim': 256,
        'keypoint_encoder': [32, 64, 128, 256],
        'GNN_layers': ['self', 'cross'] * 9,  # [self, cross, self, cross, ...] 9 in total
        'sinkhorn_iterations': 20,
        'match_threshold': 0.2,
        'with_pose': False,
        'n_layers': 9,
        'n_min_tokens': 256,
        'with_sinkhorn': True,

        'ac_fn': 'relu',
        'norm_fn': 'bn',
        'weight_path': None,
    }

    required_inputs = [
        'image0', 'keypoints0', 'scores0', 'descriptors0',
        'image1', 'keypoints1', 'scores1', 'descriptors1',
    ]

    def __init__(self, config):
        super().__init__()
        self.config = {**self.default_config, **config}
        print('gm: ', self.config)

        self.n_layers = self.config['n_layers']

        self.with_sinkhorn = self.config['with_sinkhorn']
        self.match_threshold = self.config['match_threshold']

        self.sinkhorn_iterations = self.config['sinkhorn_iterations']
        self.kenc = KeypointEncoder(
            self.config['descriptor_dim'] if self.config['descriptor_dim'] > 0 else 128,
            self.config['keypoint_encoder'],
            ac_fn=self.config['ac_fn'],
            norm_fn=self.config['norm_fn'])
        self.gnn = AttentionalGNN(
            feature_dim=self.config['descriptor_dim'] if self.config['descriptor_dim'] > 0 else 128,
            hidden_dim=self.config['hidden_dim'],
            layer_names=self.config['GNN_layers'],
            ac_fn=self.config['ac_fn'],
            norm_fn=self.config['norm_fn'],
        )

        self.final_proj = nn.ModuleList([nn.Conv1d(
            self.config['descriptor_dim'] if self.config['descriptor_dim'] > 0 else 128,
            self.config['descriptor_dim'] if self.config['descriptor_dim'] > 0 else 128,
            kernel_size=1, bias=True) for _ in range(self.n_layers)])

        bin_score = torch.nn.Parameter(torch.tensor(1.))
        self.register_parameter('bin_score', bin_score)

        self.match_net = None  # GraphLoss(config=self.config)

        self.self_prob0 = None
        self.self_prob1 = None
        self.cross_prob0 = None
        self.cross_prob1 = None

        self.desc_compressor = None

    def forward_train(self, data):
        pass

    def produce_matches(self, data, p=0.2, n_it=-1, **kwargs):
        kpts0, kpts1 = data['keypoints0'], data['keypoints1']
        scores0, scores1 = data['scores0'], data['scores1']
        if kpts0.shape[1] == 0 or kpts1.shape[1] == 0:  # no keypoints
            shape0, shape1 = kpts0.shape[:-1], kpts1.shape[:-1]
            return {
                'matches0': kpts0.new_full(shape0, -1, dtype=torch.int)[0],
                'matches1': kpts1.new_full(shape1, -1, dtype=torch.int)[0],
                'matching_scores0': kpts0.new_zeros(shape0)[0],
                'matching_scores1': kpts1.new_zeros(shape1)[0],
                'skip_train': True
            }

        if 'norm_keypoints0' in data.keys() and 'norm_keypoints1' in data.keys():
            norm_kpts0 = data['norm_keypoints0']
            norm_kpts1 = data['norm_keypoints1']
        elif 'image0' in data.keys() and 'image1' in data.keys():
            norm_kpts0 = normalize_keypoints(kpts0, data['image0'].shape)
            norm_kpts1 = normalize_keypoints(kpts1, data['image1'].shape)
        elif 'image_shape0' in data.keys() and 'image_shape1' in data.keys():
            norm_kpts0 = normalize_keypoints(kpts0, data['image_shape0'])
            norm_kpts1 = normalize_keypoints(kpts1, data['image_shape1'])
        else:
            raise ValueError('Require image shape for keypoint coordinate normalization')

        # Keypoint MLP encoder.
        enc0, enc1 = self.encode_keypoint(norm_kpts0=norm_kpts0, norm_kpts1=norm_kpts1, scores0=scores0,
                                          scores1=scores1)

        if self.config['descriptor_dim'] > 0:
            desc0, desc1 = data['descriptors0'], data['descriptors1']
            desc0 = desc0.transpose(0, 2, 1)  # [B, N, D ] -> [B, D, N]
            desc1 = desc1.transpose(0, 2, 1)  # [B, N, D ] -> [B, D, N]
            with torch.no_grad():
                if desc0.shape[1] != self.config['descriptor_dim']:
                    desc0 = self.desc_compressor(desc0)
                if desc1.shape[1] != self.config['descriptor_dim']:
                    desc1 = self.desc_compressor(desc1)
            desc0 = desc0 + enc0
            desc1 = desc1 + enc1
        else:
            desc0 = enc0
            desc1 = enc1

        desc0s, desc1s = self.gnn.predict(desc0, desc1, n_it=n_it)

        mdescs0 = self.final_proj[n_it](desc0s[-1])
        mdescs1 = self.final_proj[n_it](desc1s[-1])
        dist = torch.einsum('bdn,bdm->bnm', mdescs0, mdescs1)
        if self.config['descriptor_dim'] > 0:
            dist = dist / self.config['descriptor_dim'] ** .5
        else:
            dist = dist / 128 ** .5
        score = self.compute_score(dist=dist, dustbin=self.bin_score, iteration=self.sinkhorn_iterations)

        indices0, indices1, mscores0, mscores1 = self.compute_matches(scores=score, p=p)

        output = {
            'matches0': indices0,  # use -1 for invalid match
            'matches1': indices1,  # use -1 for invalid match
            'matching_scores0': mscores0,
            'matching_scores1': mscores1,
        }

        return output

    def forward(self, data, mode=0):
        if not self.training:
            return self.produce_matches(data=data, n_it=-1)
        return self.forward_train(data=data)

    def encode_keypoint(self, norm_kpts0, norm_kpts1, scores0, scores1):
        return self.kenc(norm_kpts0, scores0), self.kenc(norm_kpts1, scores1)

    def compute_distance(self, desc0, desc1, layer_id=-1):
        mdesc0 = self.final_proj[layer_id](desc0)
        mdesc1 = self.final_proj[layer_id](desc1)
        dist = torch.einsum('bdn,bdm->bnm', mdesc0, mdesc1)
        dist = dist / self.config['descriptor_dim'] ** .5
        return dist

    def compute_score(self, dist, dustbin, iteration):
        if self.with_sinkhorn:
            score = sink_algorithm(M=dist, dustbin=dustbin,
                                   iteration=iteration)  # [nI * nB, N, M]
        else:
            score = dual_softmax(M=dist, dustbin=dustbin)
        return score

    def compute_matches(self, scores, p=0.2):
        max0, max1 = scores[:, :-1, :-1].max(2), scores[:, :-1, :-1].max(1)
        indices0, indices1 = max0.indices, max1.indices
        mutual0 = arange_like(indices0, 1)[None] == indices1.gather(1, indices0)
        mutual1 = arange_like(indices1, 1)[None] == indices0.gather(1, indices1)
        zero = scores.new_tensor(0)
        # mscores0 = torch.where(mutual0, max0.values.exp(), zero)
        mscores0 = torch.where(mutual0, max0.values, zero)
        mscores1 = torch.where(mutual1, mscores0.gather(1, indices1), zero)
        # valid0 = mutual0 & (mscores0 > self.config['match_threshold'])
        valid0 = mutual0 & (mscores0 > p)
        valid1 = mutual1 & valid0.gather(1, indices1)
        indices0 = torch.where(valid0, indices0, indices0.new_tensor(-1))
        indices1 = torch.where(valid1, indices1, indices1.new_tensor(-1))

        return indices0, indices1, mscores0, mscores1