# from configs.common import *
from configs.backbones import *
from configs.necks import *
from configs.heads import *
from configs.losses import *
from configs.common import BaseModule,Sequential,TwoInputSequential
# from losses.contrastive_loss import ProjectedContrastiveLoss

import torch.nn as nn
import torch

import functools
from inspect import getfullargspec
from collections import abc
import numpy as np

# def build_model(cfg):
#     if isinstance(cfg, list):
#         modules = [
#             eval(cfg_.pop("type"))(**cfg_) for cfg_ in cfg
#         ]
#         return Sequential(*modules)
#     else:
#         return eval(cfg.pop("type"))(**cfg)
    
_shared_instances = {} 

# def build_model(cfg):
#     if isinstance(cfg, list):
#         modules = []
#         for cfg_ in cfg:
#             # check share
#             shared_id = cfg_.get('_shared_id', None)
#             if shared_id and shared_id in _shared_instances:
#                 modules.append(_shared_instances[shared_id])
#             else:
#                 module_type = cfg_.pop('type')
#                 module = eval(module_type)(**cfg_)
#                 if shared_id:
#                     _shared_instances[shared_id] = module
#                 modules.append(module)
#         return nn.Sequential(*modules)
#     else:
#         shared_id = cfg.get('_shared_id', None)
#         if shared_id and shared_id in _shared_instances:
#             return _shared_instances[shared_id]
#         else:
#             module_type = cfg.pop('type')
#             module = eval(module_type)(**cfg)
#             if shared_id:
#                 _shared_instances[shared_id] = module
#             return module

# def build_model(cfg):
#     if isinstance(cfg, list):
#         modules = []
#         for cfg_ in cfg:
#             shared_id = cfg_.get('_shared_id', None)
#             if shared_id and shared_id in _shared_instances:
#                 modules.append(_shared_instances[shared_id])
#             else:
#                 module_type = cfg_.pop('type')
#                 if module_type == 'Sequential':
#                     layers = []
#                     for layer_cfg in cfg_.get('module_list', []): 
#                         layers.append(build_model(layer_cfg)) 
#                     module = nn.Sequential(*layers)  
#                 else:
#                     module = eval(module_type)(**cfg_)
                
#                 if shared_id:
#                     _shared_instances[shared_id] = module
#                 modules.append(module)
#         return nn.Sequential(*modules)  # 返回构建的 nn.Sequential

#     else:
#         # 如果 cfg 是单个模块的配置
#         shared_id = cfg.get('_shared_id', None)
#         if shared_id and shared_id in _shared_instances:
#             return _shared_instances[shared_id]
#         else:
#             # 获取模块类型
#             module_type = cfg.pop('type')
#             # 如果是 Sequential 类型，处理为层的列表
#             if module_type == 'Sequential':
#                 layers = []
#                 for layer_cfg in cfg.get('module_list', []):  # 获取实际的层配置
#                     layers.append(build_model(layer_cfg))  # 递归构建层
#                 module = nn.Sequential(*layers)  # 使用 nn.Sequential 构建层
#             else:
#                 # 对于非 Sequential 类型，正常构建
#                 module = eval(module_type)(**cfg)

#             if shared_id:
#                 _shared_instances[shared_id] = module
#             return module
def build_model(cfg):
    if isinstance(cfg, list):
        modules = []
        for cfg_ in cfg:
            module_type = cfg_.pop('type')
            if module_type == 'TwoInputSequential':
                modules = [build_model(layer_cfg) for layer_cfg in cfg['module_list']]
                module = TwoInputSequential(*modules)
            else:
                module = eval(module_type)(**cfg_)
            modules.append(module)
        return nn.Sequential(*modules)
    else:
        shared_id = cfg.get('_shared_id', None)
        if shared_id and shared_id in _shared_instances:
            return _shared_instances[shared_id]
        else:
            module_type = cfg.pop('type')
            if module_type == 'TwoInputSequential':
                modules = [build_model(layer_cfg) for layer_cfg in cfg['module_list']]
                module = TwoInputSequential(*modules)
            else:
                module = eval(module_type)(**cfg)
            if shared_id:
                _shared_instances[shared_id] = module
            return module

