import torch
import torch.nn as nn
import torch.fft as fft

class FrequencyAttention(nn.Module):
    """
    频率域注意力模块
    
    核心思想：光照变化主要影响低频分量，
    高频分量（边缘/纹理）对光照更鲁棒
    """
    
    def __init__(self, channels):
        super().__init__()
        self.freq_weight = nn.Sequential(
            nn.Conv2d(channels, channels//4, 1),
            nn.ReLU(),
            nn.Conv2d(channels//4, channels, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        # FFT 变换
        freq = fft.fft2(x)
        freq_shifted = fft.fftshift(freq)
        mag = torch.abs(freq_shifted)
        
        # 频率权重
        weight = self.freq_weight(mag)
        
        # 加权
        freq_weighted = freq_shifted * weight
        
        # 逆变换
        freq_unshifted = fft.ifftshift(freq_weighted)
        out = fft.ifft2(freq_unshifted).real
        
        return out + x  # 残差


class SpatialFrequencyAttention(nn.Module):
    """空间-频率双域注意力"""
    
    def __init__(self, channels):
        super().__init__()
        self.spatial_attn = nn.Sequential(
            nn.Conv2d(channels, channels//4, 1),
            nn.ReLU(),
            nn.Conv2d(channels//4, channels, 1),
            nn.Sigmoid()
        )
        self.freq_attn = FrequencyAttention(channels)
    
    def forward(self, x):
        s_attn = self.spatial_attn(x)
        x_spatial = x * s_attn
        x_freq = self.freq_attn(x)
        
        # 融合
        return x_spatial + x_freq

class LanguageGuidedAttention(nn.Module):
    """
    语言引导注意力模块
    
    用场景描述文本引导视线区域分类
    例如："驾驶员低头看手机" → 关注手机区域
    """
    
    SCENE_DESCRIPTIONS = [
        "driver looking at road ahead",
        "driver looking at center console",
        "driver looking at phone",
        "driver looking at left mirror",
        "driver looking at right mirror",
        "driver looking down at lap",
        "driver looking at passenger",
        "driver eyes closed",
    ]
    
    def __init__(self, n_classes=8, text_dim=768):
        super().__init__()
        # 文本编码器（冻结的 CLIP）
        self.text_encoder = self._load_clip()
        
        # 跨模态注意力
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=512, num_heads=8,
            kdim=text_dim, vdim=text_dim
        )
        
        # 分类头
        self.classifier = nn.Linear(512, n_classes)
    
    def _load_clip(self):
        """加载预训练 CLIP 文本编码器"""
        # 简化：实际使用 transformers 库
        return nn.Identity()
    
    def forward(self, visual_feat, text_embed):
        """
        Args:
            visual_feat: (B, seq, 512) 视觉特征序列
            text_embed: (B, text_dim) 场景描述嵌入
        """
        # 跨模态注意力
        text = text_embed.unsqueeze(1)  # (B, 1, text_dim)
        out, _ = self.cross_attn(
            query=visual_feat,
            key=text, value=text
        )
        
        # 分类
        pooled = out.mean(dim=1)
        logits = self.classifier(pooled)
        
        return logits


# 完整 LISA 模型
class LISA(nn.Module):
    """LISA: Language-guided Interference-aware Spatial-Frequency Attention"""
    
    def __init__(self, n_gaze_classes=8):
        super().__init__()
        # 骨干网络
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
        )
        
        # 空间-频率注意力
        self.sfa = SpatialFrequencyAttention(256)
        
        # 语言引导
        self.lang_guide = LanguageGuidedAttention(n_gaze_classes)
        
    def forward(self, image, text_embed):
        feat = self.backbone(image)  # (B, 256, 4, 4)
        feat = self.sfa(feat)
        feat_flat = feat.flatten(2).permute(2, 0, 1)  # (seq, B, 256)
        
        # 投影到 512 维
        proj = nn.Linear(256, 512).to(feat.device)
        feat_proj = proj(feat_flat)  # (seq, B, 512)
        
        logits = self.lang_guide(feat_proj, text_embed)
        return logits


# 测试
if __name__ == "__main__":
    model = LISA()
    
    # 模拟输入
    image = torch.randn(2, 3, 128, 128)
    text_embed = torch.randn(2, 768)
    
    logits = model(image, text_embed)
    print(f"输出: {logits.shape}")  # (2, 8) → 8个视线区域
    print(f"参数: {sum(p.numel() for p in model.parameters()):,}")

class IMSCameraBackbone(nn.Module):
    """IMS 摄像头骨干网络 + 频域注意力"""
    
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
        )
        self.freq_attn = FrequencyAttention(128)
        
    def forward(self, x):
        feat = self.backbone(x)
        feat = self.freq_attn(feat)
        return feat