"""
Hierarchical Recurrent ANN (HR-ANN)
自上而下反馈的音视觉整合模型

基于 eLife 2026 论文

核心思想：
1. 视觉和音频各自有层级特征提取器
2. 高层语义通过反馈通路调制低层感知
3. 跨模态预测：视觉预测音频包络，音频预测视觉运动
"""

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

class FeedbackLayer(nn.Module):
    """
    反馈层：接收高层语义，生成调制信号
    
    模拟大脑皮层的反馈投射：
    - 注意力调制（gain control）
    - 预测编码（prediction error）
    - 跨模态预测
    """
    
    def __init__(self, feat_dim: int, context_dim: int):
        super().__init__()
        # 注意力调制（gain field）
        self.gate = nn.Linear(context_dim, feat_dim)
        # 预测编码（生成对低层的预测）
        self.predictor = nn.Linear(context_dim, feat_dim)
        # 跨模态预测器
        self.cross_modal = nn.Linear(context_dim, feat_dim)
        
    def forward(self, low_feat: torch.Tensor, 
                context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            low_feat: 低层特征 (B, T, F)
            context: 高层语义上下文 (B, T, C)
            
        Returns:
            modulated: 调制后的低层特征
            prediction_error: 预测误差
        """
        # 注意力增益
        gate = torch.sigmoid(self.gate(context))  # (B, T, F)
        modulated = low_feat * gate  # 增益控制
        
        # 预测编码
        prediction = self.predictor(context)
        prediction_error = low_feat - prediction  # 残差
        
        return modulated, prediction_error


class HRANN(nn.Module):
    """
    完整的 HR-ANN 模型
    
    层级结构：
    L1: 低层感知（视觉/音频各自提取）
    L2: 中层整合（跨模态交互开始）
    L3: 高层语义（整合判定）
    
    反馈：L3 → L2 → L1（自上而下调制）
    """
    
    def __init__(self, visual_dim: int = 512, audio_dim: int = 128,
                 hidden_dim: int = 256, n_classes: int = 5):
        super().__init__()
        
        # 前馈：视觉通道
        self.v_enc1 = nn.Linear(visual_dim, hidden_dim)
        self.v_enc2 = nn.Linear(hidden_dim, hidden_dim)
        
        # 前馈：音频通道
        self.a_enc1 = nn.Linear(audio_dim, hidden_dim)
        self.a_enc2 = nn.Linear(hidden_dim, hidden_dim)
        
        # 整合层
        self.fusion = nn.Linear(hidden_dim * 2, hidden_dim)
        
        # 高层语义
        self.semantic = nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
        
        # 反馈层
        self.fb_l2_visual = FeedbackLayer(hidden_dim, hidden_dim)
        self.fb_l2_audio = FeedbackLayer(hidden_dim, hidden_dim)
        self.fb_l1_visual = FeedbackLayer(hidden_dim, hidden_dim)
        self.fb_l1_audio = FeedbackLayer(hidden_dim, hidden_dim)
        
        # 输出
        self.classifier = nn.Linear(hidden_dim, n_classes)
        
        # 循环迭代次数
        self.n_iterations = 3  # 自上而下反馈迭代
        
    def forward(self, visual: torch.Tensor, audio: torch.Tensor) -> Tuple[torch.Tensor, dict]:
        """
        Args:
            visual: (B, T, visual_dim)
            audio: (B, T, audio_dim)
            
        Returns:
            logits: (B, T, n_classes)
            info: 额外信息（预测误差等）
        """
        B, T, _ = visual.shape
        
        # L1: 低层特征提取
        v_l1 = F.gelu(self.v_enc1(visual))
        a_l1 = F.gelu(self.a_enc1(audio))
        
        # L2: 中层特征
        v_l2 = F.gelu(self.v_enc2(v_l1))
        a_l2 = F.gelu(self.a_enc2(a_l1))
        
        # L3: 高层语义（初始前馈）
        fused = self.fusion(torch.cat([v_l2, a_l2], dim=-1))
        semantic_out, _ = self.semantic(fused)
        
        total_pred_error = 0
        
        # 自上而下反馈迭代
        for i in range(self.n_iterations):
            # L3 → L2 反馈
            v_l2_mod, v_pe2 = self.fb_l2_visual(v_l2, semantic_out)
            a_l2_mod, a_pe2 = self.fb_l2_audio(a_l2, semantic_out)
            
            # L2 → L1 反馈
            v_l1_mod, v_pe1 = self.fb_l1_visual(v_l1, v_l2_mod)
            a_l1_mod, a_pe1 = self.fb_l1_audio(a_l1, a_l2_mod)
            
            # 重新前馈
            v_l2 = F.gelu(self.v_enc2(v_l1_mod))
            a_l2 = F.gelu(self.a_enc2(a_l1_mod))
            
            fused = self.fusion(torch.cat([v_l2, a_l2], dim=-1))
            semantic_out, _ = self.semantic(fused)
            
            total_pred_error += v_pe2.mean() + a_pe2.mean() + v_pe1.mean() + a_pe1.mean()
        
        # 分类
        logits = self.classifier(semantic_out)
        
        return logits, {
            'prediction_error': total_pred_error / self.n_iterations,
            'semantic_repr': semantic_out,
        }


# 测试
if __name__ == "__main__":
    model = HRANN(visual_dim=512, audio_dim=128, n_classes=5)
    
    # 模拟输入
    visual = torch.randn(4, 30, 512)  # batch=4, 30帧, 512维视觉
    audio = torch.randn(4, 30, 128)   # batch=4, 30帧, 128维音频
    
    logits, info = model(visual, audio)
    print(f"Logits: {logits.shape}")
    print(f"Prediction error: {info['prediction_error']:.4f}")
    print(f"Semantic repr: {info['semantic_repr'].shape}")
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")