from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch

from common import (
    build_model,
    collect_predictions,
    get_device,
    load_config,
    make_loaders,
    mean_relative_l2_percent,
    resolve_output_path,
    save_json,
    set_seed,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(Path(__file__).resolve().parents[1] / "config" / "config.yaml"))
    parser.add_argument("--checkpoint", default=None)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    set_seed(cfg["training"]["seed"], cfg["runtime"]["deterministic"], cfg["runtime"]["benchmark"])
    device = get_device(cfg)
    checkpoint = Path(args.checkpoint) if args.checkpoint else resolve_output_path(cfg, cfg["paths"]["checkpoint"])
    results_dir = resolve_output_path(cfg, cfg["paths"]["results_dir"])
    results_dir.mkdir(parents=True, exist_ok=True)

    _, _, test_loader = make_loaders(cfg)
    model = build_model(cfg).to(device)
    payload = torch.load(checkpoint, map_location=device)
    model.load_state_dict(payload["model_state"] if "model_state" in payload else payload)
    test_l2 = mean_relative_l2_percent(model, test_loader, device)
    inputs, preds, targets = collect_predictions(model, test_loader, device)
    pred_path = results_dir / "test_predictions.npz"
    np.savez_compressed(pred_path, inputs=inputs, predictions=preds, targets=targets)
    save_json(
        results_dir / "inference_metrics.json",
        {
            "status": "complete",
            "checkpoint": str(checkpoint),
            "prediction_file": str(pred_path),
            "test_mean_relative_l2_percent": test_l2,
            "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"],
            "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"],
        },
    )
    print(f"inference_complete test_mean_relative_l2={test_l2:.6f}% predictions={pred_path}", flush=True)


if __name__ == "__main__":
    main()