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

class DistillGaze(nn.Module):
    """
    DistillGaze: 两阶段蒸馏眼动追踪
    
    Stage 1: DINOv3 → 域优化教师 (自蒸馏)
    Stage 2: 教师 → 256K学生 (跨架构蒸馏)
    
    关键创新：
    1. 设备感知增强（MTF/运动模糊/眩光）
    2. 混合监督VFM优化
    3. 特征VIC-KD + 预测伪标签
    """
    
    def __init__(self):
        super().__init__()
        
        # === Stage 1: VFM教师 ===
        self.teacher = DINOv3Encoder(freeze=False)
        
        # EMA教师（用于自蒸馏一致性）
        self.ema_teacher = DINOv3Encoder(freeze=True)
        self._update_ema()
        
        # === Stage 2: 边缘学生 ===
        self.student = FBNetStudent(model_size='256K')
        
    def _update_ema(self, decay=0.999):
        """EMA参数更新"""
        with torch.no_grad():
            for ema_p, p in zip(self.ema_teacher.parameters(), 
                                self.teacher.parameters()):
                ema_p.mul_(decay).add_(p * (1 - decay))
    
    def stage1_forward(self, synthetic_img, synthetic_label, 
                       real_img_weak, real_img_strong):
        """
        Stage 1: VFM域优化
        
        损失 = 合成监督 + 自蒸馏一致性
        """
        # 合成图像监督
        syn_feat = self.teacher(synthetic_img)
        syn_pred = self.gaze_head(syn_feat)
        sup_loss = F.mse_loss(syn_pred, synthetic_label)
        
        # 真实图像自蒸馏
        with torch.no_grad():
            real_feat_weak = self.ema_teacher(real_img_weak)
            real_pseudo = self.gaze_head(real_feat_weak)
        
        real_feat_strong = self.teacher(real_img_strong)
        real_pred = self.gaze_head(real_feat_strong)
        
        # 一致性损失
        cons_loss = F.mse_loss(real_pred, real_pseudo.detach())
        
        total_loss = sup_loss + 2.0 * cons_loss
        return total_loss
    
    def stage2_forward(self, img_weak, img_strong):
        """
        Stage 2: 边缘蒸馏
        
        教师(86M) → 学生(256K)
        损失 = 特征VIC-KD + 预测伪标签
        """
        with torch.no_grad():
            teacher_feat = self.teacher(img_weak)
            teacher_pred = self.gaze_head(teacher_feat)
        
        student_feat = self.student(img_strong)
        student_pred = self.gaze_head_student(student_feat)
        
        # 特征蒸馏 (VIC-KD)
        feat_loss = self.vic_kd_loss(student_feat, teacher_feat)
        
        # 预测蒸馏
        pred_loss = F.mse_loss(student_pred, teacher_pred.detach())
        
        total_loss = feat_loss + pred_loss
        return total_loss
    
    def vic_kd_loss(self, student_feat, teacher_feat):
        """VIC-KD: 跨架构特征蒸馏"""
        # 投影到公共空间
        s_proj = self.student_proj(student_feat)
        t_proj = self.teacher_proj(teacher_feat.detach())
        
        # 余弦相似度损失
        cos_loss = 1 - F.cosine_similarity(s_proj, t_proj, dim=-1).mean()
        return cos_loss


class DINOv3Encoder(nn.Module):
    """DINOv3 ViT-B 编码器 (86M参数)"""
    def __init__(self, freeze=True):
        super().__init__()
        self.vit = nn.Sequential(
            nn.Conv2d(3, 768, kernel_size=16, stride=16),
            *self._make_vit_layers(n_layers=12, d_model=768, n_heads=12)
        )
        if freeze:
            for p in self.parameters():
                p.requires_grad = False
    
    def _make_vit_layers(self, n_layers, d_model, n_heads):
        layers = []
        for _ in range(n_layers):
            layers.append(ViTBlock(d_model, n_heads))
        return layers
    
    def forward(self, x):
        return self.vit(x)


