#!/usr/bin/env python
from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

from dovla_cil.eval.causalstress import (  # noqa: E402
    CausalStressBenchmark,
    CausalStressConfig,
    write_metrics_json,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Evaluate a DoVLA checkpoint on CausalStress.")
    parser.add_argument("--checkpoint", type=Path, default=None)
    parser.add_argument("--backend", choices=["toy"], default="toy")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--num-tasks", type=int, default=20)
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--obs-dim", type=int, default=32)
    parser.add_argument("--action-dim", type=int, default=8)
    parser.add_argument("--action-horizon", type=int, default=4)
    parser.add_argument("--effect-dim", type=int, default=32)
    args = parser.parse_args(argv)

    config = CausalStressConfig(
        backend=args.backend,
        num_tasks=args.num_tasks,
        k=args.k,
        seed=args.seed,
        obs_dim=args.obs_dim,
        action_dim=args.action_dim,
        action_horizon=args.action_horizon,
        effect_dim=args.effect_dim,
    )
    metrics = CausalStressBenchmark(config).evaluate(args.checkpoint, device=args.device)
    metrics["config"] = {
        "backend": args.backend,
        "checkpoint": str(args.checkpoint) if args.checkpoint else None,
        "num_tasks": args.num_tasks,
        "k": args.k,
        "seed": args.seed,
    }
    write_metrics_json(metrics, args.out)
    print(f"wrote CausalStress metrics to {args.out}")
    print(
        "task_success_rate={task_success_rate:.4f} "
        "pairwise_ranking_accuracy={pairwise_ranking_accuracy:.4f} "
        "top1_action_selection={top1_action_selection:.4f}".format(**metrics)
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())