import sys
import os
import torch
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from scipy.ndimage import binary_dilation
import importlib.util

# =========================================================
# 【核心修复】自动推导当前脚本所在的绝对路径 (Ablation Experiment)
# =========================================================
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

# 确保能找到 dataset 和 PDE 引擎 (在 baseline+ 文件夹中)
sys.path.append(os.path.join(SCRIPT_DIR, "baseline++")) 
try:
    from dataset import CFDReconstructionDataset
    from model import NavierStokesURANS
except ImportError:
    print("[!] 请确保脚本与 dataset.py / model.py 处于同级目录，或修改 sys.path")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SENSOR_COUNT = 65
DT = 0.05
EVAL_FRAMES = 100 

# 精准定位数据集
PATH_UNSTEADY = os.path.join(SCRIPT_DIR, "dataset_sin", "flow_one_period.npy")
PATH_MEAN = os.path.join(SCRIPT_DIR, "dataset_sin", "mean_flow_steady.npy")

# 统一使用绝对路径动态加载字典
VARIANTS = {
    "Baseline++": {
        "arch": os.path.join(SCRIPT_DIR, "baseline++", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "baseline++", "results", "model_ep1900.pth")
    },
    "V1 (Hardcoded Gate)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_1 Hardcoded Gate", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_1 Hardcoded Gate", "results", "v1_ep1900.pth")
    },
    "V2 (Global On)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_2 Spatial Gate", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_2 Spatial Gate", "results", "v2_ep1900.pth")
    },
    "V3 (Global Off)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_3 High_Freq", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_3 High_Freq", "results", "v3_ep1900.pth")
    },
    "V4 (w/o PDE)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_4 PDE_Loss", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_4 PDE_Loss", "results", "v4_ep1900.pth")
    },
    "V5 (w/o Base Flow)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_5 BaseFlow", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_5 BaseFlow", "results", "v5_ep1900.pth")
    },
    # ================= 新增物理底层消融变体 =================
    "V6 (w/o Robust PDE)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_6 Robust PDE", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_6 Robust PDE", "results", "v6_ep1900.pth")
    },
    "V7 (w/o Safe Boundary)": {
        "arch": os.path.join(SCRIPT_DIR, "Variant_7 Safe Boundary", "architectures.py"),
        "weight": os.path.join(SCRIPT_DIR, "Variant_7 Safe Boundary", "results", "v7_ep1900.pth")
    }
}

def load_architecture_class(arch_path):
    spec = importlib.util.spec_from_file_location("dynamic_arch", arch_path)
    arch_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(arch_module)
    return arch_module.PIGU_Hybrid

def compute_vorticity(u, v, dx, dy):
    dv_dy, dv_dx = torch.gradient(v, spacing=(dy, dx), dim=(-2, -1))
    du_dy, du_dx = torch.gradient(u, spacing=(dy, dx), dim=(-2, -1))
    return dv_dx - du_dy

def get_boundary_mask(solid_mask):
    solid_np = solid_mask.cpu().numpy()
    dilated = binary_dilation(solid_np, iterations=2)
    boundary_np = dilated ^ solid_np  
    return torch.from_numpy(boundary_np).to(DEVICE)

