#!/usr/bin/env python3
"""Inference script for MHD Edge Surrogate Model."""
import argparse
import sys
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from mhd_model.dataset import (
    MhdEdgeDataset,
    DataStats,
    apply_transform,
    CHANNEL_NAMES,
)
from mhd_model.model import create_model


def denormalize_to_raw(pred_norm, stats):
    """Convert normalized predictions back to raw values."""
    pred_transformed = stats.denormalize_outputs(pred_norm)
    pred_raw = apply_transform(
        pred_transformed, stats.channel_transforms,
        asinh_scale=stats.asinh_scale, inverse=True,
    )
    return pred_raw


def load_model_and_stats(checkpoint_path, device="cpu"):
    """Load model and stats from checkpoint.

    Returns:
        model, stats, args_dict
    """
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    args_dict = ckpt.get("args", {})
    stats = DataStats.from_dict(ckpt["stats"])

    model = create_model(
        hidden_dim=args_dict.get("hidden_dim", 256),
        num_heads=args_dict.get("num_heads", 8),
        num_latent_tokens=args_dict.get("num_latent_tokens", 32),
        num_decoder_layers=args_dict.get("num_decoder_layers", 1),
    ).to(device)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()

    return model, stats, args_dict


def predict(model, stats, inputs, query_coords, device="cpu"):
    """Run inference for a single sample.

    Args:
        model: Trained CfdDecoderOnly model
        stats: DataStats for normalization
        inputs: dict with keys {Ip, c_frac, d_scale, n_core, p_input} or tensor (5,)
        query_coords: (66, 26, 2) or (N, 2) R,Z coordinates
        device: torch device

    Returns:
        dict with predicted fields:
            te: (66, 26), ti: (66, 26), rad: (66, 26),
            ni: (66, 26, 2), up: (66, 26, 2), v2ce_0: (66, 26)
    """
    from mhd_model.model import MhdCfdDecoderOnly

    # Prepare input params
    if isinstance(inputs, dict):
        input_value = torch.tensor([
            inputs["Ip"], inputs["c_frac"], inputs["d_scale"],
            inputs["n_core"], inputs["p_input"],
        ], dtype=torch.float32)
    else:
        input_value = torch.as_tensor(inputs, dtype=torch.float32)

    input_value = stats.normalize_inputs(input_value).unsqueeze(0).to(device)  # (1, 5)

    # Prepare query coords
    coords = torch.as_tensor(query_coords, dtype=torch.float32)
    if coords.dim() == 3:  # (66, 26, 2)
        spatial_shape = coords.shape[:2]
        coords = coords.reshape(-1, 2)
    else:
        spatial_shape = None

    coords_norm = stats.normalize_coords(coords).unsqueeze(0).to(device)  # (1, N, 2)
    N = coords_norm.shape[1]

    unbatch_idx, unbatch_select = MhdCfdDecoderOnly.create_unbatch_indices_from_lengths(
        [N], device=device,
    )

    with torch.no_grad():
        outputs = model(
            params=input_value,
            query_pos=coords_norm,
            unbatch_idx=unbatch_idx,
            unbatch_select=unbatch_select,
        )
        pred_norm = outputs["pred"]  # (N, 8) or (1, N, 8)
        if pred_norm.dim() == 2:
            pred_norm = pred_norm.unsqueeze(0)
        pred_raw = denormalize_to_raw(pred_norm, stats).squeeze(0).cpu()  # (N, 8)

    # Reshape to spatial grid
    if spatial_shape is not None:
        pred_raw = pred_raw.reshape(*spatial_shape, 8)

        return {
            "te": pred_raw[..., 0].numpy(),        # (66, 26) in eV
            "ti": pred_raw[..., 1].numpy(),        # (66, 26) in eV
            "rad": pred_raw[..., 2].numpy(),       # (66, 26)
            "ni": pred_raw[..., 3:5].numpy(),      # (66, 26, 2)
            "up": pred_raw[..., 5:7].numpy(),      # (66, 26, 2)
            "v2ce_0": pred_raw[..., 7].numpy(),    # (66, 26)
        }

    return {"all_channels": pred_raw.numpy()}  # (N, 8)


def predict_batch(model, stats, inputs_batch, coords_batch, device="cpu"):
    """Run inference for a batch of samples.

    Args:
        inputs_batch: (B, 5) tensor of input parameters
        coords_batch: (B, 66, 26, 2) or (B, N, 2) coordinates

    Returns:
        pred_raw: (B, N, 8) raw predictions
    """
    from mhd_model.model import MhdCfdDecoderOnly

    inputs_batch = torch.as_tensor(inputs_batch, dtype=torch.float32)
    coords_batch = torch.as_tensor(coords_batch, dtype=torch.float32)

    if coords_batch.dim() == 4:  # (B, 66, 26, 2)
        B = coords_batch.shape[0]
        coords_batch = coords_batch.reshape(B, -1, 2)

    input_norm = stats.normalize_inputs(inputs_batch).to(device)
    coords_norm = stats.normalize_coords(coords_batch).to(device)
    B, N, _ = coords_norm.shape

    unbatch_idx, unbatch_select = MhdCfdDecoderOnly.create_unbatch_indices_from_lengths(
        [N] * B, device=device,
    )

    with torch.no_grad():
        outputs = model(
            params=input_norm,
            query_pos=coords_norm,
            unbatch_idx=unbatch_idx,
            unbatch_select=unbatch_select,
        )
        pred_norm = outputs["pred"]
        if pred_norm.dim() == 2:
            pred_norm = pred_norm.view(B, N, -1)
        pred_raw = denormalize_to_raw(pred_norm, stats).cpu()

    return pred_raw.numpy()


