"""
高分辨率Patch特征抽取模块
"""
import cv2
import torch
import torch.nn as nn
import numpy as np

from .base import convnxn
from .base import MRFBlock
from .transformer import Encoder, VectorsEmbeddingModule
from .mae import MAE

from lib.models.utils.models import _tranpose_and_gather_feat

# 简单的使用CNN作为特征抽取器
class FineFeatExtractor(nn.Module):
    """
    特征抽取网络 从高清Patch中抽取细粒度特征 用于增强ReID部分
    主要要考虑到网络的简洁性
    """
    def __init__(self):
        super(FineFeatExtractor, self).__init__()


        self.first_conv = convnxn(3, 32, kernel=7, stride=1, padding=3)

        self.first_stage = nn.Sequential(
            MRFBlock(32, 1),
        )

        self.second_stage = nn.Sequential(
            convnxn(32, 64, kernel=3, stride=2, padding=1),
            MRFBlock(64, 2),
        )

        self.last_stage = nn.Sequential(
            convnxn(64, 128, kernel=3, stride=2, padding=1),
            MRFBlock(128, 3),
        )

    def forward(self, x):
        """
        x: 输入图像特征
        """
        x = self.first_conv(x)            # 32 * 112 * 112
        x = self.first_stage(x)           # 32 * 112 * 112
        x = self.second_stage(x)          # 64 * 56 * 56
        x = self.last_stage(x)            # 128 * 28 * 28

        return x

