import torch
import torch.nn as nn

class LiteFatModel(nn.Module):
    """
    LiteFat 轻量级时空图疲劳检测
    
    输入: (B, T, 68, 2) 关键点序列
    输出: (B, 2) 疲劳/清醒
    """
    
    def __init__(self, n_keypoints=68, hidden=64):
        super().__init__()
        
        # 空间图卷积（关键点间关系）
        self.spatial_gcn = nn.Sequential(
            nn.Linear(2, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
        )
        
        # 时间编码（帧间动态）
        self.temporal_encoder = nn.GRU(
            input_size=hidden * n_keypoints,
            hidden_size=hidden,
            num_layers=1,
            batch_first=True,
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden, 32), nn.ReLU(),
            nn.Linear(32, 2)
        )
    
    def forward(self, kpts_seq):
        """
        Args:
            kpts_seq: (B, T, 68, 2) 关键点序列
            
        Returns:
            logits: (B, 2) 疲劳/清醒
        """
        B, T, K, D = kpts_seq.shape
        
        # 空间编码
        x = self.spatial_gcn(kpts_seq)  # (B, T, K, hidden)
        x = x.reshape(B, T, -1)  # (B, T, K*hidden)
        
        # 时间编码
        out, hn = self.temporal_encoder(x)  # (B, T, hidden)
        
        # 取最后状态
        feat = hn[-1]  # (B, hidden)
        
        # 分类
        logits = self.classifier(feat)  # (B, 2)
        
        return logits


# 测试
if __name__ == "__main__":
    model = LiteFatModel()
    
    # 30帧关键点序列
    x = torch.randn(1, 30, 68, 2)
    out = model(x)
    
    print(f"输入: {x.shape}")
    print(f"输出: {out.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")
    print(f"模型大小(INT8): ~{sum(p.numel() for p in model.parameters()) // 1024} KB")