def visualize_comparison(pred, true_values, coords, sample_idx=0, save_path=None):
    """Create contour plots comparing predicted vs true fields.

    Args:
        pred: dict from predict() with field arrays
        true_values: (66, 26, 9) raw ground truth
        coords: (66, 26, 2) R,Z coordinates
        sample_idx: sample index for title
        save_path: path to save figure (optional)
    """
    import matplotlib.pyplot as plt

    R = coords[..., 0]
    Z = coords[..., 1]

    scalar_fields = ["te", "ti", "rad"]
    fig, axes = plt.subplots(len(scalar_fields), 3, figsize=(18, 5 * len(scalar_fields)))

    true_map = {
        "te": true_values[..., 0],
        "ti": true_values[..., 1],
        "rad": true_values[..., 2],
    }

    for row, fname in enumerate(scalar_fields):
        pred_field = pred[fname]
        true_field = true_map[fname]
        error = np.abs(pred_field - true_field)

        # True
        ax = axes[row, 0]
        c = ax.contourf(R, Z, true_field, levels=50, cmap="viridis")
        plt.colorbar(c, ax=ax)
        ax.set_title(f"{fname} - True")
        ax.set_xlabel("R [m]")
        ax.set_ylabel("Z [m]")

        # Predicted
        ax = axes[row, 1]
        c = ax.contourf(R, Z, pred_field, levels=50, cmap="viridis")
        plt.colorbar(c, ax=ax)
        ax.set_title(f"{fname} - Predicted")
        ax.set_xlabel("R [m]")

        # Error
        ax = axes[row, 2]
        c = ax.contourf(R, Z, error, levels=50, cmap="magma")
        plt.colorbar(c, ax=ax)
        ax.set_title(f"{fname} - |Error|")
        ax.set_xlabel("R [m]")

    plt.suptitle(f"Sample {sample_idx}", fontsize=16)
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Saved figure to {save_path}")
    else:
        plt.show()
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description="MHD Edge Surrogate Inference")
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--hdf5-path", type=Path,
                        default=Path.home() / "UEDGE_MHD_KSTAR_dataset.hdf5")
    parser.add_argument("--sample-idx", type=int, default=0,
                        help="Sample index from test set to evaluate")
    parser.add_argument("--output-dir", type=Path, default=Path("inference_outputs"))
    parser.add_argument("--device", type=str,
                        default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    device = torch.device(args.device)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    # Load model
    print(f"Loading checkpoint: {args.checkpoint}")
    model, stats, ckpt_args = load_model_and_stats(args.checkpoint, device)

    # Load test dataset
    test_ds = MhdEdgeDataset(args.hdf5_path, split="test", stats=stats)
    print(f"Test set: {len(test_ds)} samples")

    # Get sample
    idx = args.sample_idx
    sample = test_ds[idx]
    real_idx = test_ds.indices[idx]

    # Get raw coords (before normalization)
    coords_raw = torch.from_numpy(test_ds.coords[real_idx])  # (66, 26, 2)
    inputs_raw = torch.from_numpy(test_ds.inputs[real_idx])   # (5,)

    print(f"\nSample {idx} (dataset index {real_idx}):")
    print(f"  Inputs: {dict(zip(['Ip','c_frac','d_scale','n_core','p_input'], inputs_raw.tolist()))}")

    # Predict
    pred = predict(model, stats, inputs_raw, coords_raw, device=str(device))

    # Ground truth
    true_values = test_ds.outputs[real_idx]  # (66, 26, 8)

    # Per-channel errors
    print("\nPer-channel errors:")
    for c, name in enumerate(CHANNEL_NAMES):
        true_c = true_values[..., c]
        if name in pred:
            pred_c = pred[name]
        elif name in ["ni_0", "ni_1"]:
            pred_c = pred["ni"][..., int(name[-1])]
        elif name in ["up_0", "up_1"]:
            pred_c = pred["up"][..., int(name[-1])]
        else:
            pred_c = pred[name]

        mae = np.abs(pred_c - true_c).mean()
        rel = np.abs(pred_c - true_c).mean() / (np.abs(true_c).mean() + 1e-30)
        print(f"  {name:8s}: MAE={mae:.4e}  Rel={rel:.4f}")

    # Visualize
    save_path = args.output_dir / f"comparison_sample_{idx}.png"
    visualize_comparison(pred, true_values, coords_raw.numpy(), idx, save_path)


if __name__ == "__main__":
    main()