class BuildNet(BaseModule):
    def __init__(self,cfg):
        super(BuildNet, self).__init__()
        self.neck_cfg = cfg.get("neck")
        self.head_cfg = cfg.get("head")
        self.contrastive_head = build_model(cfg.get("contrastive_head"))
        self.backbone = build_model(cfg.get("backbone"))
        if self.neck_cfg is not None:
            self.neck = build_model(cfg.get("neck"))
        
        if self.head_cfg is not None:
            if 'losses' in self.head_cfg: 
                self.losses = [build_model(loss_cfg) for loss_cfg in self.head_cfg['losses']]
            elif 'loss' in self.head_cfg: 
                self.losses = [build_model(self.head_cfg['loss'])]
            else:
                self.losses = []
            self.head = build_model(cfg.get("head"))

        # self.contrastive_loss = ProjectedContrastiveLoss(temperature=0.7)
        self.contrastive_loss = SupConLoss(temperature=0.07)

    def freeze_layers(self,names):
        assert isinstance(names,tuple)
        for name in names:
            layers = getattr(self, name)
            # layers.eval()
            for param in layers.parameters():
                param.requires_grad = False
    
    # def extract_feat(self, img, stage='neck'):
    #     """Directly extract features from the specified stage.

    #     Args:
    #         img (Tensor): The input images. The shape of it should be
    #             ``(num_samples, num_channels, *img_shape)``.
    #         stage (str): Which stage to output the feature. Choose from
    #             "backbone", "neck" and "pre_logits". Defaults to "neck".

    #     Returns:
    #         tuple | Tensor: The output of specified stage.
    #             The output depends on detailed implementation. In general, the
    #             output of backbone and neck is a tuple and the output of
    #             pre_logits is a tensor.

    #     Examples:
    #         1. Backbone output

    #         >>> import torch
    #         >>> from mmcv import Config
    #         >>> from mmcls.models import build_classifier
    #         >>>
    #         >>> cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py').model
    #         >>> cfg.backbone.out_indices = (0, 1, 2, 3)  # Output multi-scale feature maps
    #         >>> model = build_classifier(cfg)
    #         >>> outs = model.extract_feat(torch.rand(1, 3, 224, 224), stage='backbone')
    #         >>> for out in outs:
    #         ...     print(out.shape)
    #         torch.Size([1, 64, 56, 56])
    #         torch.Size([1, 128, 28, 28])
    #         torch.Size([1, 256, 14, 14])
    #         torch.Size([1, 512, 7, 7])

    #         2. Neck output

    #         >>> import torch
    #         >>> from mmcv import Config
    #         >>> from mmcls.models import build_classifier
    #         >>>
    #         >>> cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py').model
    #         >>> cfg.backbone.out_indices = (0, 1, 2, 3)  # Output multi-scale feature maps
    #         >>> model = build_classifier(cfg)
    #         >>>
    #         >>> outs = model.extract_feat(torch.rand(1, 3, 224, 224), stage='neck')
    #         >>> for out in outs:
    #         ...     print(out.shape)
    #         torch.Size([1, 64])
    #         torch.Size([1, 128])
    #         torch.Size([1, 256])
    #         torch.Size([1, 512])

    #         3. Pre-logits output (without the final linear classifier head)

    #         >>> import torch
    #         >>> from mmcv import Config
    #         >>> from mmcls.models import build_classifier
    #         >>>
    #         >>> cfg = Config.fromfile('configs/vision_transformer/vit-base-p16_pt-64xb64_in1k-224.py').model
    #         >>> model = build_classifier(cfg)
    #         >>>
    #         >>> out = model.extract_feat(torch.rand(1, 3, 224, 224), stage='pre_logits')
    #         >>> print(out.shape)  # The hidden dims in head is 3072
    #         torch.Size([1, 3072])
    #     """  # noqa: E501
    #     assert stage in ['backbone', 'neck', 'pre_logits'], \
    #         (f'Invalid output stage "{stage}", please choose from "backbone", '
    #          '"neck" and "pre_logits"')

    #     x = self.backbone(img)

    #     if stage == 'backbone':
    #         return x

    #     if hasattr(self, 'neck') and self.neck is not None:
    #         x = self.neck(x)
    #     if stage == 'neck':
    #         return x

    def extract_feat(self, imgs, stage='neck'):

        view1, view2 = imgs

        feat1 = self.backbone(view1)
        feat2 = self.backbone(view2)

        if isinstance(feat1, (tuple, list)):
            feat1 = feat1[0]
        if isinstance(feat2, (tuple, list)):
            feat2 = feat2[0]

        # 2. 如果阶段为 backbone，直接返回
        if stage == 'backbone':
            return feat1, feat2

        # 3. 通过 neck 模块处理
        if hasattr(self, 'neck') and self.neck is not None:
            feat1, feat2 = self.neck(feat1, feat2)

        return feat1, feat2

    def extract_gray_stats(self, x):
        """
        提取每张灰度图的 mean/std，输入 x shape: [B, 1, H, W]
        返回 shape: [B, 2]
        """
        B = x.size(0)
        mean = x.view(B, -1).mean(dim=1, keepdim=True)  # [B, 1]
        std = x.view(B, -1).std(dim=1, keepdim=True)    # [B, 1]
        return torch.cat([mean, std], dim=1)            # [B, 2]

    
    # def forward(self, x, return_loss=True, train_statu=False, **kwargs):
    #     x = self.extract_feat(x)
        
    #     if not train_statu:
    #         if return_loss:
    #             return self.forward_train(x, **kwargs)
    #         else:
    #             return self.forward_test(x, **kwargs)
    #     else:
    #         return self.forward_test(x), self.forward_train(x, **kwargs)

    def forward(self, imgs, return_loss=True, train_statu=False, targets=None, **kwargs):
        view1, view2 = imgs
        feat1, feat2 = self.extract_feat((view1, view2)) 
        gray_stats = self.extract_gray_stats(view1)
        feat1 = torch.cat([feat1, gray_stats], dim=1)
        feat2 = torch.cat([feat2, gray_stats], dim=1)

        proj1 = self.contrastive_head(feat1)
        proj2 = self.contrastive_head(feat2)

        if not train_statu:
            if return_loss:
                return self.forward_train(proj1, targets, feat2=proj2, **kwargs)
            else:
                return self.forward_test(proj1, **kwargs)
        else:
            return self.forward_test(proj1, **kwargs), self.forward_train(proj1, targets, feat2=proj2, **kwargs)

        
    # def forward_train(self, x, targets, **kwargs):
         
    #     losses = dict()
    #     loss = self.head.forward_train(x, targets, **kwargs)
    #     losses.update(loss)
    #     return losses

    # def forward_train(self, x, targets, feat2=None, **kwargs):    #contrastive_loss
    #     losses = dict()
    #     # 分类loss用feat1的特征
    #     loss_cls = self.head.forward_train(x, targets, **kwargs)
    #     losses.update(loss_cls)
    
    #     # 对比loss：需要有feat2（第二视图的特征）
    #     if feat2 is not None and hasattr(self, 'contrastive_loss'):
    #         loss_contrast = self.contrastive_loss(x, feat2)
    #         losses['loss_contrast'] = loss_contrast
    #         # 总loss合并
    #         losses['loss'] = losses.get('loss_cls', 0) + 0.5 * loss_contrast
    #     else:
    #         losses['loss'] = losses.get('loss_cls', 0)

    #     return losses

    def forward_train(self, proj1, targets, feat2=None, **kwargs):
        losses = dict()
        loss_cls = self.head.forward_train(proj1, targets, **kwargs)
        losses.update(loss_cls)

        # ✅ SupConLoss expects [B, 2, D] feature + label
        if feat2 is not None and hasattr(self, 'contrastive_loss'):
            proj1 = proj1.unsqueeze(1)  # [B, 1, D]
            proj2 = feat2.unsqueeze(1)  # [B, 1, D]
            features = torch.cat([proj1, proj2], dim=1)  # [B, 2, D]
            loss_contrast = self.contrastive_loss(features, targets)
            losses['loss_contrast'] = loss_contrast
            losses['loss'] = losses.get('loss_cls', 0) + loss_contrast
        else:
            losses['loss'] = losses.get('loss_cls', 0)

        return losses

        
    def forward_test(self, x, **kwargs):
        
        out = self.head.simple_test(x,**kwargs)
        return out
    
