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

class FocusGate(nn.Module):
    """
    FocusGate: 可靠性门控的先验融合模块
    
    论文: arXiv:2610.08663 (ICASSP 2027投稿)
    
    功能:
    1. 从散焦图计算3个形状统计量
    2. 基于阈值判断帧是否可靠
    3. 可靠时融合散焦先验，不可靠时退回基线
    """
    
    def __init__(
        self,
        spread_thresh: float = 0.55,
        sep_thresh: float = 0.45,
        sharp_min: float = 0.03,
        sharp_max: float = 0.70,
        defocus_weight: float = 0.25
    ):
        super().__init__()
        self.spread_thresh = spread_thresh
        self.sep_thresh = sep_thresh
        self.sharp_min = sharp_min
        self.sharp_max = sharp_max
        self.defocus_weight = defocus_weight
    
    def compute_shape_stats(self, defocus_map: torch.Tensor) -> dict:
        """
        计算散焦图的3个形状统计量
        
        Args:
            defocus_map: (B, H, W) in-focus概率, 值域[0,1]
        
        Returns:
            stats字典
        """
        B = defocus_map.shape[0]
        flat = defocus_map.reshape(B, -1)
        
        # Spread: p90 - p10
        p10 = torch.quantile(flat, 0.10, dim=-1)
        p90 = torch.quantile(flat, 0.90, dim=-1)
        spread = p90 - p10
        
        # Separation: 双峰分离度
        high_mean = (flat * (flat > 0.5).float()).sum(-1) / \
                    ((flat > 0.5).float().sum(-1) + 1e-8)
        low_mean = (flat * (flat <= 0.5).float()).sum(-1) / \
                   ((flat <= 0.5).float().sum(-1) + 1e-8)
        separation = high_mean - low_mean
        
        # Sharp area: 锐利区域占比
        sharp_frac = (flat > 0.5).float().mean(-1)
        
        return {
            'spread': spread,
            'separation': separation,
            'sharp_frac': sharp_frac
        }
    
    def forward(
        self,
        base_pred: torch.Tensor,    # (B, H, W) 基线预测
        defocus_map: torch.Tensor   # (B, H, W) 散焦图
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        门控融合
        
        Args:
            base_pred: 基线预测 (B, H, W)
            defocus_map: 散焦图 (B, H, W)
        
        Returns:
            fused: (B, H, W) 融合后预测
            gate: (B,) 门控状态
        """
        stats = self.compute_shape_stats(defocus_map)
        
        # 门控判定
        gate = (
            (stats['spread'] >= self.spread_thresh) &
            (stats['separation'] >= self.sep_thresh) &
            (stats['sharp_frac'] >= self.sharp_min) &
            (stats['sharp_frac'] <= self.sharp_max)
        ).float()  # (B,)
        
        # rank归一化散焦图
        defocus_rank = self._rank_normalize(defocus_map)
        
        # 门控融合: gate=1时加入散焦先验，gate=0时退回基线
        fused = base_pred + gate.unsqueeze(-1).unsqueeze(-1) * \
                self.defocus_weight * defocus_rank
        
        return fused, gate
    
    def _rank_normalize(self, x: torch.Tensor) -> torch.Tensor:
        """逐帧rank归一化到[0,1]"""
        B = x.shape[0]
        flat = x.reshape(B, -1)
        # 对每帧的像素排序
        sorted_vals, _ = torch.sort(flat, dim=-1)
        # 计算rank（midrank处理ties）
        ranks = torch.zeros_like(flat)
        for b in range(B):
            ranks[b] = torch.searchsorted(sorted_vals[b], flat[b])
        ranks = ranks / (flat.shape[-1] - 1)
        return ranks.reshape(x.shape)


class DefocusEstimator(nn.Module):
    """
    散焦估计器
    
    论文中两个版本:
    1. Scratch estimator: U-Net, 2.76M参数
    2. Recipe estimator: MobileNetV3-L + 深度蒸馏 + 高频通道
    """
    
    def __init__(self, backbone='mobilenetv3'):
        super().__init__()
        
        if backbone == 'mobilenetv3':
            # Recipe estimator (论文中效果更好)
            import torchvision.models as models
            self.encoder = models.mobilenet_v3_large(weights='DEFAULT')
            self.encoder.classifier = nn.Identity()
            
            self.head = nn.Sequential(
                nn.Conv2d(960, 256, 1),
                nn.BatchNorm2d(256),
                nn.ReLU6(inplace=True),
                nn.Conv2d(256, 1, 1),
                nn.Sigmoid()
            )
            
            # 高频输入通道
            self.high_freq_conv = nn.Conv2d(1, 3, 1)
        else:
            # Scratch U-Net
            self.encoder = nn.Sequential(
                nn.Conv2d(4, 64, 3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2),
                nn.Conv2d(64, 128, 3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2),
                nn.Conv2d(128, 256, 3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2),
            )
            self.decoder = nn.Sequential(
                nn.ConvTranspose2d(256, 128, 2, stride=2),
                nn.ReLU(inplace=True),
                nn.ConvTranspose2d(128, 64, 2, stride=2),
                nn.ReLU(inplace=True),
                nn.ConvTranspose2d(64, 1, 2, stride=2),
                nn.Sigmoid()
            )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """预测in-focus概率图"""
        if hasattr(self, 'high_freq_conv'):
            # 高频通道
            gray = x.mean(dim=1, keepdim=True)
            high_freq = gray - torch.nn.functional.avg_pool2d(
                gray, 3, stride=1, padding=1
            )
            x_hf = self.high_freq_conv(high_freq)
            x = torch.cat([x, x_hf], dim=1)
        
        if hasattr(self, 'decoder'):
            feat = self.encoder(x)
            out = self.decoder(feat)
        else:
            feat = self.encoder(x)
            out = self.head(feat)
        
        return out.squeeze(1)


class FullFocusGate(nn.Module):
    """
    完整FocusGate系统
    
    包含:
    1. 多个gaze-free先验（中心高斯、人脸、谱残差等）
    2. 散焦估计器
    3. 可靠性门控
    """
    
    def __init__(self):
        super().__init__()
        self.defocus_estimator = DefocusEstimator(backbone='mobilenetv3')
        self.gate = FocusGate()
        
        # 先验权重 (论文中by coordinate ascent fitted)
        self.weights = {
            'center': 1.0,
            'face': 2.0,  # gamma=0.5
            'spectral_residual': 0.0,  # swept to zero
            'dino_attn': 1.5,  # gamma=3
            'depth': 0.0,  # swept to zero
            'flow': 1.0,
            'defocus': 0.25,
        }
    
    def forward(self, frame, priors):
        """
        Args:
            frame: (B, 3, H, W) 输入帧
            priors: dict of (B, H, W) 先验图
        """
        # 1. 基线融合
        base = torch.zeros_like(priors['center'])
        for name, (w, gamma) in [
            ('center', (1.0, 1)),
            ('face', (2.0, 0.5)),
            ('dino_attn', (1.5, 3)),
            ('flow', (1.0, 1)),
        ]:
            base += w * priors[name].clamp(0, 1).pow(gamma)
        
        # 2. 散焦估计
        defocus = self.defocus_estimator(frame)
        
        # 3. 门控融合
        fused, gate = self.gate(base, defocus)
        
        return fused, gate


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("FocusGate 可靠性门控融合 - 测试")
    print("=" * 60)
    
    B, H, W = 2, 224, 224
    
    # 模拟散焦图
    # 场景1: 浅景深（双峰，门控开启）
    defocus_shallow = torch.zeros(B, H, W)
    defocus_shallow[:, :H//2, :] = 0.9  # 上半部分锐利
    defocus_shallow[:, H//2:, :] = 0.1  # 下半部分模糊
    
    # 场景2: 深景深（均匀，门控关闭）
    defocus_deep = torch.full((B, H, W), 0.5)
    defocus_deep += torch.randn(B, H, W) * 0.05
    
    base_pred = torch.randn(B, H, W).sigmoid()
    
    gate_module = FocusGate()
    
    # 浅景深
    fused_shallow, gate_shallow = gate_module(base_pred, defocus_shallow)
    print(f"\n浅景深场景:")
    print(f"  门控状态: {gate_shallow.tolist()}")
    print(f"  融合变化: {(fused_shallow - base_pred).abs().mean():.6f}")
    
    # 深景深
    fused_deep, gate_deep = gate_module(base_pred, defocus_deep)
    print(f"\n深景深场景:")
    print(f"  门控状态: {gate_deep.tolist()}")
    print(f"  融合变化: {(fused_deep - base_pred).abs().mean():.6f}")
    
    print(f"\n门控效果验证:")
    print(f"  浅景深门控开启: {gate_shallow.mean():.2f} (应≈1)")
    print(f"  深景深门控关闭: {gate_deep.mean():.2f} (应≈0)")