class FBNetStudent(nn.Module):
    """FBNet 轻量化学生模型 (256K参数)"""
    def __init__(self, model_size='256K'):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.BatchNorm2d(16), nn.ReLU6(),
            
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32), nn.ReLU6(),
            
            nn.Conv2d(32, 48, 3, stride=2, padding=1),
            nn.BatchNorm2d(48), nn.ReLU6(),
            
            nn.Conv2d(48, 64, 3, stride=1, padding=1),
            nn.BatchNorm2d(64), nn.ReLU6(),
            
            nn.AdaptiveAvgPool2d(1),
        )
        self.proj = nn.Linear(64, 64)
        
    def forward(self, x):
        return self.proj(self.features(x).flatten(1))


class ViTBlock(nn.Module):
    """简化版ViT Block"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model*4),
            nn.GELU(),
            nn.Linear(d_model*4, d_model)
        )
    
    def forward(self, x):
        B, C, H, W = x.shape
        x = x.flatten(2).transpose(1, 2)  # (B, HW, C)
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.ffn(self.norm2(x))
        return x.transpose(1, 2).view(B, C, H, W)


# === 设备感知增强 ===
class DeviceAwareAugmentation:
    """设备感知数据增强"""
    
    def __init__(self, strong=True):
        self.strong = strong
        self.p = 0.3  # 触发概率
    
    def __call__(self, img):
        if self.strong:
            # 强增强
            img = self._apply_camera_aug(img)
            img = self._apply_illumination_aug(img)
        else:
            # 弱增强
            img = self._gamma_jitter(img)
            img = self._random_scale(img)
        return img
    
    def _apply_camera_aug(self, img):
        """相机退化模拟"""
        if torch.rand(1) < self.p:
            # MTF滤波（调制传递函数）
            img = self._mtf_filter(img)
        if torch.rand(1) < self.p:
            # 运动模糊
            img = self._motion_blur(img)
        if torch.rand(1) < self.p:
            # 压缩伪影
            img = self._compression(img)
        return img
    
    def _apply_illumination_aug(self, img):
        """照明退化"""
        if torch.rand(1) < self.p:
            # 角膜眩光
            img = self._glint_inpaint(img)
        if torch.rand(1) < self.p:
            # 亮度/对比度
            img = self._brightness_contrast(img)
        return img
    
    def _gamma_jitter(self, img):
        gamma = torch.uniform(0.8, 1.2, (1,)).item()
        return img ** gamma
    
    def _random_scale(self, img):
        s = torch.uniform(0.9, 1.1, (1,)).item()
        return F.interpolate(img.unsqueeze(0), scale_factor=s).squeeze(0)
    
    def _mtf_filter(self, img):
        return img  # 简化
    
    def _motion_blur(self, img):
        return img  # 简化
    
    def _compression(self, img):
        return img  # 简化
    
    def _glint_inpaint(self, img):
        return img  # 简化
    
    def _brightness_contrast(self, img):
        return img * torch.uniform(0.8, 1.2, (1,)).item()


# 测试
if __name__ == "__main__":
    model = DistillGaze()
    
    # 模拟输入
    B = 4
    syn_img = torch.randn(B, 3, 64, 64)
    syn_label = torch.randn(B, 2)  # [yaw, pitch]
    real_weak = torch.randn(B, 3, 64, 64)
    real_strong = torch.randn(B, 3, 64, 64)
    
    # Stage 1
    loss1 = model.stage1_forward(syn_img, syn_label, real_weak, real_strong)
    print(f"Stage 1 损失: {loss1.item():.4f}")
    
    # Stage 2
    loss2 = model.stage2_forward(real_weak, real_strong)
    print(f"Stage 2 损失: {loss2.item():.4f}")
    
    # 学生模型大小
    student_params = sum(p.numel() for p in model.student.parameters())
    teacher_params = sum(p.numel() for p in model.teacher.parameters())
    print(f"\n教师参数量: {teacher_params/1e6:.1f}M")
    print(f"学生参数量: {student_params/1e3:.1f}K")
    print(f"压缩比: {teacher_params/student_params:.0f}x")
    
    # 增强测试
    aug = DeviceAwareAugmentation(strong=True)
    augmented = aug(syn_img[0])
    print(f"\n增强后形状: {augmented.shape}")