# 特征细化网络 (使用CNN抽取器的旧版本)
class FineFeatNet(nn.Module):
    """
    利用Transformer结构进行细粒度特征编码
    """
    def __init__(self):
        super(FineFeatNet, self).__init__()

        self.fine_extract = FineFeatExtractor()
        # ========================================================================
        # TODO 根据传参设置而不是硬编码  需要在dataset中设置输出的最大尺寸是112*112
        # ========================================================================
        self.pos_emb = VectorsEmbeddingModule(pool_size=(4, 4, 0), patch_num=49, embedding_dim=128, dropout=0.1)
        self.encoder = Encoder(layers=1, embedding_dim=128, reduced_dim=128, n_head=8, hidden_unit=512, dropout=0.1)

    def forward(self, low_outputs, batch):
        # =====================================================================================================
        # !!! 原始项目代码里需要累计梯度导致输出是一个列表 但实际上并没有用到梯度累计 有时间修改这部分内容 重构一下代码
        # =====================================================================================================
        if self.training:
            for i in range(len(low_outputs)):
                # 训练模式 patch从batch中直接加载
                # batch * k  * reid_dim 其中包括存在对象的也有不存在对象的
                low_output = low_outputs[i]
                id_head = _tranpose_and_gather_feat(low_output['id'], batch['ind'])
                # batch * k * 3 * h * w
                hd_patches = batch['hd_patches']
                # batch * k (0 / 1指示变量)
                selected_mask = batch['reg_mask']
                # reshape
                low_id_vectors = id_head[selected_mask > 0].contiguous()   # .view(-1, reid_dim)
                hd_patches_x = hd_patches[selected_mask > 0].contiguous()  # .view(-1, c, h, w)
                # (batch * total_k) * dim * patch * patch
                # 仅在此处放入显存
                hd_patches_x = hd_patches_x.to(low_id_vectors.device)
                fine_feat = self.fine_extract(hd_patches_x)
                # (batch * total_k) * (patch_nums + 1) * dim
                embeddings = self.pos_emb(fine_feat, low_id_vectors)
                # (batch * total_k) * (patch_nums + 1) * dim
                output = self.encoder(embeddings)

                # 将强化后的特征添加到原来的输出当中
                # (batch * total_k) *  dim
                fine_vectors = output[:, 0]
                low_outputs[i]["fine_id"] = torch.cat([low_id_vectors, fine_vectors], dim = 0)
            return low_outputs
        else:
            
            # 测试模式 patch需要从low_output里面decode出来
            # 1. 根据header的输出进行decode, 得到经过筛选之后的有效的位置的初始reid特征 (只对置信度大于某个阈值的位置进行计算)
            # 2. 根据解出的bbox得到在原始图像中的patch
            # 3. 利用low_reid和patch进行特征增强 选中的位置用于特征增强 没有选中的位置保持原来的low_reid
            # 4. 返回修改之后的特征 (batch * reid_dim * output_h * output_w)

            img0 = batch['img0']
            thres = batch['thres']
            transforms = batch['transforms']
            patch_size = batch['patch_size']
            # thres = 0.8
            origin_height, origin_width = img0.shape[0], img0.shape[1]
            # print(origin_height, origin_width)
            for i in range(len(low_outputs)):
                low_output = low_outputs[i]
                # reid_feat = low_output['id']
                # hm = low_output['hm']
                # reg = low_output['reg']
                # wh = low_output['wh']

                # print(reid_feat.size())
                # print(hm.size())
                # print(reg.size())
                # print(wh.size())

                hm = low_output['hm']
                # for simplest evaluate
                b, c, h, w = hm.size()
                assert b == 1

                hm = hm.clone().view(b, c, -1).squeeze(0).squeeze(0)
                sigmoid_hm = torch.sigmoid(hm)
                
                # print("====================================================================")
                selected_mask = (sigmoid_hm >= thres).int().cpu().numpy()

                # b * 2 * h * w
                reg = low_output['reg'].clone().view(b, 2, -1).squeeze(0).cpu().numpy()
                wh = low_output['wh'].clone().view(b, 4, -1).squeeze(0).cpu().numpy()

                # print(selected_mask)
                # print(selected_mask.sum())
                # print(hm.size())
                # print(reg.size())
                # print(wh.size())

                selected_reg = reg[:, selected_mask == 1].transpose()
                selected_wh = wh[:, selected_mask == 1].transpose()
                indices = np.argwhere(selected_mask > 0)
                # print(indices.shape)
                # print(selected_reg.shape)
                # print(selected_wh.shape)

                # main decode process
                pacthes_tensors = []
                for i in range(indices.shape[0]):
                    # print("-------------------- bbox i -----------------------")
                    center_id_pos = indices[i, 0]
                    center_reg = selected_reg[i, :]
                    wh_ltrb = selected_wh[i, :]

                    # print(center_id_pos)
                    # print(center_reg)
                    # print(wh_ltrb)
                    # 对应于输出图中的中心位置
                    center_y = center_id_pos // w
                    center_x = center_id_pos % w

                    center_y += center_reg[0]
                    center_x += center_reg[1]

                    lt_x = max((center_x - wh_ltrb[0]), 0)
                    lt_y = max((center_y - wh_ltrb[1]), 0)
                    rb_x = min((center_x + wh_ltrb[2]), w)
                    rb_y = min((center_y + wh_ltrb[3]), h)

                    lt_x = int(lt_x / w * origin_width)
                    lt_y = int(lt_y / h * origin_height)
                    rb_x = int(rb_x / w * origin_width)
                    rb_y = int(rb_y / h * origin_height)
                    # print(lt_x, lt_y, rb_x, rb_y)

                   
                    crop_img = img0[lt_y:rb_y, lt_x:rb_x].copy()
                    if crop_img.shape[0] == 0 or crop_img.shape[1] == 0:
                        # invalid patch avoiding empty error
                        crop_img = img0[0:10, 0:10].copy()
                    crop_img = cv2.resize(crop_img, (patch_size, patch_size), cv2.INTER_CUBIC)
                    
                    # visual
                    # cv2.imwrite("./crop/crop_%d_.jpg" % (self.num, ), crop_img)
                    # self.num += 1

                    # BGR -> RGB
                    crop_img = np.ascontiguousarray(crop_img[:, :, ::-1])
                    crop_img = transforms(crop_img).unsqueeze(0)
                    pacthes_tensors.append(crop_img)

                if len(pacthes_tensors) > 0:

                    low_reid_feat = low_output['id'].clone()
                    b, reid_dim, h, w = low_reid_feat.size()

                    # (batch * h * w) * reid_dim
                    low_reid_feat = low_reid_feat.permute(0, 2, 3, 1).view(-1, reid_dim)
                    
                    # print(low_reid_feat.size())
                    indices = indices.reshape((-1,))
                    selected_low_reid_vectors = low_reid_feat[indices, :]

                    hr_patches_input = torch.cat(pacthes_tensors, dim=0).to(selected_low_reid_vectors.device)
                    fine_feat = self.fine_extract(hr_patches_input)
                    embeddings = self.pos_emb(fine_feat, selected_low_reid_vectors)
                    fine_vectors = self.encoder(embeddings)[:, 0]

                    # patches * 128
                    for (k, ind) in zip(range(len(indices)), indices):
                        centery = ind // w
                        centerx = ind % w
                        # instead
                        low_output['id'][0, :, centery, centerx] = fine_vectors[k]

            return low_outputs

