#!/usr/bin/env python3
"""Run one same-GPU seed of the frozen Rain EP bias confirmation."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import subprocess
import sys
import time


ROOT = Path(__file__).resolve().parents[1]
ENDPOINT = ROOT / "experiments" / "rain_ep_bias_train.py"
RESULT_ROOT = ROOT / "results" / "ep_bias" / "c1"
AUTHOR_REVISION = "6b253fd8a5d267535f58ab79992256ef10031ceb"
SEEDS = (1989, 1990, 1991, 1992, 1993)
CONDITIONS = (
    ("clean", "clean"),
    ("raw", "raw"),
    ("same_rms_noise", "noise"),
    ("constant", "constant"),
    ("innovation", "innovation"),
    ("oracle", "oracle"),
)


def revision(path: Path) -> str:
    return subprocess.check_output(
        ["git", "-C", str(path), "rev-parse", "HEAD"], text=True).strip()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--author-root", type=Path, required=True)
    parser.add_argument("--seed", type=int, choices=SEEDS, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    author_root = args.author_root.resolve()
    if revision(author_root) != AUTHOR_REVISION:
        raise ValueError("Rain author revision changed")
    RESULT_ROOT.mkdir(parents=True, exist_ok=True)
    started = time.time()
    outputs = []
    for mode, file_mode in CONDITIONS:
        output = RESULT_ROOT / f"rain-ep-c1-s{args.seed}-{file_mode}.json"
        if output.exists():
            raise FileExistsError(output)
        command = [
            sys.executable, str(ENDPOINT),
            "--author-root", str(author_root),
            "--device", "cuda", "--adapter", "layer", "--mode", mode,
            "--bias-ratio", "0.01", "--predictor-rate", "0.2",
            "--layer-calibration-steps", "1", "--epochs", "3",
            "--train-limit", "10000", "--test-limit", "2000",
            "--evaluation-split", "train_holdout", "--data-seed", "6100",
            "--batch-size", "128", "--training-iterations", "12",
            "--inference-iterations", "30", "--seed", str(args.seed),
            "--output", str(output),
        ]
        subprocess.run(command, cwd=author_root, check=True)
        outputs.append(str(output.relative_to(ROOT)))
    launch = {
        "stage": "rain_ep_bias_c1",
        "seed": args.seed,
        "conditions": [condition for condition, _ in CONDITIONS],
        "outputs": outputs,
        "author_revision": revision(author_root),
        "sdil_revision": revision(ROOT),
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "wall_seconds": time.time() - started,
    }
    path = RESULT_ROOT / f"launch-s{args.seed}.json"
    path.write_text(json.dumps(launch, indent=2, sort_keys=True) + "\n")
    print(json.dumps(launch, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()