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

class PhysicalEyeRenderer:
    """
    物理眼渲染器
    
    功能：给定4D gaze标签，生成对应的瞳孔/虹膜几何
    
    原理：
    1. 3D眼球模型（球体+虹膜圆盘）
    2. gaze方向 → 眼球旋转 → 瞳孔在2D投影中的位置和形状
    3. 变化解剖学参数（角膜曲率、瞳孔大小、眼睑开度）
    4. 变化相机参数（视角、距离）
    
    这实现了"无需源图像的多样合成"
    """
    
    def __init__(self, eye_radius=12e-3, iris_radius=5e-3,
                 pupil_radius_range=(1e-3, 4e-3)):
        self.eye_radius = eye_radius  # 眼球半径 ~12mm
        self.iris_radius = iris_radius  # 虹膜半径 ~5mm
        self.pupil_range = pupil_radius_range
        
    def gaze_to_eyeball_rotation(self, pitch: float, yaw: float) -> np.ndarray:
        """
        将gaze (pitch, yaw)转换为眼球3D旋转矩阵
        
        Args:
            pitch: 俯仰角 (rad)
            yaw: 偏航角 (rad)
        
        Returns:
            R: (3, 3) 旋转矩阵
        """
        # Pitch绕X轴，Yaw绕Y轴
        Rx = np.array([
            [1, 0, 0],
            [0, np.cos(pitch), -np.sin(pitch)],
            [0, np.sin(pitch), np.cos(pitch)]
        ])
        Ry = np.array([
            [np.cos(yaw), 0, np.sin(yaw)],
            [0, 1, 0],
            [-np.sin(yaw), 0, np.cos(yaw)]
        ])
        return Rx @ Ry
    
    def project_pupil(self, R: np.ndarray, camera_dist: float = 0.5,
                      focal: float = 800, cx: float = 320, 
                      cy: float = 240) -> Tuple[np.ndarray, float]:
        """
        投影瞳孔到2D图像平面
        
        Args:
            R: 眼球旋转矩阵
            camera_dist: 相机距离(m)
            focal: 焦距(px)
            cx, cy: 主点
        
        Returns:
            pupil_center_2d: (2,) 瞳孔中心2D坐标
            pupil_radius_2d: 投影后瞳孔半径
        """
        # 瞳孔在眼球坐标系中的位置（前方）
        pupil_3d = np.array([0, 0, self.eye_radius])
        
        # 旋转
        pupil_rotated = R @ pupil_3d
        
        # 平移到相机坐标系
        pupil_camera = pupil_rotated + np.array([0, 0, -camera_dist])
        
        # 透视投影
        if pupil_camera[2] >= 0:
            return np.array([cx, cy]), 0  # 瞳孔在相机后方
        
        pupil_2d_x = focal * pupil_camera[0] / (-pupil_camera[2]) + cx
        pupil_2d_y = focal * pupil_camera[1] / (-pupil_camera[2]) + cy
        pupil_center = np.array([pupil_2d_x, pupil_2d_y])
        
        # 投影后瞳孔半径（简化：假设正圆投影）
        pupil_radius_3d = 2e-3  # 2mm
        pupil_radius_2d = focal * pupil_radius_3d / (-pupil_camera[2])
        
        return pupil_center, pupil_radius_2d
    
    def render_geometry_map(self, pitch_l: float, yaw_l: float,
                            pitch_r: float, yaw_r: float,
                            img_size=(256, 256)) -> np.ndarray:
        """
        生成空间几何条件图（用于扩散模型条件化）
        
        Args:
            pitch_l, yaw_l: 左眼gaze
            pitch_r, yaw_r: 右眼gaze
            img_size: 输出图像大小
        
        Returns:
            geometry_map: (H, W, C) 几何条件图
        """
        H, W = img_size
        cx, cy = W // 4, H // 2  # 左眼中心
        
        # 左眼
        R_l = self.gaze_to_eyeball_rotation(pitch_l, yaw_l)
        pupil_l, radius_l = self.project_pupil(R_l, cx=cx, cy=cy)
        
        # 右眼
        cx_r = 3 * W // 4
        R_r = self.gaze_to_eyeball_rotation(pitch_r, yaw_r)
        pupil_r, radius_r = self.project_pupil(R_r, cx=cx_r, cy=cy)
        
        # 创建几何条件图
        # 通道1: 左眼瞳孔位置+半径
        # 通道2: 右眼瞳孔位置+半径
        # 通道3: 虹膜mask
        geom_map = np.zeros((H, W, 3), dtype=np.float32)
        
        # 绘制左眼瞳孔
        Y, X = np.ogrid[:H, :W]
        dist_l = np.sqrt((Y - pupil_l[1])**2 + (X - pupil_l[0])**2)
        geom_map[:, :, 0] = (dist_l < radius_l).astype(np.float32)
        geom_map[:, :, 1] = (dist_l < radius_l * 2.5).astype(np.float32)  # 虹膜
        
        # 绘制右眼瞳孔
        dist_r = np.sqrt((Y - pupil_r[1])**2 + (X - pupil_r[0])**2)
        geom_map[:, :, 0] += (dist_r < radius_r).astype(np.float32) * 0.5
        geom_map[:, :, 2] = (dist_r < radius_r * 2.5).astype(np.float32)  # 虹膜
        
        return geom_map
    
    def sample_diverse_geometry(self, pitch_l, yaw_l, pitch_r, yaw_r,
                                 n_samples=10):
        """
        通过变化解剖学参数生成多样化几何
        
        这实现了"无需源图像的多样合成"
        
        Args:
            n_samples: 采样数量
        
        Returns:
            geometries: list of geometry maps
        """
        geometries = []
        for i in range(n_samples):
            # 随机变化解剖学参数
            pupil_size = np.random.uniform(*self.pupil_range)
            eye_openness = np.random.uniform(0.6, 1.0)
            camera_angle = np.random.uniform(-5, 5)  # degrees
            
            # 生成几何
            geom = self.render_geometry_map(
                pitch_l, yaw_l, pitch_r, yaw_r
            )
            geometries.append(geom)
        
        return geometries


