#!/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.scaling import (  # noqa: E402
    ScalingExperiment,
    parse_k_values,
    run_scaling_experiment,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Run scaling-law experiments over intervention multiplicity K."
    )
    parser.add_argument("--backend", choices=["toy"], default="toy")
    parser.add_argument("--tasks", default="builtins", help="'builtins' or TaskSpec JSON/JSONL path.")
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--total-records", type=int, default=4096)
    parser.add_argument("--k-values", default="1,2,4,8,16,32")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--shard-size", type=int, default=1000)
    parser.add_argument("--batch-groups", type=int, default=8)
    parser.add_argument("--records-per-group", type=int, default=8)
    parser.add_argument("--hidden-dim", type=int, default=256)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--device", default="auto")
    parser.add_argument(
        "--eval-num-tasks",
        type=int,
        default=20,
        help="Number of toy CausalStress groups per K.",
    )
    parser.add_argument(
        "--eval-k",
        type=int,
        default=None,
        help="Override CausalStress K. Defaults to the current scaling K.",
    )
    args = parser.parse_args(argv)

    config = ScalingExperiment(
        backend=args.backend,
        tasks=args.tasks,
        output_dir=args.out,
        total_records=args.total_records,
        k_values=parse_k_values(args.k_values),
        epochs=args.epochs,
        seed=args.seed,
        shard_size=args.shard_size,
        batch_groups=args.batch_groups,
        records_per_group=args.records_per_group,
        hidden_dim=args.hidden_dim,
        learning_rate=args.lr,
        device=args.device,
        eval_num_tasks=args.eval_num_tasks,
        eval_k=args.eval_k,
    )
    print("planned runs:")
    for run in config.planned_runs():
        print(run)
    summary = run_scaling_experiment(config)
    print(f"wrote aggregate CSV to {summary['aggregate_csv']}")
    print(f"wrote plots to {args.out}")
    for metric, values in summary["regression"].items():
        print(f"{metric}: beta_log_k={values['beta_log_k']:.6g}")
    return 0


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