#!/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.experiments.baselines import (  # noqa: E402
    BaselineConfig,
    list_baselines,
    train_baseline,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run a DoVLA-CIL baseline experiment.")
    parser.add_argument("--baseline", choices=list_baselines(), required=True)
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--backend", choices=["toy"], default="toy")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-groups", type=int, default=4)
    parser.add_argument("--records-per-group", type=int, default=8)
    parser.add_argument("--hidden-dim", type=int, default=128)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--shard-size", type=int, default=1024)
    parser.add_argument("--eval-num-tasks", type=int, default=6)
    parser.add_argument("--eval-k", type=int, default=4)
    args = parser.parse_args(argv)

    summary = train_baseline(
        BaselineConfig(
            baseline=args.baseline,
            dataset=args.dataset,
            out=args.out,
            backend=args.backend,
            epochs=args.epochs,
            batch_groups=args.batch_groups,
            records_per_group=args.records_per_group,
            hidden_dim=args.hidden_dim,
            lr=args.lr,
            device=args.device,
            seed=args.seed,
            shard_size=args.shard_size,
            eval_num_tasks=args.eval_num_tasks,
            eval_k=args.eval_k,
        )
    )
    eval_metrics = summary["eval"]
    print(f"baseline={summary['baseline']}")
    print(f"prepared_dataset={summary['prepared_dataset']}")
    print(f"checkpoint={summary['checkpoint']}")
    print(
        "task_success_rate={task_success_rate:.4f} "
        "pairwise_ranking_accuracy={pairwise_ranking_accuracy:.4f}".format(**eval_metrics)
    )
    return 0


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