"""Evaluate an existing stage-two checkpoint without any optimization."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from easycom_head import OFFICIAL_TEST_SESSIONS
from egospeak_omni_finetune import load_stage_two_checkpoint
from egospeak_omni_residual import align_source_predictions, sha256_file
from train_egospeak_omni_finetune import raw_endpoint_metrics
from train_egospeak_omni_residual import (
    evaluate_predictions,
    load_sessions,
    normalize_sessions,
    omni_contract_signature,
    predict_session,
    print_metrics,
    resolve_device,
)


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--sessions", nargs="+", default=None)
    parser.add_argument("--egospeak-feature-root", type=Path)
    parser.add_argument("--omni-feature-root", type=Path)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--eval-chunk-size", type=int, default=0)
    parser.add_argument("--no-progress", action="store_true")
    args = parser.parse_args(argv)
    if args.eval_chunk_size < 0:
        raise ValueError("Evaluation chunk size must be nonnegative")
    model, payload = load_stage_two_checkpoint(args.checkpoint)
    config = payload["training_config"]
    args.egospeak_feature_root = args.egospeak_feature_root or Path(
        config["egospeak_feature_root"]
    )
    args.omni_feature_root = args.omni_feature_root or Path(config["omni_feature_root"])
    args.omni_layer_fusion = config["omni_layer_fusion"]
    args.expected_omni_mode = config["expected_omni_mode"]
    names = normalize_sessions(args.sessions, OFFICIAL_TEST_SESSIONS)
    sessions = load_sessions(args, names)
    if {int(s.omni.shape[-1]) for s in sessions} != {model.omni_dim}:
        raise ValueError("Inference Omni dimension differs from checkpoint")
    known = payload.get("omni_sources", {})
    for session in sessions:
        signature = omni_contract_signature(session)
        for source in known.values():
            c = source["config"]
            expected = (
                source["model_id"],
                source["mode"],
                c.get("memory_seconds"),
                c.get("video_fps"),
                c.get("readout_last_n_layers"),
                c.get("readout_layer_fusion"),
                tuple(source["state_shape_before_layer_fusion"][1:]),
            )
            if signature != expected:
                raise ValueError(
                    "Inference Omni extraction differs from checkpoint provenance"
                )
    output = args.output_dir.expanduser().resolve()
    if output.exists() and (not output.is_dir() or any(output.iterdir())):
        raise FileExistsError(f"Refusing to overwrite output: {output}")
    for name in ("pred", "valid", "raw_eval"):
        (output / name).mkdir(parents=True)
    device = resolve_device(args.device)
    model.requires_grad_(False)
    model.to(device).eval()
    predictions, valid, targets = {}, {}, {}
    for session in sessions:
        raw, _ = predict_session(
            model,
            session,
            device,
            chunk_size=args.eval_chunk_size,
            amp=False,
            show_progress=not args.no_progress,
        )
        aligned, mask = align_source_predictions(raw, session.num_steps)
        predictions[session.name], valid[session.name] = aligned, mask
        targets[session.name] = np.asarray(session.targets, dtype=np.float32)
        np.save(output / "raw_eval" / f"{session.name}.npy", raw)
        np.save(output / "pred" / f"{session.name}.npy", aligned)
        np.save(output / "valid" / f"{session.name}.npy", mask)
    dense = evaluate_predictions(predictions, valid, targets)
    endpoints = raw_endpoint_metrics(predictions, valid, targets)
    metrics = {
        **dense,
        "raw_endpoint_anticipation": endpoints,
        "sessions": names,
        "train_sessions": payload["train_sessions"],
        "seed": config["seed"],
        "checkpoint": str(args.checkpoint.resolve()),
        "checkpoint_sha256": sha256_file(args.checkpoint),
        "inference_only": True,
        "budgets": payload["budgets"],
        "omni_sources": {s.name: s.omni_metadata for s in sessions},
    }
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    (output / "README.md").write_text(
        "# Stage-two inference-only evaluation\n\n"
        f"Checkpoint: {args.checkpoint}. Sessions: {names}. No optimization.\n\n"
        "metrics.json stores dense/raw endpoint AP as fractions; console uses percent. "
        "Repeated official-test evaluation is exploratory, not dev selection.\n"
    )
    print_metrics("stage_two inference only", dense)
    print(
        "Raw endpoint mean AP (percent):",
        {
            family: None
            if endpoints[f"mean_{family}_ap"] is None
            else endpoints[f"mean_{family}_ap"] * 100
            for family in ("onset", "offset")
        },
    )
    print(f"Saved {output / 'metrics.json'}; checkpoint unchanged")


if __name__ == "__main__":
    main()