import torch.nn as nn
import torch
import torch.nn.functional as F

from lib.model.frozen_bn import FrozenBatchNorm2d


def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1,
         freeze_bn=False):
    if freeze_bn:
        return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, bias=True),
            FrozenBatchNorm2d(out_planes),
            nn.ReLU(inplace=True))
    else:
        return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, bias=True),
            nn.BatchNorm2d(out_planes),
            nn.ReLU(inplace=True))


class Corner_Predictor(nn.Module):
    """ Corner Predictor module"""

    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
        super(Corner_Predictor, self).__init__()
        self.feat_sz = feat_sz
        self.stride = stride
        self.img_sz = self.feat_sz * self.stride
        '''top-left corner'''
        self.conv1_tl = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_tl = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_tl = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_tl = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_tl = nn.Conv2d(channel // 8, 1, kernel_size=1)

        '''bottom-right corner'''
        self.conv1_br = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_br = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_br = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_br = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_br = nn.Conv2d(channel // 8, 1, kernel_size=1)

        '''about coordinates and indexs'''
        with torch.no_grad():
            self.indice = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride
            # generate mesh-grid
            self.coord_x = self.indice.repeat((self.feat_sz, 1)) \
                .view((self.feat_sz * self.feat_sz,)).float().cuda()
            self.coord_y = self.indice.repeat((1, self.feat_sz)) \
                .view((self.feat_sz * self.feat_sz,)).float().cuda()

    def forward(self, x, return_dist=False, softmax=True):
        """ Forward pass with input x. """
        score_map_tl, score_map_br = self.get_score_map(x)
        if return_dist:
            coorx_tl, coory_tl, prob_vec_tl = self.soft_argmax(score_map_tl, return_dist=True, softmax=softmax)
            coorx_br, coory_br, prob_vec_br = self.soft_argmax(score_map_br, return_dist=True, softmax=softmax)
            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz, prob_vec_tl, prob_vec_br
        else:
            coorx_tl, coory_tl = self.soft_argmax(score_map_tl)
            coorx_br, coory_br = self.soft_argmax(score_map_br)
            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz

    def get_score_map(self, x):
        # top-left branch
        x_tl1 = self.conv1_tl(x)
        x_tl2 = self.conv2_tl(x_tl1)
        x_tl3 = self.conv3_tl(x_tl2)
        x_tl4 = self.conv4_tl(x_tl3)
        score_map_tl = self.conv5_tl(x_tl4)

        # bottom-right branch
        x_br1 = self.conv1_br(x)
        x_br2 = self.conv2_br(x_br1)
        x_br3 = self.conv3_br(x_br2)
        x_br4 = self.conv4_br(x_br3)
        score_map_br = self.conv5_br(x_br4)
        return score_map_tl, score_map_br

    def soft_argmax(self, score_map, return_dist=False, softmax=True):
        """ get soft-argmax coordinate for a given heatmap """
        score_vec = score_map.view((-1, self.feat_sz * self.feat_sz))  # (batch, feat_sz * feat_sz)
        prob_vec = nn.functional.softmax(score_vec, dim=1)
        exp_x = torch.sum((self.coord_x * prob_vec), dim=1)
        exp_y = torch.sum((self.coord_y * prob_vec), dim=1)
        if return_dist:
            if softmax:
                return exp_x, exp_y, prob_vec
            else:
                return exp_x, exp_y, score_vec
        else:
            return exp_x, exp_y


def _sigmoid(x):
    y = torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4)
    return y

class CenterPredictor(nn.Module, ):
    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False, multi_object=False):
        super(CenterPredictor, self).__init__()
        self.feat_sz = feat_sz
        self.stride = stride
        self.img_sz = self.feat_sz * self.stride
        self.multi_object = multi_object
        self.inference = False

        # center predict
        self.conv1_ctr = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_ctr = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_ctr = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_ctr = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_ctr = nn.Conv2d(channel // 8, 1, kernel_size=1)

        # size regress
        self.conv1_size = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_size = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_size = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_size = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_size = nn.Conv2d(channel // 8, 2, kernel_size=1)

        # offset regress
        self.conv1_offset = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_offset = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_offset = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_offset = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_offset = nn.Conv2d(channel // 8, 2, kernel_size=1)

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, x, gt_score_map=None):
        """ Forward pass with input x. """
        score_map_ctr_raw, size_map_raw, offset_map = self.get_score_map(x)

        if torch.onnx.is_in_onnx_export():
            return torch.sigmoid(score_map_ctr_raw), torch.sigmoid(size_map_raw), offset_map
        else:
            score_map_ctr = _sigmoid(score_map_ctr_raw)
            size_map = _sigmoid(size_map_raw)
            if self.multi_object: 
                if self.inference:
                    bboxs, scores = self.get_all_bbox(score_map_ctr, size_map, offset_map, thr=0.1)
                    return bboxs, scores
                else:
                    batch_indices, idx_x, idx_y, peak_scores = self.find_peaks(score_map_ctr)
                    bbox = self.compute_bboxes(batch_indices, idx_x, idx_y, size_map, offset_map)
                    return score_map_ctr, bbox, batch_indices, size_map, offset_map, peak_scores
            else:
                if gt_score_map is None:
                    bbox = self.cal_bbox(score_map_ctr, size_map, offset_map)
                else:
                    bbox = self.cal_bbox(gt_score_map.unsqueeze(1), size_map, offset_map)

            return score_map_ctr, bbox, size_map, offset_map

    def get_all_bbox(self, score_map, size_map, offset_map, thr = 0.1):
        flat_score_map = score_map.flatten(0)
        # Reshape size_map and offset_map to [bs, 2, stride*stride]
        size_map_flat = size_map.flatten(2)  # [bs, 2, stride*stride]
        offset_map_flat = offset_map.flatten(2)  # [bs, 2, stride*stride]

        # Compute flattened index for each (idx_x, idx_y) coordinate
        stride = size_map.shape[2]  # assuming square feature map

        bboxes = []
        scores = []
        for idx in range(flat_score_map.shape[0]):
            if flat_score_map[idx] > thr:
                # Gather size and offset values for each detected peak, based on batch index
                size = size_map_flat[0, :, idx]  # [n_peaks, 2]
                offset = offset_map_flat[0, :, idx]  # [n_peaks, 2]

                # Compute bbox coordinates (cx, cy, w, h)
                idx_y = idx // self.feat_sz
                idx_x = idx % self.feat_sz
                cx = (idx_x + offset[0]) / stride
                cy = (idx_y + offset[1]) / stride
                w = size[0]
                h = size[1]

                # Concatenate into a single tensor [n_peaks, 4]
                bboxes.append([float(cx), float(cy), float(w), float(h)])
                scores.append(float(flat_score_map[idx]))
        return bboxes, scores
        

    def cal_bbox(self, score_map_ctr, size_map, offset_map, return_score=False):
        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
        idx_y = idx // self.feat_sz
        idx_x = idx % self.feat_sz

        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
        size = size_map.flatten(2).gather(dim=2, index=idx)
        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)

        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
        # cx, cy, w, h
        bbox = torch.cat([(idx_x.to(torch.float) + offset[:, :1]) / self.feat_sz,
                          (idx_y.to(torch.float) + offset[:, 1:]) / self.feat_sz,
                          size.squeeze(-1)], dim=1)

        if return_score:
            return bbox, max_score
        return bbox

    def get_pred(self, score_map_ctr, size_map, offset_map):
        max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
        idx_y = idx // self.feat_sz
        idx_x = idx % self.feat_sz

        idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
        size = size_map.flatten(2).gather(dim=2, index=idx)
        offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)

        # bbox = torch.cat([idx_x - size[:, 0] / 2, idx_y - size[:, 1] / 2,
        #                   idx_x + size[:, 0] / 2, idx_y + size[:, 1] / 2], dim=1) / self.feat_sz
        return size * self.feat_sz, offset

    def get_score_map(self, x):


        # ctr branch
        x_ctr1 = self.conv1_ctr(x)
        x_ctr2 = self.conv2_ctr(x_ctr1)
        x_ctr3 = self.conv3_ctr(x_ctr2)
        x_ctr4 = self.conv4_ctr(x_ctr3)
        score_map_ctr = self.conv5_ctr(x_ctr4)

        # offset branch
        x_offset1 = self.conv1_offset(x)
        x_offset2 = self.conv2_offset(x_offset1)
        x_offset3 = self.conv3_offset(x_offset2)
        x_offset4 = self.conv4_offset(x_offset3)
        score_map_offset = self.conv5_offset(x_offset4)

        # size branch
        x_size1 = self.conv1_size(x)
        x_size2 = self.conv2_size(x_size1)
        x_size3 = self.conv3_size(x_size2)
        x_size4 = self.conv4_size(x_size3)
        score_map_size = self.conv5_size(x_size4)
        return score_map_ctr, score_map_size, score_map_offset

    def find_peaks(self, score_map, threshold=0.05, kernel_size=3):
        """
        Detect local peaks in the score_map by applying max-pooling.
        
        Args:
            score_map (torch.Tensor): [bs, 1, stride, stride], the heatmap output from the network
            threshold (float): Minimum score threshold to filter low-confidence detections
            kernel_size (int): The window size for max-pooling to find local maxima
            
        Returns:
            batch_indices (torch.Tensor): Batch indices of detected peaks [n_peaks]
            idx_x (torch.Tensor): Detected x coordinates of object centers [n_peaks]
            idx_y (torch.Tensor): Detected y coordinates of object centers [n_peaks]
            peak_scores (torch.Tensor): Scores of the detected peaks [n_peaks]
        """
        # Apply max-pooling to the score_map
        max_pool = F.max_pool2d(score_map, kernel_size, stride=1, padding=kernel_size // 2)
    
        # Identify local maxima (locations where score_map is equal to max_pool)
        peaks = (score_map == max_pool) & (score_map > threshold)
        
        # Get the coordinates of the peaks
        idx = torch.nonzero(peaks, as_tuple=False)  # [n_peaks, 4] where last two dims are (y, x)
        
        # Extract batch indices and (y, x) coordinates of the detected peaks
        batch_indices = idx[:, 0]
        idx_y, idx_x = idx[:, 2], idx[:, 3]

        # Extract the scores of the detected peaks
        peak_scores = score_map[batch_indices, 0, idx_y, idx_x]

        return batch_indices, idx_x, idx_y, peak_scores
    
    def compute_bboxes(self, batch_indices, idx_x, idx_y, size_map, offset_map):
        """
        Compute bounding boxes for multiple detected objects.
        
        Args:
            batch_indices (torch.Tensor): Batch indices of detected peaks [n_peaks]
            idx_x (torch.Tensor): Detected x coordinates of object centers [n_peaks]
            idx_y (torch.Tensor): Detected y coordinates of object centers [n_peaks]
            size_map (torch.Tensor): Size map output from the network [bs, 2, stride, stride]
            offset_map (torch.Tensor): Offset map output from the network [bs, 2, stride, stride]
        
        Returns:
            bboxes (torch.Tensor): Bounding boxes in normalized form [n_peaks, 4 (cx, cy, w, h)]
        """
        n_peaks = batch_indices.shape[0]

        # Reshape size_map and offset_map to [bs, 2, stride*stride]
        size_map_flat = size_map.flatten(2)  # [bs, 2, stride*stride]
        offset_map_flat = offset_map.flatten(2)  # [bs, 2, stride*stride]

        # Compute flattened index for each (idx_x, idx_y) coordinate
        stride = size_map.shape[2]  # assuming square feature map
        idx = idx_y * stride + idx_x  # [n_peaks]

        # Gather size and offset values for each detected peak, based on batch index
        size = size_map_flat[batch_indices, :, idx]  # [n_peaks, 2]
        offset = offset_map_flat[batch_indices, :, idx]  # [n_peaks, 2]

        # Compute bbox coordinates (cx, cy, w, h)
        cx = (idx_x.to(torch.float) + offset[:, 0]) / stride
        cy = (idx_y.to(torch.float) + offset[:, 1]) / stride
        w = size[:, 0]
        h = size[:, 1]

        # Concatenate into a single tensor [n_peaks, 4]
        bboxes = torch.stack([cx, cy, w, h], dim=1)

        return bboxes



class MLP(nn.Module):
    """ Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, BN=False):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        if BN:
            self.layers = nn.ModuleList(nn.Sequential(nn.Linear(n, k), nn.BatchNorm1d(k))
                                        for n, k in zip([input_dim] + h, h + [output_dim]))
        else:
            self.layers = nn.ModuleList(nn.Linear(n, k)
                                        for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x


class Corner_Predictor_v3(nn.Module):
    """ Corner Predictor module"""

    def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=False):
        super(Corner_Predictor_v3, self).__init__()
        self.feat_sz = feat_sz
        self.stride = stride
        self.img_sz = self.feat_sz * self.stride
        '''top-left corner'''
        self.conv1_tl = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_tl = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_tl = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_tl = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_tl = nn.Conv2d(channel // 8, 1, kernel_size=1)

        '''bottom-right corner'''
        self.conv1_br = conv(inplanes, channel, freeze_bn=freeze_bn)
        self.conv2_br = conv(channel, channel // 2, freeze_bn=freeze_bn)
        self.conv3_br = conv(channel // 2, channel // 4, freeze_bn=freeze_bn)
        self.conv4_br = conv(channel // 4, channel // 8, freeze_bn=freeze_bn)
        self.conv5_br = nn.Conv2d(channel // 8, 1, kernel_size=1)

        '''about coordinates and indexs'''
        with torch.no_grad():
            self.indice_tl = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride
            self.indice_br = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride + self.feat_sz - 1
            # generate mesh-grid
            self.coord_x_tl = self.indice_tl.repeat((self.feat_sz, 1)) \
                .view((self.feat_sz * self.feat_sz,)).float()#.cuda()
            self.coord_y_tl = self.indice_tl.repeat((1, self.feat_sz)) \
                .view((self.feat_sz * self.feat_sz,)).float()#.cuda()
            self.coord_x_br = self.indice_br.repeat((self.feat_sz, 1)) \
                .view((self.feat_sz * self.feat_sz,)).float()#.cuda()
            self.coord_y_br = self.indice_br.repeat((1, self.feat_sz)) \
                .view((self.feat_sz * self.feat_sz,)).float()#.cuda()
            # self.coord_x_tl = self.indice_tl.repeat((self.feat_sz, 1)) \
            #     .view((self.feat_sz * self.feat_sz,)).float()
            # self.coord_y_tl = self.indice_tl.repeat((1, self.feat_sz)) \
            #     .view((self.feat_sz * self.feat_sz,)).float()
            # self.coord_x_br = self.indice_br.repeat((self.feat_sz, 1)) \
            #     .view((self.feat_sz * self.feat_sz,)).float()
            # self.coord_y_br = self.indice_br.repeat((1, self.feat_sz)) \
            #     .view((self.feat_sz * self.feat_sz,)).float()

    def forward(self, x, pos_last, return_dist=False, softmax=True):
        """ Forward pass with input x. """
        score_map_tl, score_map_br = self.get_score_map(x)

        if return_dist:
            coorx_tl, coory_tl, prob_vec_tl = self.soft_argmax_tl(score_map_tl, return_dist=True, softmax=softmax)
            coorx_br, coory_br, prob_vec_br = self.soft_argmax_br(score_map_br, return_dist=True, softmax=softmax)
            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz, prob_vec_tl, prob_vec_br
        else:
            coorx_tl, coory_tl = self.soft_argmax_tl(score_map_tl)
            coorx_br, coory_br = self.soft_argmax_br(score_map_br)
            return torch.stack((coorx_tl, coory_tl, coorx_br, coory_br), dim=1) / self.img_sz

    def get_score_map(self, x):
        # top-left branch
        x_tl1 = self.conv1_tl(x)
        x_tl2 = self.conv2_tl(x_tl1)
        x_tl3 = self.conv3_tl(x_tl2)
        x_tl4 = self.conv4_tl(x_tl3)
        score_map_tl = self.conv5_tl(x_tl4)

        # bottom-right branch
        x_br1 = self.conv1_br(x)
        x_br2 = self.conv2_br(x_br1)
        x_br3 = self.conv3_br(x_br2)
        x_br4 = self.conv4_br(x_br3)
        score_map_br = self.conv5_br(x_br4)
        return score_map_tl, score_map_br

    def soft_argmax_tl(self, score_map, return_dist=False, softmax=True):
        """ get soft-argmax coordinate for a given heatmap """
        score_vec = score_map.view((-1, self.feat_sz * self.feat_sz))  # (batch, feat_sz * feat_sz)
        prob_vec = nn.functional.softmax(score_vec, dim=1)
        # import cv2
        # import numpy as np
        # print(score_map.shape)
        # img = np.array((score_map[0][0]).cpu())
        # img = cv2.resize(img, (384, 384))
        # cv2.imshow('img', img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        exp_x = torch.sum((self.coord_x_tl * prob_vec), dim=1)
        exp_y = torch.sum((self.coord_y_tl * prob_vec), dim=1)
        if return_dist:
            if softmax:
                return exp_x, exp_y, prob_vec
            else:
                return exp_x, exp_y, score_vec
        else:
            return exp_x, exp_y

    def soft_argmax_br(self, score_map, return_dist=False, softmax=True):
        """ get soft-argmax coordinate for a given heatmap """
        score_vec = score_map.view((-1, self.feat_sz * self.feat_sz))  # (batch, feat_sz * feat_sz)
        prob_vec = nn.functional.softmax(score_vec, dim=1)
        exp_x = torch.sum((self.coord_x_br * prob_vec), dim=1)
        exp_y = torch.sum((self.coord_y_br * prob_vec), dim=1)
        if return_dist:
            if softmax:
                return exp_x, exp_y, prob_vec
            else:
                return exp_x, exp_y, score_vec
        else:
            return exp_x, exp_y


def build_box_head(cfg, hidden_dim):
    stride = cfg.MODEL.BACKBONE.STRIDE

    if cfg.MODEL.HEAD.TYPE == "MLP":
        mlp_head = MLP(hidden_dim, hidden_dim, 4, 3)  # dim_in, dim_hidden, dim_out, 3 layers
        return mlp_head
    elif "CORNER" in cfg.MODEL.HEAD.TYPE:
        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
        channel = getattr(cfg.MODEL, "NUM_CHANNELS", 256)
        print("head channel: %d" % channel)
        if cfg.MODEL.HEAD.TYPE == "CORNER":
            corner_head = Corner_Predictor(inplanes=cfg.MODEL.HIDDEN_DIM, channel=channel,
                                           feat_sz=feat_sz, stride=stride)
        elif cfg.MODEL.HEAD.TYPE == "CORNERv3":
            corner_head = Corner_Predictor_v3(inplanes=hidden_dim, channel=cfg.MODEL.HEAD.NUM_CHANNELS,
                                           feat_sz=feat_sz, stride=stride)
        else:
            raise ValueError()
        return corner_head
    elif cfg.MODEL.HEAD.TYPE == "CENTER":
        in_channel = hidden_dim
        out_channel = cfg.MODEL.HEAD.NUM_CHANNELS
        feat_sz = int(cfg.DATA.SEARCH.SIZE / stride)
        center_head = CenterPredictor(inplanes=in_channel, channel=out_channel,
                                      feat_sz=feat_sz, stride=stride, multi_object=True)
        return center_head
    else:
        raise ValueError("HEAD TYPE %s is not supported." % cfg.MODEL.HEAD.TYPE)