# 特征细化网络 (使用MAE抽取器的新版本)
class MAEFeatNet(nn.Module):
    def __init__(
        self,
        mae_pretrain_path = None
    ):
        """
        MAE细粒度特征抽取器 参数设置需要和MAE预训练时保持一致
        Args:
            mae_pretrain_path: 预训练权重路径
        """
        # ========================================================================
        # TODO 根据传参设置而不是硬编码  需要在dataset中设置输出的最大尺寸是128*128
        # ========================================================================
        super(MAEFeatNet, self).__init__()
        # 设置MAE抽取器
        self.mae = MAE(
            128,
            16, 
            128, 
            4, 
            2,
            0.75
        )
        if self.training:
            # 利用预训练模型来初始化进行完整训练
            self.mae.load_state_dict(torch.load(mae_pretrain_path))
            print("MAE loading finished.")
        # 去除MAE预训练时的解码器
        del self.mae.decoder

        # 设置ReID任务时的新解码器
        self.pos_embeddings = nn.parameter.Parameter(torch.zeros(size = (16 ** 2 + 1, 128), requires_grad=True))
        self.transformer = Encoder(1, 128, 128, 8, 512, 0.1)

    def forward(self, low_outputs, batch):
        if self.training:

            # 由于代码遗留问题low_outputs实际上是一个里面只有一个字典元素的列表

            for i in range(len(low_outputs)):
                # 训练模式 patch从batch中直接加载
                # batch * k  * reid_dim 其中包括存在对象的也有不存在对象的
                low_output = low_outputs[i]
                id_head = _tranpose_and_gather_feat(low_output['id'], batch['ind'])
                # batch * k * 3 * h * w
                hd_patches = batch['hd_patches']
                # batch * k (0 / 1指示变量)
                selected_mask = batch['reg_mask']
                # reshape
                low_id_vectors = id_head[selected_mask > 0].contiguous()   # .view(-1, reid_dim)
                hd_patches_x = hd_patches[selected_mask > 0].contiguous()  # .view(-1, c, h, w)
                # (batch * total_k) * dim * patch * patch
                # 仅在此处放入显存
                hd_patches_x = hd_patches_x.to(low_id_vectors.device)

                # 使用MAE从高清图像Patch中获取细粒度特征
                fine_feat = self.mae(hd_patches_x, encoder_mode = True)
                # (batch * total_k) * (patch_nums + 1) * dim
                embeddings = torch.cat([low_id_vectors.unsqueeze(1), fine_feat], dim=1) + self.pos_embeddings
                # (batch * total_k) * (patch_nums + 1) * dim
                output = self.transformer(embeddings)

                # 将强化后的特征添加到原来的输出当中
                # (batch * total_k) *  dim
                fine_vectors = output[:, 0]

                # 将原本没有增强过的和已经增强过的同时返回以供下一步处理
                low_outputs[i]["fine_id"] = torch.cat([low_id_vectors, fine_vectors], dim = 0)

            return low_outputs
        else:
            # !!! Note: 所谓的可插拔即low_outputs是否使用下面的增强 如果不增强完全可以不使用整个MAE部分

            # 测试模式 patch需要从low_output里面decode出来
            # 1. 根据header的输出进行decode, 得到经过筛选之后的有效的位置的初始reid特征 (只对置信度大于某个阈值的位置进行计算)
            # 2. 根据解出的bbox得到在原始图像中的patch
            # 3. 利用low_reid和patch进行特征增强 选中的位置用于特征增强 没有选中的位置保持原来的low_reid
            # 4. 返回修改之后的特征 (batch * reid_dim * output_h * output_w)

            img0 = batch['img0']
            thres = batch['thres']
            transforms = batch['transforms']
            patch_size = 128
            # thres = 0.8
            origin_height, origin_width = img0.shape[0], img0.shape[1]
            # print(origin_height, origin_width)
            for i in range(len(low_outputs)):
                low_output = low_outputs[i]
                # reid_feat = low_output['id']
                # hm = low_output['hm']
                # reg = low_output['reg']
                # wh = low_output['wh']

                # print(reid_feat.size())
                # print(hm.size())
                # print(reg.size())
                # print(wh.size())

                hm = low_output['hm']
                # for simplest evaluate
                b, c, h, w = hm.size()
                assert b == 1

                hm = hm.clone().view(b, c, -1).squeeze(0).squeeze(0)
                sigmoid_hm = torch.sigmoid(hm)
                
                # print("====================================================================")
                selected_mask = (sigmoid_hm >= thres).int().cpu().numpy()

                # b * 2 * h * w
                reg = low_output['reg'].clone().view(b, 2, -1).squeeze(0).cpu().numpy()
                wh = low_output['wh'].clone().view(b, 4, -1).squeeze(0).cpu().numpy()

                # print(selected_mask)
                # print(selected_mask.sum())
                # print(hm.size())
                # print(reg.size())
                # print(wh.size())

                selected_reg = reg[:, selected_mask == 1].transpose()
                selected_wh = wh[:, selected_mask == 1].transpose()
                indices = np.argwhere(selected_mask > 0)
                # print(indices.shape)
                # print(selected_reg.shape)
                # print(selected_wh.shape)

                # main decode process
                pacthes_tensors = []
                for i in range(indices.shape[0]):
                    # print("-------------------- bbox i -----------------------")
                    center_id_pos = indices[i, 0]
                    center_reg = selected_reg[i, :]
                    wh_ltrb = selected_wh[i, :]

                    # print(center_id_pos)
                    # print(center_reg)
                    # print(wh_ltrb)
                    # 对应于输出图中的中心位置
                    center_y = center_id_pos // w
                    center_x = center_id_pos % w

                    center_y += center_reg[0]
                    center_x += center_reg[1]

                    lt_x = max((center_x - wh_ltrb[0]), 0)
                    lt_y = max((center_y - wh_ltrb[1]), 0)
                    rb_x = min((center_x + wh_ltrb[2]), w)
                    rb_y = min((center_y + wh_ltrb[3]), h)

                    lt_x = int(lt_x / w * origin_width)
                    lt_y = int(lt_y / h * origin_height)
                    rb_x = int(rb_x / w * origin_width)
                    rb_y = int(rb_y / h * origin_height)
                    # print(lt_x, lt_y, rb_x, rb_y)

                    crop_img = img0[lt_y:rb_y, lt_x:rb_x].copy()
                    if crop_img.shape[0] == 0 or crop_img.shape[1] == 0:
                        # invalid patch avoiding empty error
                        crop_img = img0[0:10, 0:10].copy()
                    crop_img = cv2.resize(crop_img, (patch_size, patch_size), cv2.INTER_CUBIC)
                    
                    # visual
                    # cv2.imwrite("./crop/crop_%d_.jpg" % (self.num, ), crop_img)
                    # self.num += 1

                    # BGR -> RGB
                    crop_img = np.ascontiguousarray(crop_img[:, :, ::-1])
                    crop_img = transforms(crop_img).unsqueeze(0)
                    pacthes_tensors.append(crop_img)

                if len(pacthes_tensors) > 0:

                    low_reid_feat = low_output['id'].clone()
                    b, reid_dim, h, w = low_reid_feat.size()

                    # (batch * h * w) * reid_dim
                    low_reid_feat = low_reid_feat.permute(0, 2, 3, 1).view(-1, reid_dim)
                    
                    # print(low_reid_feat.size())
                    indices = indices.reshape((-1,))
                    selected_low_reid_vectors = low_reid_feat[indices, :]

                    hr_patches_input = torch.cat(pacthes_tensors, dim=0).to(selected_low_reid_vectors.device)
                    
                    # 使用mae抽取特征
                    fine_feat = self.mae(hr_patches_input, encoder_mode = True)
                    embeddings = torch.cat([selected_low_reid_vectors.unsqueeze(1), fine_feat], dim=1) + self.pos_embeddings
                    fine_vectors = self.transformer(embeddings)[:, 0]

                    # patches * 128
                    for (k, ind) in zip(range(len(indices)), indices):
                        centery = ind // w
                        centerx = ind % w
                        # instead
                        low_output['id'][0, :, centery, centerx] = fine_vectors[k]

            return low_outputs

    