class GazeDiTConditioning(nn.Module):
    """
    GazeDiT条件化模块
    
    将4D gaze标签和几何条件图融合为扩散模型的条件信号
    """
    
    def __init__(self, gaze_dim=4, geom_channels=3,
                 hidden_dim=256):
        super().__init__()
        
        # 4D gaze嵌入
        self.gaze_embed = nn.Sequential(
            nn.Linear(gaze_dim, 64),
            nn.GELU(),
            nn.Linear(64, 128)
        )
        
        # 几何条件图编码（简化版，实际用SegFormer）
        self.geom_encoder = nn.Sequential(
            nn.Conv2d(geom_channels, 32, 3, stride=2, padding=1),
            nn.GELU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.GELU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.GELU(),
            nn.AdaptiveAvgPool2d(1)
        )
        
        # 融合
        self.fusion = nn.Sequential(
            nn.Linear(128 + 128, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
    
    def forward(self, gaze_4d: torch.Tensor, 
                geom_map: torch.Tensor) -> torch.Tensor:
        """
        Args:
            gaze_4d: (B, 4) 4D binocular gaze [pitch_l, yaw_l, pitch_r, yaw_r]
            geom_map: (B, 3, H, W) 几何条件图
        
        Returns:
            conditioning: (B, hidden_dim) 扩散模型条件向量
        """
        # gaze嵌入
        gaze_feat = self.gaze_embed(gaze_4d)  # (B, 128)
        
        # 几何编码
        geom_feat = self.geom_encoder(geom_map).flatten(1)  # (B, 128)
        
        # 融合
        conditioning = self.fusion(
            torch.cat([gaze_feat, geom_feat], dim=-1)
        )
        
        return conditioning


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("GazeDiT 物理眼渲染器 + 条件化模块 测试")
    print("=" * 60)
    
    # 物理渲染器
    renderer = PhysicalEyeRenderer()
    
    # 生成几何条件图
    pitch_l, yaw_l = np.radians(-5), np.radians(10)
    pitch_r, yaw_r = np.radians(-5), np.radians(10)
    
    geom = renderer.render_geometry_map(pitch_l, yaw_l, pitch_r, yaw_r)
    print(f"\n几何条件图: shape={geom.shape}, 非零比例={np.mean(geom > 0):.4f}")
    
    # 多样化采样
    geoms = renderer.sample_diverse_geometry(
        pitch_l, yaw_l, pitch_r, yaw_r, n_samples=5
    )
    print(f"多样化采样: {len(geoms)}个几何条件图")
    
    # 条件化模块
    conditioning = GazeDiTConditioning()
    
    gaze_4d = torch.tensor([[pitch_l, yaw_l, pitch_r, yaw_r]])
    geom_tensor = torch.from_numpy(geom).permute(2, 0, 1).unsqueeze(0).float()
    
    cond = conditioning(gaze_4d, geom_tensor)
    print(f"\n条件化输出: shape={cond.shape}")
    print(f"参数量: {sum(p.numel() for p in conditioning.parameters()):,}")