from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch

from common import (
    ROOT,
    UnitTransformer,
    build_model,
    load_config,
    load_pipe_tensors,
    relative_l2_per_sample,
    resolve,
    save_json,
    select_device,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Run HO-FNO Pipe inference with the best checkpoint.")
    parser.add_argument("--config", default=str(ROOT / "config" / "config.yaml"))
    parser.add_argument("--checkpoint", default=None)
    args = parser.parse_args()

    cfg = load_config(args.config)
    output_dir = resolve(cfg["paths"]["output_dir"])
    results_dir = resolve(cfg["paths"]["results_dir"], output_dir)
    results_dir.mkdir(parents=True, exist_ok=True)
    checkpoint = resolve(args.checkpoint or cfg["paths"]["checkpoint"], output_dir)
    device = select_device(require_accelerator=False)

    payload = torch.load(checkpoint, map_location=device)
    ckpt_cfg = payload.get("config", cfg)
    ckpt_cfg["paths"]["data_dir"] = cfg["paths"]["data_dir"]
    model = build_model(ckpt_cfg).to(device)
    model.load_state_dict(payload["model_state"])
    model.eval()
    x_norm = UnitTransformer.from_state_dict(payload["x_normalizer"]).to(device)
    y_norm = UnitTransformer.from_state_dict(payload["y_normalizer"]).to(device)

    _, _, x_test, y_test = load_pipe_tensors(ckpt_cfg)
    batch_size = int(ckpt_cfg["training"]["batch_size"])
    predictions = []
    targets = []
    inputs = []
    rel_values = []
    with torch.no_grad():
        for start in range(0, x_test.shape[0], batch_size):
            x = x_test[start : start + batch_size].to(device)
            y = y_test[start : start + batch_size].to(device)
            pred = y_norm.decode(model(x_norm.encode(x)))
            rel = relative_l2_per_sample(pred, y)
            predictions.append(pred.cpu().numpy())
            targets.append(y.cpu().numpy())
            inputs.append(x.cpu().numpy())
            rel_values.append(rel.cpu().numpy())

    pred_arr = np.concatenate(predictions, axis=0)
    target_arr = np.concatenate(targets, axis=0)
    input_arr = np.concatenate(inputs, axis=0)
    rel_arr = np.concatenate(rel_values, axis=0)
    np.save(results_dir / "predictions.npy", pred_arr)
    np.save(results_dir / "targets.npy", target_arr)
    np.save(results_dir / "inputs.npy", input_arr)
    np.save(results_dir / "relative_l2_per_sample.npy", rel_arr)
    metrics = {
        "checkpoint": str(checkpoint),
        "test_samples": int(rel_arr.shape[0]),
        "mean_relative_l2": float(rel_arr.mean()),
        "median_relative_l2": float(np.median(rel_arr)),
        "max_relative_l2": float(rel_arr.max()),
        "min_relative_l2": float(rel_arr.min()),
    }
    save_json(results_dir / "inference_metrics.json", metrics)
    print(
        "inference_complete "
        f"mean_relative_l2={metrics['mean_relative_l2']:.8f} "
        f"predictions={results_dir / 'predictions.npy'}",
        flush=True,
    )


if __name__ == "__main__":
    main()