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


class WeakLabelDataProcessor(nn.Module):
    def __init__(self, img_feat_dim=512, text_feat_dim=768):
        super().__init__()
        # 1. 自监督预训练分支（图像+文本双模态自监督，学习基础特征）
        self.img_ssl_encoder = self._build_img_ssl_encoder()  # 图像自监督编码器
        self.text_ssl_encoder = self._build_text_ssl_encoder()# 文本自监督编码器
        
        # 2. 注意力引导样本筛选器（计算样本置信度，筛除低质量弱标注样本）
        self.attention_scorer = nn.Sequential(
            nn.Linear(img_feat_dim + text_feat_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 1),
            nn.Sigmoid()  # 输出0~1的样本置信度
        )
        self.filter_threshold = 0.6  # 样本筛选阈值（高质量样本阈值）

    def _build_img_ssl_encoder(self):
        """图像自监督预训练编码器（SimCLR风格，无标注学习图像局部特征）"""
        return nn.Sequential(
            nn.Conv2d(3, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 512)
        )

    def _build_text_ssl_encoder(self):
        """文本自监督预训练编码器（掩码自监督，学习文本语义特征）"""
        return nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Linear(512, 512)
        )

    def forward(self, img_data, text_data, weak_label=None):
        # 步骤1：自监督预训练提取双模态基础特征
        img_feat = self.img_ssl_encoder(img_data)  # 图像局部特征
        text_feat = self.text_ssl_encoder(text_data)# 文本语义特征
        fusion_feat = torch.cat([img_feat, text_feat], dim=-1)

        # 步骤2：注意力引导样本筛选（保留高置信度弱标注样本）
        sample_confidence = self.attention_scorer(fusion_feat)
        high_quality_mask = sample_confidence > self.filter_threshold  # 高质量样本掩码
        
        # 输出：高质量样本的双模态特征+筛选掩码
        return (img_feat[high_quality_mask.squeeze()], 
                text_feat[high_quality_mask.squeeze()], 
                high_quality_mask)


class MultiModalCrossAttentionFusion(nn.Module):
    def __init__(self, feat_dim=512, num_heads=8):
        super().__init__()
        # 多模态交叉注意力核心（图像→文本、文本→图像双向注意力）
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=feat_dim, num_heads=num_heads, batch_first=True
        )
        # 细粒度特征增强层
        self.fusion_norm = nn.LayerNorm(feat_dim)
        self.fusion_mlp = nn.Sequential(
            nn.Linear(feat_dim, feat_dim * 4),
            nn.GELU(),
            nn.Linear(feat_dim * 4, feat_dim)
        )

    def forward(self, img_local_feat, text_semantic_feat):
        """
        输入：img_local_feat-图像局部特征 [B, N, D]，text_semantic_feat-文本语义特征 [B, L, D]
        输出：细粒度融合特征（消除模态鸿沟）
        """
        # 交叉注意力融合：图像特征关注文本语义，文本特征关注图像局部
        attn_feat, _ = self.cross_attention(
            query=img_local_feat, key=text_semantic_feat, value=text_semantic_feat
        )
        # 残差+归一化+MLP增强
        fusion_feat = self.fusion_norm(attn_feat + img_local_feat)
        fusion_feat = self.fusion_mlp(fusion_feat) + fusion_feat
        return fusion_feat


class CrossDomainContrastiveLoss(nn.Module):
    def __init__(self, temperature=0.07):
        super().__init__()
        self.temp = temperature  # 对比学习温度系数

    def forward(self, source_feat, target_feat, domain_label):
        """
        跨域对比损失：拉近同源/同语义特征，推远异域特征，挖掘域不变特征
        source_feat: 源域融合特征
        target_feat: 目标域融合特征
        domain_label: 域标签（源域=0，目标域=1）
        """
        # 拼接源域+目标域特征，构建对比样本对
        all_feat = torch.cat([source_feat, target_feat], dim=0)
        all_label = torch.cat([domain_label, 1 - domain_label], dim=0)
        
        # 计算特征相似度矩阵
        sim_matrix = F.cosine_similarity(all_feat.unsqueeze(1), all_feat.unsqueeze(0), dim=-1)
        sim_matrix = torch.exp(sim_matrix / self.temp)

        # 正样本对（同域/同语义）、负样本对（异域）
        pos_mask = (all_label.unsqueeze(1) == all_label.unsqueeze(0)).float()
        neg_mask = 1 - pos_mask
        pos_sim = (sim_matrix * pos_mask).sum(dim=-1)
        neg_sim = (sim_matrix * neg_mask).sum(dim=-1)

        # 跨域对比损失（对齐域分布，挖掘域不变特征）
        contrast_loss = -torch.log(pos_sim / (pos_sim + neg_sim)).mean()
        return contrast_loss


class IntegratedAlgorithmFramework(nn.Module):
    def __init__(self):
        super().__init__()
        # 模块1：弱标注多源数据处理
        self.data_processor = WeakLabelDataProcessor()
        # 模块2：多模态交叉注意力融合
        self.fusion_model = MultiModalCrossAttentionFusion()
        # 模块3：跨域迁移（对比损失+域自适应）
        self.domain_align_loss = CrossDomainContrastiveLoss()
        # 细粒度分类头
        self.classifier = nn.Linear(512, 10)  # 按任务修改类别数

    def forward(self, source_img, source_text, target_img=None, target_text=None, is_train=True):
        """
        训练模式：源域+目标域数据，完成数据处理→融合→跨域迁移
        推理模式：仅源域数据，完成细粒度识别
        """
        # 1. 弱标注数据处理（自监督预训练+注意力筛选）
        src_img_feat, src_text_feat, src_mask = self.data_processor(source_img, source_text)
        
        # 2. 多模态细粒度融合
        src_fusion_feat = self.fusion_model(src_img_feat.unsqueeze(1), src_text_feat.unsqueeze(1))
        src_pred = self.classifier(src_fusion_feat.squeeze(1))

        # 训练模式：加入跨域迁移
        if is_train and target_img is not None and target_text is not None:
            tar_img_feat, tar_text_feat, tar_mask = self.data_processor(target_img, target_text)
            tar_fusion_feat = self.fusion_model(tar_img_feat.unsqueeze(1), tar_text_feat.unsqueeze(1))
            # 计算跨域对比损失（对齐源/目标域分布）
            domain_loss = self.domain_align_loss(
                src_fusion_feat.squeeze(1), 
                tar_fusion_feat.squeeze(1), 
                torch.zeros(src_fusion_feat.shape[0]).to(source_img.device)
            )
            return src_pred, domain_loss
        
        # 推理模式：仅输出分类结果
        return src_pred