from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import torch

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

from model.ContinuousFlowOperator import ContinuousFlowOperator, rmse, trajectory_relative_l2  # noqa: E402
from scripts.common import ensure_dirs, load_burgers_splits, load_config, write_json  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run full-test inference for CFO Burgers")
    parser.add_argument("--config", default=str(PROJECT_ROOT / "config" / "config.yaml"))
    parser.add_argument("--checkpoint", default=None)
    return parser.parse_args()


@torch.no_grad()
def rollout_all(model: ContinuousFlowOperator, data: np.ndarray, device: torch.device, steps_per_segment: int, batch_size: int) -> np.ndarray:
    preds = []
    for start in range(0, len(data), batch_size):
        x0 = torch.from_numpy(data[start : start + batch_size, 0]).to(device=device, dtype=torch.float32)
        preds.append(model.rollout(x0, data.shape[1], steps_per_segment=steps_per_segment).numpy())
    return np.concatenate(preds, axis=0)


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    ensure_dirs(cfg)
    ckpt_path = Path(args.checkpoint or cfg["inference"]["checkpoint"])
    if not ckpt_path.exists():
        raise FileNotFoundError(f"checkpoint not found: {ckpt_path}")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    splits, tspan = load_burgers_splits(cfg)
    test_data = splits["test"]
    model = ContinuousFlowOperator(input_size=cfg["model"]["input_size"], gamma=cfg["model"]["gamma"]).to(device)
    ckpt = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()
    pred = rollout_all(
        model,
        test_data,
        device,
        steps_per_segment=int(cfg["inference"]["steps_per_segment"]),
        batch_size=int(cfg["inference"]["eval_batch_size"]),
    )
    per_sample = []
    for i in range(len(test_data)):
        per_sample.append(trajectory_relative_l2(test_data[i : i + 1], pred[i : i + 1]))
    metrics = {
        "checkpoint": str(ckpt_path),
        "test_shape": list(test_data.shape),
        "prediction_shape": list(pred.shape),
        "trajectory_relative_l2": trajectory_relative_l2(test_data, pred),
        "rmse": rmse(test_data, pred),
        "per_sample_relative_l2_mean": float(np.mean(per_sample)),
        "per_sample_relative_l2_std": float(np.std(per_sample)),
        "per_sample_relative_l2_max": float(np.max(per_sample)),
        "worst_sample_index": int(np.argmax(per_sample)),
        "solver": cfg["inference"]["solver"],
        "steps_per_segment": int(cfg["inference"]["steps_per_segment"]),
    }
    out_dir = Path(cfg["paths"]["results_dir"])
    np.savez_compressed(out_dir / "test_predictions.npz", prediction=pred, target=test_data, tspan=tspan, per_sample_relative_l2=np.asarray(per_sample))
    write_json(out_dir / "test_metrics.json", metrics)
    print("Full test inference complete.", flush=True)
    print(metrics, flush=True)


if __name__ == "__main__":
    main()