def evaluate_variant(name, config, dataset, pde_engine, dx_val, dy_val, valid_mask, boundary_mask):
    print(f"\n[*] Evaluating Quantitative Metrics for: {name}")
    
    if not os.path.exists(config['arch']) or not os.path.exists(config['weight']):
        print(f"  [!] Missing files for {name}. Please check paths. Skipping.")
        return None

    PIGU_Hybrid_Class = load_architecture_class(config['arch'])
    
    # 动态适配：如果是新架构，它自己会接受 out_dim=4 的默认值；如果是老架构则为 3。
    model = PIGU_Hybrid_Class(sensor_in_dim=3, sensor_count=SENSOR_COUNT).to(DEVICE)
    try:
        model.load_state_dict(torch.load(config['weight'], map_location=DEVICE, weights_only=True))
    except Exception as e:
        print(f"  [!] Failed to load weights for {name}: {e}")
        return None
        
    model.eval()

    metrics = {
        "L2_Velocity": [], "Continuity_Res": [], "Vorticity_L2": [], "TKE_L2": [], "Boundary_Pressure": [] 
    }
    
    stats_max = dataset.stats['max'].to(DEVICE)
    stats_min = dataset.stats['min'].to(DEVICE)
    
    # 【修复 1】：兼容 3通道和 4通道的的反归一化
    def denormalize(norm_tensor):
        uvp_norm = norm_tensor[:, :3, :, :]
        uvp_phys = (uvp_norm + 1) / 2 * (stats_max - stats_min) + stats_min
        if norm_tensor.shape[1] >= 4:
            nu_t_raw = norm_tensor[:, 3:4, :, :]
            return torch.cat([uvp_phys, nu_t_raw], dim=1)
        return uvp_phys

    dx_tensor = torch.tensor(dx_val, device=DEVICE)
    dy_tensor = torch.tensor(dy_val, device=DEVICE)

    num_eval = min(EVAL_FRAMES, len(dataset))
    
    with torch.no_grad():
        for idx in range(num_eval):
            s_val_t, s_pos, grid_pos_norm, s_val_next, mean_flow = dataset[idx]
            
            s_val_t, s_pos = s_val_t.unsqueeze(0).to(DEVICE), s_pos.unsqueeze(0).to(DEVICE)
            grid_pos_norm = grid_pos_norm.unsqueeze(0).to(DEVICE)
            s_val_next, mean_flow = s_val_next.unsqueeze(0).to(DEVICE), mean_flow.unsqueeze(0).to(DEVICE)
            true_norm_t = dataset.data[idx].unsqueeze(0).to(DEVICE)

            pred_norm_t = model(s_val_t, s_pos, grid_pos_norm, base_flow=mean_flow)
            pred_norm_next = model(s_val_next, s_pos, grid_pos_norm, base_flow=mean_flow)
            
            p_phys = denormalize(pred_norm_t)
            t_phys = denormalize(true_norm_t)
            mean_phys = denormalize(mean_flow)
            
            u_p, v_p, pres_p = p_phys[0, 0], p_phys[0, 1], p_phys[0, 2]
            u_t, v_t, pres_t = t_phys[0, 0], t_phys[0, 1], t_phys[0, 2]
            u_m, v_m = mean_phys[0, 0], mean_phys[0, 1]

            # 1. L2 Velocity
            uv_p = p_phys[0, 0:2][:, valid_mask]
            uv_t = t_phys[0, 0:2][:, valid_mask]
            metrics["L2_Velocity"].append((torch.norm(uv_p - uv_t) / (torch.norm(uv_t) + 1e-8)).item())

            # 【修复 2】：自动探测并计算 Beta 掩码，提供给 PDE Engine
            beta_mask = None
            if hasattr(model, 'projector') and hasattr(model.projector, 'gate_scale'):
                u_mean = mean_flow[:, 0:1, :, :]
                u_flat = u_mean.view(mean_flow.shape[0], -1, 1)
                beta_flat = torch.sigmoid(model.projector.gate_scale * (model.projector.gate_threshold - u_flat))
                beta_mask = beta_flat.view(mean_flow.shape[0], 1, dataset.H, dataset.W)

            # 2. Continuity
            # 无论是什么架构，我们都使用外部导入的 Baseline++ 的 pde_engine 作为统一的裁判标准
            _, _, res_c = pde_engine(p_phys, denormalize(pred_norm_next), None, dx=dx_tensor, dy=dy_tensor, beta_mask=beta_mask)
            metrics["Continuity_Res"].append(torch.sqrt(torch.mean(res_c[0, 0][valid_mask]**2)).item())

            # 3. Vorticity
            vort_p = compute_vorticity(u_p, v_p, dx_val, dy_val)
            vort_t = compute_vorticity(u_t, v_t, dx_val, dy_val)
            metrics["Vorticity_L2"].append((torch.norm(vort_p[valid_mask] - vort_t[valid_mask]) / (torch.norm(vort_t[valid_mask]) + 1e-8)).item())

            # 4. TKE
            tke_p = 0.5 * ((u_p - u_m)**2 + (v_p - v_m)**2)
            tke_t = 0.5 * ((u_t - u_m)**2 + (v_t - v_m)**2)
            metrics["TKE_L2"].append((torch.norm(tke_p[valid_mask] - tke_t[valid_mask]) / (torch.norm(tke_t[valid_mask]) + 1e-8)).item())

            # 5. Boundary Pressure
            metrics["Boundary_Pressure"].append((torch.norm(pres_p[boundary_mask] - pres_t[boundary_mask]) / (torch.norm(pres_t[boundary_mask]) + 1e-8)).item())

    result_row = {"Variant": name}
    for k, v in metrics.items(): result_row[k] = np.mean(v)
    return result_row

def main():
    print("[*] Initializing Unified Evaluation Script...")
    dataset = CFDReconstructionDataset(PATH_UNSTEADY, PATH_MEAN, SENSOR_COUNT, dt=DT)
    
    # 实例化裁判引擎：Baseline++ 版本的标准物理检验器
    pde_engine = NavierStokesURANS(dataset.stats).to(DEVICE)
    
    dx_val = dataset.box_len[0] / (dataset.W - 1)
    dy_val = dataset.box_len[1] / (dataset.H - 1)

    raw_mean = np.load(PATH_MEAN).astype(np.float32)
    tree = cKDTree(raw_mean[:, :2])
    grid_pts = np.stack([dataset.grid_X.ravel(), dataset.grid_Y.ravel()], axis=-1)
    dist, _ = tree.query(grid_pts)
    is_solid_wall = (dist > max(dx_val, dy_val) * 2.5).reshape(dataset.H, dataset.W)
    
    valid_mask = ~torch.from_numpy(is_solid_wall).to(DEVICE)
    boundary_mask = get_boundary_mask(torch.from_numpy(is_solid_wall))

    results_list = []
    for name, config in VARIANTS.items():
        res = evaluate_variant(name, config, dataset, pde_engine, dx_val, dy_val, valid_mask, boundary_mask)
        if res: results_list.append(res)

    if results_list:
        df = pd.DataFrame(results_list)
        pd.options.display.float_format = '{:.3e}'.format  
        print("\n" + "="*85)
        print("🚀 Final Quantitative Ablation Study Results (Including PDEs) 🚀")
        print("="*85)
        print(df.to_string(index=False, justify='center'))
        print("="*85)
        
        os.makedirs(os.path.join(SCRIPT_DIR, "eval_results"), exist_ok=True)
        csv_path = os.path.join(SCRIPT_DIR, "eval_results", "comprehensive_ablation_metrics.csv")
        df.to_csv(csv_path, index=False)
        print(f"[+] Results saved to {csv_path}")

if __name__ == "__main__":
    main()