import torch
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.spatial import cKDTree
import os

from dataset import CFDReconstructionDataset
from architectures import PIGU_Hybrid
from model import NavierStokesURANS 

# ==========================================
# 1. 动画与模型配置区 (专为 Variant 5 设置)
# ==========================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SENSOR_COUNT = 65
DT = 0.05
FRAMES_TO_RENDER = 60  

# 指向 Variant 5 的权重路径和输出文件
MODEL_WEIGHTS_PATH = "results_ablation_v5/model_ep400.pth"  
OUTPUT_FILENAME = "eval_results/var5_animation_with_metrics.gif" 

PATH_UNSTEADY = "../../Ablation Experiment/dataset_sin/flow_one_period.npy"
PATH_MEAN = "../../Ablation Experiment/dataset_sin/mean_flow_steady.npy"

VAR_IDX = 0 
VAR_NAME = ['u-velocity', 'v-velocity', 'Pressure'][VAR_IDX]

def main():
    os.makedirs(os.path.dirname(OUTPUT_FILENAME), exist_ok=True)
    
    print(f"[*] Initializing CFD Animation & Evaluation for Variant 5...")
    dataset = CFDReconstructionDataset(PATH_UNSTEADY, PATH_MEAN, SENSOR_COUNT, dt=DT)
    
    stats_max = dataset.stats['max'].to(DEVICE)
    stats_min = dataset.stats['min'].to(DEVICE)
    
    dx_val = dataset.box_len[0] / (dataset.W - 1)
    dy_val = dataset.box_len[1] / (dataset.H - 1)
    dx_tensor = torch.tensor(dx_val, device=DEVICE)
    dy_tensor = torch.tensor(dy_val, device=DEVICE)
    
    def denormalize_batch(norm_tensor):
        return (norm_tensor + 1) / 2 * (stats_max - stats_min) + stats_min

    print(f"[*] Loading Model Weights from: {MODEL_WEIGHTS_PATH}")
    model = PIGU_Hybrid(sensor_in_dim=3, sensor_count=SENSOR_COUNT).to(DEVICE)
    
    if not os.path.exists(MODEL_WEIGHTS_PATH):
        print(f"[!] Warning: Weights {MODEL_WEIGHTS_PATH} not found. Please train V5 first.")
        return
        
    model.load_state_dict(torch.load(MODEL_WEIGHTS_PATH, map_location=DEVICE))
    model.eval()
    
    pde_engine = NavierStokesURANS(dataset.stats).to(DEVICE)

    # 提取传感器的物理坐标用于绘图
    sensor_phys_x = [dataset.grid_X[p[0], p[1]] for p in dataset.sensor_indices]
    sensor_phys_y = [dataset.grid_Y[p[0], p[1]] for p in dataset.sensor_indices]

    print("[*] Building Geometry Mask...")
    raw_mean = np.load(PATH_MEAN).astype(np.float32)
    coords_mean = raw_mean[:, :2]
    tree = cKDTree(coords_mean)
    
    grid_pts = np.stack([dataset.grid_X.ravel(), dataset.grid_Y.ravel()], axis=-1)
    dist, _ = tree.query(grid_pts)
    threshold = max(dx_val, dy_val) * 2.5 
    is_solid_wall = (dist > threshold).reshape(dataset.H, dataset.W)
    is_solid_wall_float = is_solid_wall.astype(float)
    
    valid_mask_torch = ~torch.from_numpy(is_solid_wall).to(DEVICE)

    print(f"[*] Running Inference and Calculating Metrics for {FRAMES_TO_RENDER} frames...")
    true_frames, pred_frames, err_frames = [], [], []
    
    total_l2 = 0.0
    total_rmse = 0.0
    total_resc = 0.0
    num_frames = min(FRAMES_TO_RENDER, len(dataset))

    with torch.no_grad():
        for idx in range(num_frames):
            s_val_t, s_pos, grid_pos_norm, s_val_next, mean_flow = dataset[idx]
            true_norm_t = dataset.data[idx] 
            
            s_val_t = s_val_t.unsqueeze(0).to(DEVICE)
            s_pos = s_pos.unsqueeze(0).to(DEVICE)
            grid_pos_norm = grid_pos_norm.unsqueeze(0).to(DEVICE)
            s_val_next = s_val_next.unsqueeze(0).to(DEVICE)
            mean_flow = mean_flow.unsqueeze(0).to(DEVICE)
            true_norm_t = true_norm_t.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)
            
            pred_phys_t = denormalize_batch(pred_norm_t)
            true_phys_t = denormalize_batch(true_norm_t)
            pred_phys_next = denormalize_batch(pred_norm_next)
            
            p_var_torch = pred_phys_t[0, VAR_IDX]
            t_var_torch = true_phys_t[0, VAR_IDX]
            
            diff_valid = (p_var_torch - t_var_torch)[valid_mask_torch]
            t_valid = t_var_torch[valid_mask_torch]
            
            frame_l2 = torch.norm(diff_valid) / (torch.norm(t_valid) + 1e-8)
            frame_rmse = torch.sqrt(torch.mean(diff_valid**2))
            
            _, _, res_c = pde_engine(pred_phys_t, pred_phys_next, None, dx=dx_tensor, dy=dy_tensor)
            res_c_valid = res_c[0, 0][valid_mask_torch]
            frame_resc = torch.sqrt(torch.mean(res_c_valid**2))
            
            total_l2 += frame_l2.item()
            total_rmse += frame_rmse.item()
            total_resc += frame_resc.item()
            
            p_var = p_var_torch.cpu().numpy()
            t_var = t_var_torch.cpu().numpy()
            
            p_var[is_solid_wall] = np.nan
            t_var[is_solid_wall] = np.nan
            err_var = np.abs(p_var - t_var)
            
            true_frames.append(t_var)
            pred_frames.append(p_var)
            err_frames.append(err_var)

    avg_l2 = total_l2 / num_frames
    avg_rmse = total_rmse / num_frames
    avg_resc = total_resc / num_frames
    
    print(f"\n[Metrics] L2 Rel Error: {avg_l2:.4f} | RMSE: {avg_rmse:.4f} | Cont. Res: {avg_resc:.2e}")

    # ==========================================
    # 3. 顶级学术期刊 CFD 风格动画渲染
    # ==========================================
    print(f"[*] Generating Professional CFD-Style Animation...")
    
    vmin = np.nanmin(true_frames)
    vmax = np.nanmax(true_frames)
    err_vmax = np.nanmax(err_frames) * 0.8 

    fig, axes = plt.subplots(1, 3, figsize=(18, 5.5), facecolor='white')
    fig.subplots_adjust(bottom=0.2) 
    fig.suptitle(f'Variant 5 (w/o Base Flow) Reconstruction - {VAR_NAME}', fontsize=18, fontweight='bold', y=0.98)
    
    extent = [dataset.grid_X.min(), dataset.grid_X.max(), dataset.grid_Y.min(), dataset.grid_Y.max()]

    for ax in axes:
        ax.set_facecolor('#c0c0c0') 
        ax.set_aspect('equal')
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
    
    im_true = axes[0].imshow(true_frames[0], cmap='turbo', origin='lower', extent=extent, vmin=vmin, vmax=vmax, interpolation='bicubic')
    axes[0].contour(is_solid_wall_float, levels=[0.5], extent=extent, colors='black', linewidths=1.5, origin='lower')
    axes[0].set_title("CFD Ground Truth", fontsize=14, pad=10)
    
    im_pred = axes[1].imshow(pred_frames[0], cmap='turbo', origin='lower', extent=extent, vmin=vmin, vmax=vmax, interpolation='bicubic')
    axes[1].contour(is_solid_wall_float, levels=[0.5], extent=extent, colors='black', linewidths=1.5, origin='lower')
    axes[1].set_title("Direct Full-Flow Prediction", fontsize=14, pad=10)
    axes[1].scatter(sensor_phys_x, sensor_phys_y, c='white', edgecolors='black', s=18, alpha=0.9, zorder=5)
    
    divider = make_axes_locatable(axes[1])
    cax = divider.append_axes("right", size="5%", pad=0.1)
    fig.colorbar(im_pred, cax=cax, label=f'{VAR_NAME} Magnitude')
    
    im_err = axes[2].imshow(err_frames[0], cmap='inferno', origin='lower', extent=extent, vmin=0, vmax=err_vmax, interpolation='bicubic')
    axes[2].contour(is_solid_wall_float, levels=[0.5], extent=extent, colors='white', linewidths=1.5, origin='lower')
    axes[2].set_title("Absolute Error", fontsize=14, pad=10)
    
    divider_err = make_axes_locatable(axes[2])
    cax_err = divider_err.append_axes("right", size="5%", pad=0.1)
    fig.colorbar(im_err, cax=cax_err, label='Error Magnitude')

    metrics_text = (
        f"Quantitative Evaluation ({FRAMES_TO_RENDER} frames average)\n"
        f"$L_2$ Relative Error: {avg_l2:.4f}  |  RMSE: {avg_rmse:.4f}  |  Continuity Residual: {avg_resc:.2e}"
    )
    fig.text(0.5, 0.05, metrics_text, ha='center', va='center', fontsize=14, fontweight='bold',
             bbox=dict(facecolor='#f8f9fa', alpha=0.9, edgecolor='#adb5bd', boxstyle='round,pad=0.5'))

    def update(frame_idx):
        im_true.set_array(true_frames[frame_idx])
        im_pred.set_array(pred_frames[frame_idx])
        im_err.set_array(err_frames[frame_idx])
        return [im_true, im_pred, im_err]

    ani = animation.FuncAnimation(fig, update, frames=len(true_frames), interval=100, blit=True)

    if OUTPUT_FILENAME.endswith('.mp4'):
        ani.save(OUTPUT_FILENAME, writer='ffmpeg', fps=10, dpi=200)
    else:
        ani.save(OUTPUT_FILENAME, writer='pillow', fps=10, dpi=150)
        
    print(f"[+] Animation successfully saved to: {OUTPUT_FILENAME}")

if __name__ == "__main__":
    main()