import torch
import torch.nn as nn
from typing import Tuple

class TrajectoryAwareKinematicEncoder(nn.Module):
    """
    EyeTAG论文 Section 3.2 的TAKE模块
    
    将过去预测的差分序列编码为运动学token
    
    Args:
        prev_dim: 历史窗口大小 T
        gaze_dim: 视线向量维度 (默认3, 3D单位向量)
        hidden_dim: MLP隐藏层维度
        out_dim: 输出token维度
    """
    
    def __init__(self, prev_dim: int = 48, gaze_dim: int = 3,
                 hidden_dim: int = 256, out_dim: int = 256):
        super().__init__()
        self.prev_dim = prev_dim
        self.gaze_dim = gaze_dim
        
        # 差分计算：将绝对历史转为差分序列
        # Δg_t = g_t - g_{t-1}
        
        # MLP编码差分序列
        self.delta_mlp = nn.Sequential(
            nn.Linear(prev_dim * gaze_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, out_dim)
        )
        
        # LayerNorm 稳定训练
        self.norm = nn.LayerNorm(out_dim)
        
    def compute_deltas(self, prev_gazes: torch.Tensor) -> torch.Tensor:
        """
        计算差分序列
        
        Args:
            prev_gazes: 过去T帧的视线预测, shape=(B, T, 3)
            
        Returns:
            deltas: 差分序列, shape=(B, T-1, 3)
        """
        # Δg_t = g_t - g_{t-1}
        deltas = prev_gazes[:, 1:] - prev_gazes[:, :-1]
        return deltas
    
    def forward(self, prev_gazes: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        
        Args:
            prev_gazes: 过去T帧视线预测, shape=(B, T, 3)
            
        Returns:
            kinematic_token: 运动学token, shape=(B, out_dim)
        """
        B = prev_gazes.shape[0]
        
        # 步骤1: 计算差分
        deltas = self.compute_deltas(prev_gazes)  # (B, T-1, 3)
        
        # 步骤2: 展平
        deltas_flat = deltas.reshape(B, -1)  # (B, (T-1)*3)
        
        # 步骤3: MLP编码
        token = self.delta_mlp(deltas_flat)  # (B, out_dim)
        
        # 步骤4: LayerNorm
        token = self.norm(token)
        
        return token


# 测试
if __name__ == "__main__":
    # 模拟参数
    B, T, D = 4, 48, 3  # batch=4, 窗口=48帧, 3D视线
    
    # 创建TAKE模块
    take = TrajectoryAwareKinematicEncoder(prev_dim=T-1, gaze_dim=D)
    
    # 模拟过去48帧的视线预测 (3D单位向量)
    prev_gazes = torch.randn(B, T, D)
    prev_gazes = prev_gazes / prev_gazes.norm(dim=-1, keepdim=True)  # 归一化
    
    # 前向传播
    token = take(prev_gazes)
    print(f"输入: prev_gazes shape={prev_gazes.shape}")
    print(f"差分序列: deltas shape={take.compute_deltas(prev_gazes).shape}")
    print(f"输出: kinematic_token shape={token.shape}")
    print(f"参数量: {sum(p.numel() for p in take.parameters()):,}")

class EyeTAGFusionModule(nn.Module):
    """
    EyeTAG论文 Section 3.3 的融合模块
    
    使用交叉注意力对齐面部和眼部特征，
    再通过因果Transformer解码器融合运动学token
    """
    
    def __init__(self, face_dim=2048, eye_dim=512, kinematic_dim=256,
                 hidden_dim=256, num_heads=4, num_layers=2, max_len=48):
        super().__init__()
        
        # 维度对齐
        self.face_proj = nn.Linear(face_dim, hidden_dim)
        self.eye_proj = nn.Linear(eye_dim, hidden_dim)
        
        # 交叉注意力: Eye查询Face
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            batch_first=True
        )
        
        # 因果Transformer解码器
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            batch_first=True,
            activation='gelu'
        )
        self.transformer_decoder = nn.TransformerDecoder(
            decoder_layer,
            num_layers=num_layers
        )
        
        # 输出头: 3D视线方向
        self.gaze_head = nn.Linear(hidden_dim, 3)
        
        # 位置编码
        self.pos_embed = nn.Parameter(torch.randn(1, max_len, hidden_dim) * 0.02)
        
        # 因果掩码
        self.register_buffer(
            'causal_mask',
            torch.triu(torch.ones(max_len, max_len) * float('-inf'), diagonal=1)
        )
    
    def forward(self, face_feat, eye_feat, kinematic_token, 
                seq_len=None):
        """
        Args:
            face_feat: (B, T, face_dim) 面部特征序列
            eye_feat: (B, T, eye_dim) 眼部特征序列
            kinematic_token: (B, hidden_dim) 运动学token
            seq_len: 实际序列长度
        Returns:
            gaze_3d: (B, 3) 3D视线方向
        """
        B = face_feat.shape[0]
        T = seq_len if seq_len is not None else face_feat.shape[1]
        
        # 维度对齐
        face = self.face_proj(face_feat)  # (B, T, hidden)
        eye = self.eye_proj(eye_feat)    # (B, T, hidden)
        
        # 交叉注意力: Eye查询Face
        attn_out, _ = self.cross_attn(
            query=eye,    # (B, T, hidden)
            key=face,     # (B, T, hidden)
            value=face    # (B, T, hidden)
        )
        
        # 融合后的视觉特征
        visual_feat = eye + attn_out  # 残差连接
        
        # 位置编码
        visual_feat = visual_feat + self.pos_embed[:, :T, :]
        
        # 运动学token扩展为序列前缀
        kinematic_seq = kinematic_token.unsqueeze(1)  # (B, 1, hidden)
        memory = torch.cat([kinematic_seq, visual_feat], dim=1)  # (B, T+1, hidden)
        
        # 因果Transformer解码
        # tgt是memory本身（自回归）
        causal_out = self.transformer_decoder(
            tgt=memory,
            memory=memory,
            tgt_mask=self.causal_mask[:memory.shape[1], :memory.shape[1]]
        )
        
        # 取最后一帧的输出预测3D视线
        last_hidden = causal_out[:, -1, :]  # (B, hidden)
        gaze_3d = self.gaze_head(last_hidden)  # (B, 3)
        
        # 归一化为单位向量
        gaze_3d = gaze_3d / (gaze_3d.norm(dim=-1, keepdim=True) + 1e-8)
        
        return gaze_3d


# 测试
if __name__ == "__main__":
    B, T = 4, 48
    
    fusion = EyeTAGFusionModule()
    
    # 模拟特征
    face_feat = torch.randn(B, T, 2048)
    eye_feat = torch.randn(B, T, 512)
    kin_token = torch.randn(B, 256)
    
    gaze = fusion(face_feat, eye_feat, kin_token)
    print(f"Face特征: {face_feat.shape}")
    print(f"Eye特征: {eye_feat.shape}")
    print(f"运动学token: {kin_token.shape}")
    print(f"输出3D视线: {gaze.shape}")
    print(f"视线向量范数: {gaze.norm(dim=-1)}")
    print(f"融合模块参数量: {sum(p.numel() for p in fusion.parameters()):,}")

class ScheduledSampling:
    """
    EyeTAG论文 Section 3.4 的计划采样策略
    
    训练时，输入TAKE的视线历史以概率α来自真实标签，
    否则来自模型自身预测缓存。α从1.0线性退火到0.0。
    
    关键：缓存每个epoch重建一次，无需在线rollout
    """
    
    def __init__(self, total_epochs=5, start_alpha=1.0, end_alpha=0.0):
        self.start_alpha = start_alpha
        self.end_alpha = end_alpha
        self.total_epochs = total_epochs
    
    def get_alpha(self, epoch: int) -> float:
        """获取当前epoch的α值"""
        if epoch >= self.total_epochs:
            return self.end_alpha
        progress = epoch / self.total_epochs
        alpha = self.start_alpha - (self.start_alpha - self.end_alpha) * progress
        return max(alpha, self.end_alpha)
    
    def sample_history(self, gt_gazes, pred_cache, alpha):
        """
        混合真实和预测的历史
        
        Args:
            gt_gazes: (B, T, 3) 真实视线标签
            pred_cache: (B, T, 3) 模型预测缓存
            alpha: 使用真实标签的概率
            
        Returns:
            mixed_history: (B, T, 3) 混合历史
        """
        B, T, D = gt_gazes.shape
        # 每帧独立采样
        use_gt = torch.rand(B, T, 1, device=gt_gazes.device) < alpha
        mixed = torch.where(use_gt, gt_gazes, pred_cache)
        return mixed


# 退火曲线示例
if __name__ == "__main__":
    import matplotlib
    matplotlib.use('Agg')
    
    ss = ScheduledSampling(total_epochs=5)
    print("Epoch | α值   | 真实标签比例")
    print("------|-------|------------")
    for ep in range(8):
        alpha = ss.get_alpha(ep)
        print(f"  {ep:2d}  | {alpha:.2f}  | {alpha*100:.0f}%")

"""
EyeTAG 完整推理示例

论文：EyeTAG: Eye Trajectory-Aware Gaze Estimation
会议：BMVC 2026
arXiv: 2610.00922
GitHub: https://github.com/peter8366/EyeTAG

依赖：
    pip install torch torchvision decord h5py
    
    # 需要VGGFace2预训练权重
    # 下载: https://github.com/cydonia999/VGGFace2-pytorch
    # 放置于: checkpoints/resnet50_ft_weight.pkl
"""

import torch
import torch.nn as nn
import torchvision.models as models
from typing import Optional
import numpy as np


class EyeTAGInference(nn.Module):
    """
    EyeTAG完整推理模型
    
    使用方法:
        model = EyeTAGInference(num_frames=48)
        model.load_state_dict(torch.load('eyetag_gaze360.pth'))
        
        # 流式推理
        face, left_eye, right_eye = preprocess(frame)
        gaze = model(face, left_eye, right_eye)
    """
    
    def __init__(self, num_frames: int = 48):
        super().__init__()
        self.num_frames = num_frames
        
        # Face Encoder: ResNet-50 (VGGFace2预训练)
        self.face_encoder = models.resnet50(weights=None)
        self.face_encoder.fc = nn.Identity()  # 移除分类头
        self.face_proj = nn.Linear(2048, 256)
        
        # Eye Encoder: ResNet-18
        self.eye_encoder = models.resnet18(weights=None)
        self.eye_encoder.fc = nn.Identity()
        self.eye_proj = nn.Linear(512, 256)
        
        # 运动学编码器
        self.take = TrajectoryAwareKinematicEncoder(
            prev_dim=num_frames - 1,
            gaze_dim=3,
            hidden_dim=256,
            out_dim=256
        )
        
        # 融合模块
        self.fusion = EyeTAGFusionModule(
            face_dim=256,
            eye_dim=256,
            kinematic_dim=256,
            hidden_dim=256,
            num_heads=4,
            num_layers=2,
            max_len=num_frames
        )
        
        # 历史预测缓存（流式推理用）
        self.register_buffer(
            'gaze_history',
            torch.zeros(1, num_frames, 3)
        )
    
    def forward(self, face: torch.Tensor, 
                left_eye: torch.Tensor,
                right_eye: torch.Tensor) -> torch.Tensor:
        """
        单帧流式推理
        
        Args:
            face: (1, 3, 224, 224) 面部裁剪
            left_eye: (1, 3, 64, 64) 左眼裁剪
            right_eye: (1, 3, 64, 64) 右眼裁剪
            
        Returns:
            gaze_3d: (1, 3) 3D视线单位向量
        """
        # 编码
        face_feat = self.face_proj(self.face_encoder(face))  # (1, 256)
        eye_feat = self.eye_proj(
            self.eye_encoder(torch.cat([left_eye, right_eye], dim=0).mean(dim=0, keepdim=True))
        )  # (1, 256)
        
        # 扩展为序列形式（流式：只取当前帧）
        face_seq = face_feat.unsqueeze(1)  # (1, 1, 256)
        eye_seq = eye_feat.unsqueeze(1)    # (1, 1, 256)
        
        # 运动学token
        kin_token = self.take(self.gaze_history)  # (1, 256)
        
        # 融合预测
        gaze = self.fusion(face_seq, eye_seq, kin_token, seq_len=1)  # (1, 3)
        
        # 更新历史缓存（滚动窗口）
        with torch.no_grad():
            self.gaze_history = torch.cat([
                self.gaze_history[:, 1:, :],
                gaze.detach()
            ], dim=1)
        
        return gaze
    
    def reset_history(self):
        """重置历史缓存（新被试时调用）"""
        self.gaze_history.zero_()


def preprocess_frame(frame: np.ndarray, 
                     face_bbox: tuple,
                     eye_bboxes: dict) -> tuple:
    """
    预处理单帧图像
    
    Args:
        frame: (H, W, 3) BGR图像
        face_bbox: (x, y, w, h) 面部边界框
        eye_bboxes: {'left': (x,y,w,h), 'right': (x,y,w,h)}
    
    Returns:
        face_tensor: (1, 3, 224, 224)
        left_eye_tensor: (1, 3, 64, 64)
        right_eye_tensor: (1, 3, 64, 64)
    """
    import cv2
    
    def crop_resize(img, bbox, size):
        x, y, w, h = bbox
        crop = img[y:y+h, x:x+w]
        crop = cv2.resize(crop, (size, size))
        crop = crop[:, :, ::-1]  # BGR -> RGB
        crop = crop.astype(np.float32) / 255.0
        # ImageNet标准化
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        crop = (crop - mean) / std
        crop = np.transpose(crop, (2, 0, 1))
        return torch.from_numpy(crop).unsqueeze(0).float()
    
    face = crop_resize(frame, face_bbox, 224)
    left = crop_resize(frame, eye_bboxes['left'], 64)
    right = crop_resize(frame, eye_bboxes['right'], 64)
    
    return face, left, right


# 完整演示
if __name__ == "__main__":
    print("=" * 60)
    print("EyeTAG 推理演示")
    print("=" * 60)
    
    # 初始化模型
    model = EyeTAGInference(num_frames=48)
    model.eval()
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,} ({total_params/1e6:.2f}M)")
    print(f"预期: ~37.91M (论文报告)")
    
    # 模拟单帧推理
    face = torch.randn(1, 3, 224, 224)
    left_eye = torch.randn(1, 3, 64, 64)
    right_eye = torch.randn(1, 3, 64, 64)
    
    # 重置历史
    model.reset_history()
    
    # 流式推理10帧
    print("\n流式推理10帧:")
    print(f"{'帧':>4} | {'pitch(°)':>10} | {'yaw(°)':>10} | {'|g|':>8}")
    print("-" * 42)
    
    for i in range(10):
        with torch.no_grad():
            gaze = model(face, left_eye, right_eye)
        
        # 转为角度
        pitch = torch.asin(gaze[0, 1]).item() * 180 / np.pi
        yaw = torch.atan2(gaze[0, 0], gaze[0, 2]).item() * 180 / np.pi
        norm = gaze[0].norm().item()
        
        print(f"{i:4d} | {pitch:10.2f} | {yaw:10.2f} | {norm:8.6f}")
    
    print(f"\n历史缓存状态: shape={model.gaze_history.shape}")
    print(f"历史缓存非零帧数: {(model.gaze_history.abs().sum(-1) > 0).sum().item()}")

# 边缘部署优化建议
# 1. 视觉编码器缓存（流式部署）
class CachedEyeTAG:
    """
    流式部署优化：视觉编码器跨帧缓存
    
    论文说明：GFLOPs figure是无缓存上限
    流式部署中，视觉编码器可跨滑动窗口缓存
    """
    
    def __init__(self, model, cache_size=48):
        self.model = model
        self.cache_size = cache_size
        self.face_cache = []  # 缓存面部特征
        self.eye_cache = []   # 缓存眼部特征
    
    def step(self, face_img, left_eye, right_eye):
        # 编码当前帧
        with torch.no_grad():
            face_feat = self.model.face_proj(
                self.model.face_encoder(face_img)
            )
            eye_feat = self.model.eye_proj(
                self.model.eye_encoder(
                    torch.cat([left_eye, right_eye]).mean(0, keepdim=True)
                )
            )
        
        # 更新缓存
        self.face_cache.append(face_feat)
        self.eye_cache.append(eye_feat)
        if len(self.face_cache) > self.cache_size:
            self.face_cache.pop(0)
            self.eye_cache.pop(0)
        
        # 只计算运动学token和融合
        # 省去重复的视觉编码计算
        # 实际FLOPs可降低到原来的1/48
        pass