import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# ==========================================
# 1. 傅里叶特征与注意力机制组件 (与 Baseline++ 保持一致)
# ==========================================
class HybridEmbedding(nn.Module):
    def __init__(self, in_dim=2, mapping_size=64, high_scale=12.0):
        super().__init__()
        half_dim = mapping_size // 2
        self.raw_proj = nn.Linear(in_dim, half_dim)
        self.register_buffer('B_high', torch.randn(in_dim, half_dim) * high_scale)

    def forward(self, x):
        feat_smooth = F.silu(self.raw_proj(x)) 
        proj_high = 2 * np.pi * x @ self.B_high
        feat_detail = torch.cat([torch.sin(proj_high), torch.cos(proj_high)], dim=-1)
        return feat_smooth, feat_detail

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True)
        )
    def forward(self, x): return self.double_conv(x)

class GraphFeatureEncoder(nn.Module):
    def __init__(self, in_features, hidden_dim):
        super().__init__()
        self.mlp1 = nn.Sequential(
            nn.Linear(in_features, hidden_dim), nn.LayerNorm(hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim)
        )
        self.attn = nn.MultiheadAttention(embed_dim=hidden_dim, num_heads=4, batch_first=True)
        self.norm = nn.LayerNorm(hidden_dim)
    def forward(self, x):
        h = self.mlp1(x) 
        h_attn, _ = self.attn(h, h, h)
        h = self.norm(h + h_attn)
        return h

class SpatiallyModulatedHybridAttention(nn.Module):
    def __init__(self, coord_dim=2, feature_dim=64):
        super().__init__()
        self.sensor_embedding = HybridEmbedding(in_dim=coord_dim, mapping_size=64, high_scale=12.0)
        self.key_proj_smooth = nn.Linear(32, feature_dim)
        self.key_proj_detail = nn.Linear(64, feature_dim)
        
        self.query_mlp = nn.Sequential(
            nn.Linear(coord_dim, 64), nn.SiLU(), nn.Linear(64, feature_dim * 2)
        )
        self.scale = feature_dim ** -0.5
        
        self.gate_scale = nn.Parameter(torch.tensor([15.0]))
        self.gate_threshold = nn.Parameter(torch.tensor([-0.10])) 

    def forward(self, grid_coords, sensor_coords, sensor_feats, base_flow):
        B, H, W, _ = grid_coords.shape
        flat_grid = grid_coords.view(B, -1, 2)
        
        Q_total = self.query_mlp(flat_grid)
        Q_smooth = Q_total[..., :64]
        Q_detail = Q_total[..., 64:]
        
        u_mean = base_flow[:, 0:1, :, :]
        u_flat = u_mean.view(B, -1, 1)
        
        beta = torch.sigmoid(self.gate_scale * (self.gate_threshold - u_flat))
        
        Q_detail = Q_detail * beta
        
        k_raw_smooth, k_raw_detail = self.sensor_embedding(sensor_coords)
        K_smooth = self.key_proj_smooth(k_raw_smooth)
        K_detail = self.key_proj_detail(k_raw_detail)
        
        score_smooth = torch.matmul(Q_smooth, K_smooth.transpose(-2, -1))
        score_detail = torch.matmul(Q_detail, K_detail.transpose(-2, -1))
        
        attn_logits = (score_smooth + score_detail) * self.scale
        attn = F.softmax(attn_logits, dim=-1)
        
        out = torch.matmul(attn, sensor_feats)
        return out.view(B, H, W, -1).permute(0, 3, 1, 2)

# ==========================================
# 2. 主模型：PIGU-Hybrid (V5 核心修改区)
# ==========================================
class PIGU_Hybrid(nn.Module):
    def __init__(self, sensor_in_dim=3, sensor_count=65, hidden_dim=64, out_dim=4):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.sensor_encoder = GraphFeatureEncoder(sensor_in_dim + 2, hidden_dim)
        self.projector = SpatiallyModulatedHybridAttention(coord_dim=2, feature_dim=hidden_dim)
        
        # 🛑 V5 修改：移除了 base_flow 的 3 个通道
        # 输入通道数从 (hidden_dim + 2 + hidden_dim + 3) 变为 (hidden_dim + 2 + hidden_dim)
        self.in_conv = DoubleConv(hidden_dim + 2 + hidden_dim, 64)
        
        self.down1 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(64, 128))
        self.down2 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(128, 256))
        self.up1 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv_up1 = DoubleConv(256 + 128, 128)
        self.up2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv_up2 = DoubleConv(128 + 64, 64)
        self.out_conv = nn.Conv2d(64, out_dim, kernel_size=1)

    def forward(self, sensor_vals, sensor_coords, grid_coords, base_flow):
        B = sensor_vals.shape[0]
        H, W = grid_coords.shape[1], grid_coords.shape[2]
        node_in = torch.cat([sensor_vals, sensor_coords], dim=-1)
        node_feats = self.sensor_encoder(node_in) 
        global_context = torch.mean(node_feats, dim=1) 
        global_map = global_context.view(B, self.hidden_dim, 1, 1).expand(-1, -1, H, W)
        
        grid_feats = self.projector(grid_coords, sensor_coords, node_feats, base_flow) 
        coords_map = grid_coords.permute(0, 3, 1, 2) 
        
        # 🛑 V5 修改：输入特征拼接处彻底丢弃 base_flow
        x = torch.cat([grid_feats, coords_map, global_map], dim=1) 
        
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x = self.up1(x3)
        x = torch.cat([x, x2], dim=1)
        x = self.conv_up1(x)
        x = self.up2(x)
        x = torch.cat([x, x1], dim=1)
        x = self.conv_up2(x)
        
        # fluc_out 现在直接预测全量的 u, v, p，不再是脉动量！
        full_out = self.out_conv(x)
        
        # 🛑 V5 修改：去掉了均值流的叠加
        uvp_out = full_out[:, :3, :, :]
        
        # 使用 Softplus 防止原始粘性为负，保持物理稳定性
        nu_t_out = F.softplus(full_out[:, 3:, :, :])
        
        return torch.cat([uvp_out, nu_t_out], dim=1)