#!/usr/bin/env python3
"""Convenience launcher for mutation and depth sweep runs.

Usage:
    python scripts/run_sweep.py mutation     # Run all 3 mutation rate experiments
    python scripts/run_sweep.py depth        # Run all 3 depth cap experiments
    python scripts/run_sweep.py all          # Run all 6 experiments
    python scripts/run_sweep.py smoke        # Quick smoke test (100k each)
"""

import sys
import time
from pathlib import Path

from abiogenesis.experiment import Experiment


MUTATION_CONFIGS = [
    {
        "seed": 42,
        "mutation_rate": 1e-5,
        "output_dir": "data/bff_mutation_1e-5/",
        "label": "mutation_1e-5",
    },
    {
        "seed": 42,
        "mutation_rate": 1e-3,
        "output_dir": "data/bff_mutation_1e-3/",
        "label": "mutation_1e-3",
    },
    {
        "seed": 42,
        "mutation_rate": 1e-1,
        "output_dir": "data/bff_mutation_1e-1/",
        "label": "mutation_1e-1",
    },
]

DEPTH_CONFIGS = [
    {
        "seed": 42,
        "max_tree_depth": 10,
        "output_dir": "data/bff_depth_10/",
        "label": "depth_cap_10",
    },
    {
        "seed": 42,
        "max_tree_depth": 20,
        "output_dir": "data/bff_depth_20/",
        "label": "depth_cap_20",
    },
    {
        "seed": 42,
        "max_tree_depth": 0,
        "output_dir": "data/bff_depth_uncapped/",
        "label": "depth_uncapped",
    },
]


def run_configs(configs: list[dict], n_interactions: int = 10_000_000):
    """Run a list of experiment configs sequentially."""
    for i, cfg in enumerate(configs):
        run_cfg = {k: v for k, v in cfg.items() if k != "label"}
        label = cfg.get("label", f"run_{i}")
        print(f"\n{'='*60}")
        print(f"  {label} ({i+1}/{len(configs)})")
        print(f"{'='*60}\n")

        # Check for existing checkpoints to resume from
        out_dir = Path(run_cfg.get("output_dir", "data/"))
        existing = sorted(out_dir.glob("checkpoint_*.npz")) if out_dir.exists() else []
        if existing:
            last_cp = str(existing[-1])
            print(f"  Resuming from {last_cp}")
            exp = Experiment.load_checkpoint(last_cp)
            remaining = n_interactions - exp.interaction_count
            if remaining <= 0:
                print(f"  Already complete ({exp.interaction_count:,} interactions)")
                print(f"  Final compression: {exp.snapshot_compression[-1]:.4f}")
                continue
        else:
            exp = Experiment(config=run_cfg)
            remaining = n_interactions

        t0 = time.time()
        exp.run(n_interactions=remaining)
        elapsed = time.time() - t0

        print(f"\n  Completed {label} in {elapsed:.1f}s")
        print(f"  Final compression: {exp.snapshot_compression[-1]:.4f}")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    mode = sys.argv[1].lower()

    if mode == "mutation":
        run_configs(MUTATION_CONFIGS)
    elif mode == "depth":
        run_configs(DEPTH_CONFIGS)
    elif mode == "all":
        run_configs(MUTATION_CONFIGS)
        run_configs(DEPTH_CONFIGS)
    elif mode == "smoke":
        print("Running smoke test (100k interactions each)...")
        all_configs = MUTATION_CONFIGS + DEPTH_CONFIGS
        run_configs(all_configs, n_interactions=100_000)
    else:
        print(f"Unknown mode: {mode}")
        print(__doc__)
        sys.exit(1)


if __name__ == "__main__":
    main()
