from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from model.neural_convection_fno import build_model_from_config
from scripts.train import load_pair, normalize_x, normalize_y, ssim_2d, tensor_to_numpy


class InferenceDataset(Dataset):
    def __init__(self, x: torch.Tensor, y: torch.Tensor, stats: dict[str, torch.Tensor]) -> None:
        self.x = x
        self.y = y
        self.stats = stats

    def __len__(self) -> int:
        return int(self.x.shape[0])

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return normalize_x(self.x[idx], self.stats), normalize_y(self.y[idx], self.stats)


def load_config(path: str | Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def resolve(path: str | Path) -> Path:
    path = Path(path)
    return path if path.is_absolute() else PROJECT_ROOT / path


def select_device(name: str) -> torch.device:
    if name == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(name)


def save_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)


@torch.no_grad()
def main() -> None:
    parser = argparse.ArgumentParser(description="Run compact FNO inference on a NeuralConvection split.")
    parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml"))
    parser.add_argument("--split", default=None)
    parser.add_argument("--checkpoint", default=None)
    parser.add_argument("--device", default=None)
    args = parser.parse_args()

    config = load_config(args.config)
    data_cfg = config["data"]
    infer_cfg = config["inference"]
    split = args.split or infer_cfg["split"]
    device = select_device(args.device or config["experiment"].get("device", "auto"))
    root = resolve(data_cfg["root_dir"])
    stats = torch.load(resolve(data_cfg["stats_path"]), map_location="cpu")
    x, y = load_pair(root, split, data_cfg["x_file"], data_cfg["y_file"])
    loader = DataLoader(
        InferenceDataset(x, y, stats),
        batch_size=int(infer_cfg["batch_size"]),
        shuffle=False,
        num_workers=int(data_cfg.get("num_workers", 0)),
        pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda",
    )

    model = build_model_from_config(config).to(device)
    ckpt_path = resolve(args.checkpoint or infer_cfg["checkpoint_path"])
    checkpoint = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    preds: list[np.ndarray] = []
    targets: list[np.ndarray] = []
    inputs: list[np.ndarray] = []
    start = time.perf_counter()
    for x_batch, y_batch in loader:
        x_batch = x_batch.to(device, non_blocking=True)
        pred = tensor_to_numpy(model(x_batch))
        preds.append(pred)
        targets.append(tensor_to_numpy(y_batch))
        inputs.append(tensor_to_numpy(x_batch))
    elapsed = time.perf_counter() - start
    pred_all = np.concatenate(preds, axis=0)
    target_all = np.concatenate(targets, axis=0)
    input_all = np.concatenate(inputs, axis=0)
    channel_names = list(data_cfg["channel_names"])
    metrics = {
        "split": split,
        "checkpoint_path": str(ckpt_path),
        "num_samples": int(pred_all.shape[0]),
        "time_seconds_per_sample": elapsed / max(1, pred_all.shape[0]),
        "per_channel": {},
        "paper_table5_compact_fno_2d": config["paper_table5_compact_fno_2d"],
    }
    for idx, name in enumerate(channel_names):
        mse = float(np.mean((pred_all[:, idx] - target_all[:, idx]) ** 2))
        ssim = float(np.mean([ssim_2d(pred_all[i, idx], target_all[i, idx]) for i in range(pred_all.shape[0])]))
        metrics["per_channel"][name] = {"mse": mse, "ssim": ssim}
        paper = config["paper_table5_compact_fno_2d"]["channels"].get(name, {})
        print(
            f"{name}: mse={mse:.8e} ssim={ssim:.6f} paper_mse={paper.get('mse')} paper_ssim={paper.get('ssim')}",
            flush=True,
        )

    out_path = resolve(infer_cfg["output_path"])
    out_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(out_path, pred_norm=pred_all, target_norm=target_all, input_norm=input_all, channel_names=np.array(channel_names))
    save_json(resolve(infer_cfg["metrics_path"]), metrics)
    print(f"saved_predictions={out_path}", flush=True)
    print(f"saved_metrics={resolve(infer_cfg['metrics_path'])}", flush=True)


if __name__ == "__main__":
    main()