#!/usr/bin/env python3
"""Run the locked RAID-only pairwise-CE neologism training recipe."""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class ModelProfile:
    model: str
    ai_token: str
    human_token: str
    pilot_lr: float
    continuation_steps: int
    continuation_lr: float | None = None
    continuation_pairs: int | None = None
    schedule_total_steps: int | None = None


PROFILES = {
    "gemma": ModelProfile(
        model="google/gemma-4-E4B-it",
        ai_token="<ai>",
        human_token="<human>",
        pilot_lr=1e-3,
        continuation_steps=75,
        continuation_lr=1e-4,
        continuation_pairs=5_000,
        schedule_total_steps=625,
    ),
    "llama": ModelProfile(
        model="meta-llama/Llama-3.1-8B-Instruct",
        ai_token="<AIGEN>",
        human_token="<REAL>",
        pilot_lr=1e-4,
        continuation_steps=0,
    ),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-family", choices=PROFILES, required=True)
    parser.add_argument("--dataset-disk", type=Path, required=True)
    parser.add_argument("--ai-init", type=Path, required=True)
    parser.add_argument("--human-init", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--model", help="Override the Hugging Face model ID/path.")
    parser.add_argument("--runner", type=Path)
    parser.add_argument("--beemo-pairs", type=int, default=2_163)
    parser.add_argument("--bootstrap-resamples", type=int, default=1_000)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def run(command: list[str], *, repo_root: Path, dry_run: bool) -> None:
    printable = " ".join(subprocess.list2cmdline([arg]) for arg in command)
    print(f"+ {printable}", flush=True)
    if dry_run:
        return
    env = os.environ.copy()
    env["PYTHONPATH"] = os.pathsep.join(
        part
        for part in (str(repo_root), env.get("PYTHONPATH", ""))
        if part
    )
    subprocess.run(command, cwd=repo_root, env=env, check=True)


def common_command(
    *,
    python: str,
    runner: Path,
    args: argparse.Namespace,
    profile: ModelProfile,
    output_dir: Path,
    ai_init: Path,
    human_init: Path,
) -> list[str]:
    return [
        python,
        "-u",
        str(runner),
        "--dataset-disk",
        str(args.dataset_disk),
        "--train-split",
        "standard_train_expanded6",
        "--test-split",
        "standard_test",
        "--output-dir",
        str(output_dir),
        "--model",
        args.model or profile.model,
        "--ai-token",
        profile.ai_token,
        "--human-token",
        profile.human_token,
        "--prompt-template",
        "Write {token} text.",
        "--ai-init",
        str(ai_init),
        "--human-init",
        str(human_init),
        "--objective",
        "ai_pairwise",
        "--max-length",
        "512",
        "--batch-size",
        "8",
        "--eval-batch-size",
        "8",
        "--min-lr",
        "1e-5",
        "--warmup-steps",
        "20",
        "--beta",
        "1",
        "--bootstrap-resamples",
        str(args.bootstrap_resamples),
        "--seed",
        str(args.seed),
    ]


def main() -> None:
    args = parse_args()
    profile = PROFILES[args.model_family]
    repo_root = Path(__file__).resolve().parents[1]
    runner = args.runner or repo_root / "scripts/run_gemma_expanded_pairwise_ce.py"
    output_dir = args.output_dir.resolve()
    train_dir = output_dir / "train"
    eval_dir = output_dir / "eval"

    if output_dir.exists() and any(output_dir.iterdir()):
        raise FileExistsError(f"Refusing to overwrite non-empty {output_dir}")
    train_dir.mkdir(parents=True, exist_ok=True)

    train = common_command(
        python=sys.executable,
        runner=runner,
        args=args,
        profile=profile,
        output_dir=train_dir,
        ai_init=args.ai_init,
        human_init=args.human_init,
    )
    train.extend(
        [
            "--pilot-pairs",
            "500",
            "--pilot-epochs",
            "2",
            "--pilot-lrs",
            str(profile.pilot_lr),
            "--pilot-beemo-pairs",
            "250",
        ]
    )

    if args.model_family == "gemma":
        assert profile.continuation_lr is not None
        assert profile.continuation_pairs is not None
        assert profile.schedule_total_steps is not None
        train.extend(
            [
                "--full-pairs",
                str(profile.continuation_pairs),
                "--full-epochs",
                "1",
                "--full-lr",
                str(profile.continuation_lr),
                "--stop-after-steps",
                str(profile.continuation_steps),
                "--schedule-total-steps",
                str(profile.schedule_total_steps),
                "--eval-every",
                str(profile.continuation_steps),
                "--beemo-pairs",
                str(args.beemo_pairs),
            ]
        )
        run(train, repo_root=repo_root, dry_run=args.dry_run)
        return

    train.append("--pilot-only")
    run(train, repo_root=repo_root, dry_run=args.dry_run)

    # Llama continuation consistently regressed the fixed monitor, so evaluate
    # the 125-update pilot directly on both authoritative test populations.
    selected_ai = train_dir / "pilot/lr_0.0001/final/tokens/ai_token.pt"
    eval_dir.mkdir(parents=True, exist_ok=True)
    evaluate = common_command(
        python=sys.executable,
        runner=runner,
        args=args,
        profile=profile,
        output_dir=eval_dir,
        ai_init=selected_ai,
        human_init=args.human_init,
    )
    evaluate.extend(
        [
            "--eval-only",
            "--eval-targets",
            "beemo,raid",
            "--beemo-pairs",
            str(args.beemo_pairs),
        ]
    )
    run(evaluate, repo_root=repo_root, dry_run=args.dry_run)


if __name__ == "__main__":
    main()