#!/usr/bin/env python3
"""
BabyLM Challenge 2026 - Batch Experiment Runner

Reads experiments.csv, generates YAML configs, runs training, and collects results.

Usage:
    # Run all Phase 1 experiments
    python run_experiments.py --phase 1

    # Run specific experiments by ID
    python run_experiments.py --ids P1.1 P1.2 P1.3

    # Run all experiments in a phase range
    python run_experiments.py --phase 1 2 3A 3B

    # Dry run (generate configs but don't train)
    python run_experiments.py --phase 1 --dry-run

    # Resume (skip already completed experiments)
    python run_experiments.py --phase 1 --resume
"""

import argparse
import csv
import json
import os
import subprocess
import sys
import time
import yaml
from pathlib import Path
from datetime import datetime

ROOT = Path(__file__).resolve().parent
EXPERIMENTS_CSV = ROOT / "experiments.csv"
CONFIGS_DIR = ROOT / "experiment_configs"
RESULTS_DIR = ROOT / "experiment_results"
RESULTS_JSON = RESULTS_DIR / "all_results.json"


# ═══════════════════════════════════════════════════════════════════════
# CSV → YAML config conversion
# ═══════════════════════════════════════════════════════════════════════

# Maps CSV column → (yaml_section, yaml_key)
CSV_TO_YAML = {
    # model section
    "arch":               ("model", "arch"),
    "hidden_size":        ("model", "hidden_size"),
    "num_layers":         ("model", "num_layers"),
    "num_heads":          ("model", "num_heads"),
    "intermediate_size":  ("model", "intermediate_size"),
    "use_rope":           ("model", "use_rope"),
    "use_geglu":          ("model", "use_geglu"),
    "use_pre_norm":       ("model", "use_pre_norm"),
    "use_attention_gate": ("model", "use_attention_gate"),
    "use_dwa":            ("model", "use_dwa"),
    "use_moe":            ("model", "use_moe"),
    "moe_num_experts":    ("model", "moe_num_experts"),
    "moe_top_k":          ("model", "moe_top_k"),
    "moe_expert_size":    ("model", "moe_expert_size"),
    "moe_freq_penalty":   ("model", "moe_freq_penalty"),
    "use_attn_res":       ("model", "use_attn_res"),
    "attn_res_num_blocks":("model", "attn_res_num_blocks"),
    "dropout":            ("model", "dropout"),
    "position_bucket_size":("model", "position_bucket_size"),
    "z_loss_weight":      ("model", "z_loss_weight"),
    "rope_theta":         ("model", "rope_theta"),
    "rtd_lambda":         ("model", "rtd_lambda"),
    "gen_size_ratio":     ("model", "gen_size_ratio"),
    # data section
    "data":               ("data", "train_file"),  # needs path mapping
    "tokenizer":          ("data", "tokenizer"),
    "seq_len":            ("data", "max_seq_len"),
    # embedding section
    "embedding":          ("embedding", "type"),
    "embedding_init":     ("embedding", "init"),
    # training section
    "objective":          ("training", "objective"),
    "epochs":             ("training", "epochs"),
    "batch_size":         ("training", "batch_size"),
    "lr":                 ("training", "learning_rate"),
    "weight_decay":       ("training", "weight_decay"),
    "warmup_ratio":       ("training", "warmup_ratio"),
    "max_grad_norm":      ("training", "max_grad_norm"),
    "grad_accum":         ("training", "gradient_accumulation_steps"),
    "mntp_ratio":         ("training", "mntp_ratio"),
    "seed":               ("training", "seed"),
    # masking section
    "masking":            ("masking", "type"),
    "mask_ratio":         ("masking", "mask_ratio"),
    "mask_ratio_end":     ("masking", "mask_ratio_end"),
    "amlm_lambda":        ("masking", "amlm_lambda"),
    "amlm_update_interval":("masking", "amlm_update_interval"),
    "amlm_min_ratio":     ("masking", "amlm_min_ratio"),
    "amlm_max_ratio":     ("masking", "amlm_max_ratio"),
    # optimizer section
    "optimizer":          ("optimizer", "type"),
    "optimizer_betas":    ("optimizer", "betas"),  # needs parsing
    "forgetter":          ("optimizer", "forgetter"),
    # distillation section
    "kd_enabled":         ("distillation", "enabled"),
    "kd_teacher":         ("distillation", "teacher_model"),
    "kd_temperature":     ("distillation", "temperature"),
    "kd_alpha":           ("distillation", "alpha"),
    # checkpoint section
    "ckpt_averaging":     ("checkpoint", "averaging"),
    "ckpt_avg_last_k":    ("checkpoint", "avg_last_k"),
}

# Data ID → train file path
DATA_PATH_MAP = {
    "sample_A":          "data/8_sample_A/train.txt",
    "sample_B":          "data/8_sample_B/train.txt",
    "sample_C":          "data/8_sample_C/train.txt",
    "sample_D":          "data/8_sample_D/train.txt",
    "champion_replica":  "data/8_champion_replica/train.txt",
    "B_paraphrase":      "data/8_B_paraphrase/train.txt",
    "B_variation_sets":  "data/8_B_variation_sets/train.txt",
    "B_recombitext":     "data/8_B_recombitext/train.txt",
    "B_cd_synth":        "data/8_B_cd_synth/train.txt",
    "B_mattr":           "data/8_B_mattr/train.txt",
    "eval_mixed":        "data/8_eval_mixed/train.txt",
    "no_eval":           "data/8_no_eval/train.txt",
    "full_augment":      "data/8_full_augment/train.txt",
    "B_fineweb33":       "data/8_B_fineweb33/train.txt",
    "B_fineweb67":       "data/8_B_fineweb67/train.txt",
    "B_knowledge":       "data/8_B_knowledge/train.txt",
    "B_para_fineweb":    "data/8_B_para_fineweb/train.txt",
    "3way_equal":        "data/8_3way_equal/train.txt",
}


def parse_value(val: str):
    """Parse a CSV string value to the appropriate Python type."""
    if val == "":
        return None
    if val in ("True", "true"):
        return True
    if val in ("False", "false"):
        return False
    if val == "none":
        return None
    try:
        return int(val)
    except ValueError:
        pass
    try:
        return float(val)
    except ValueError:
        pass
    return val


def parse_betas(val: str):
    """Parse betas string like '(0.9,0.98)' to list [0.9, 0.98]."""
    if not val or val == "None":
        return [0.9, 0.98]
    cleaned = val.strip("()")
    parts = [float(x.strip()) for x in cleaned.split(",")]
    return parts


def csv_row_to_yaml(row: dict) -> dict:
    """Convert a CSV experiment row to a YAML config dict."""
    yaml_dict = {"name": row["name"]}
    sections = {}

    for csv_col, (section, key) in CSV_TO_YAML.items():
        val = row.get(csv_col)
        if val is None or val == "":
            continue

        if section not in sections:
            sections[section] = {}

        # Special handling
        if csv_col == "data":
            # Map data ID to file path
            data_id = val
            path = DATA_PATH_MAP.get(data_id, f"data/8_{data_id}/train.txt")
            sections[section][key] = path
            continue

        if csv_col == "optimizer_betas":
            sections[section][key] = parse_betas(val)
            continue

        if csv_col == "masking" and val == "none":
            sections[section][key] = "standard"
            continue

        sections[section][key] = parse_value(val)

    yaml_dict.update(sections)
    return yaml_dict


# ═══════════════════════════════════════════════════════════════════════
# Load experiments from CSV
# ═══════════════════════════════════════════════════════════════════════

def load_experiments(csv_path: Path) -> list[dict]:
    """Load experiments from CSV, skipping comment rows."""
    experiments = []
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            if row.get("id", "").startswith("##"):
                continue
            experiments.append(row)
    return experiments


def filter_experiments(experiments: list[dict],
                       phases: list[str] = None,
                       ids: list[str] = None) -> list[dict]:
    """Filter experiments by phase or ID."""
    if ids:
        return [e for e in experiments if e["id"] in ids]
    if phases:
        return [e for e in experiments if e.get("phase") in phases]
    return experiments


# ═══════════════════════════════════════════════════════════════════════
# Run a single experiment
# ═══════════════════════════════════════════════════════════════════════

def run_experiment(exp: dict, dry_run: bool = False) -> dict:
    """Generate YAML config, run training, return result info."""
    exp_id = exp["id"]
    exp_name = exp["name"]

    # Generate YAML config
    CONFIGS_DIR.mkdir(parents=True, exist_ok=True)
    yaml_config = csv_row_to_yaml(exp)
    yaml_path = CONFIGS_DIR / f"{exp_id.replace('.', '_')}.yaml"

    with open(yaml_path, "w") as f:
        yaml.dump(yaml_config, f, default_flow_style=False, allow_unicode=True)

    print(f"\n{'='*60}")
    print(f"  Experiment {exp_id}: {exp_name}")
    print(f"  Config: {yaml_path}")
    print(f"  Phase: {exp.get('phase', '?')}")
    print(f"{'='*60}")

    if dry_run:
        print(f"  [DRY RUN] Would run: python -m scripts.03_training.train --config {yaml_path}")
        return {"id": exp_id, "name": exp_name, "status": "dry_run", "config": str(yaml_path)}

    # Check if data file exists
    data_id = exp.get("data", "sample_B")
    data_path = ROOT / DATA_PATH_MAP.get(data_id, f"data/8_{data_id}/train.txt")
    if not data_path.exists():
        print(f"  [SKIP] Data file not found: {data_path}")
        return {"id": exp_id, "name": exp_name, "status": "skipped", "reason": f"data not found: {data_path}"}

    # Run training
    start_time = time.time()
    cmd = [
        sys.executable, "-m", "scripts.03_training.train",
        "--config", str(yaml_path),
    ]

    try:
        result = subprocess.run(
            cmd, cwd=str(ROOT),
            capture_output=True, text=True, timeout=7200,  # 2 hour timeout
        )
        elapsed = time.time() - start_time

        # Save stdout/stderr
        RESULTS_DIR.mkdir(parents=True, exist_ok=True)
        log_path = RESULTS_DIR / f"{exp_id.replace('.', '_')}.log"
        with open(log_path, "w") as f:
            f.write(f"=== STDOUT ===\n{result.stdout}\n\n=== STDERR ===\n{result.stderr}")

        status = "completed" if result.returncode == 0 else "failed"
        exp_result = {
            "id": exp_id, "name": exp_name, "status": status,
            "returncode": result.returncode,
            "elapsed_sec": round(elapsed),
            "config": str(yaml_path), "log": str(log_path),
            "timestamp": datetime.now().isoformat(),
        }

        # ── Capture eval results if training succeeded ──
        if status == "completed":
            eval_results_file = ROOT / "models" / exp_name / "eval_results" / "eval_results.json"
            if eval_results_file.exists():
                try:
                    with open(eval_results_file) as ef:
                        eval_data = json.load(ef)
                    summary = eval_data.get("summary", {})
                    exp_result["eval_summary"] = summary
                    # Extract key scores for easy comparison
                    if "blimp" in summary:
                        exp_result["blimp"] = summary["blimp"]
                    if "ewok" in summary:
                        exp_result["ewok"] = summary["ewok"]
                    if "avg_zero_shot" in summary:
                        exp_result["avg_zero_shot"] = summary["avg_zero_shot"]
                    print(f"  Eval results captured: BLiMP={summary.get('blimp', 'N/A')}, "
                          f"Avg={summary.get('avg_zero_shot', 'N/A')}")
                except Exception as e:
                    print(f"  WARNING: Could not read eval results: {e}")
            else:
                print(f"  NOTE: No eval results found at {eval_results_file}")

        return exp_result

    except subprocess.TimeoutExpired:
        return {"id": exp_id, "name": exp_name, "status": "timeout"}
    except Exception as e:
        return {"id": exp_id, "name": exp_name, "status": "error", "error": str(e)}


# ═══════════════════════════════════════════════════════════════════════
# Results management
# ═══════════════════════════════════════════════════════════════════════

def load_results() -> dict:
    """Load existing results from JSON."""
    if RESULTS_JSON.exists():
        with open(RESULTS_JSON) as f:
            return json.load(f)
    return {}


def save_result(result: dict):
    """Append one result to the results JSON."""
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    results = load_results()
    results[result["id"]] = result
    with open(RESULTS_JSON, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)


def is_completed(exp_id: str) -> bool:
    """Check if an experiment was already completed."""
    results = load_results()
    return results.get(exp_id, {}).get("status") == "completed"


# ═══════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="BabyLM Batch Experiment Runner")
    parser.add_argument("--phase", nargs="*", help="Phase(s) to run (e.g., 1 2 3A)")
    parser.add_argument("--ids", nargs="*", help="Specific experiment IDs (e.g., P1.1 P2.3)")
    parser.add_argument("--dry-run", action="store_true", help="Generate configs only, don't train")
    parser.add_argument("--resume", action="store_true", help="Skip already completed experiments")
    parser.add_argument("--csv", default=str(EXPERIMENTS_CSV), help="Path to experiments CSV")
    parser.add_argument("--list", action="store_true", help="List experiments without running")
    args = parser.parse_args()

    # Load and filter
    experiments = load_experiments(Path(args.csv))
    filtered = filter_experiments(experiments, phases=args.phase, ids=args.ids)

    if not filtered:
        print("No experiments matched the filter criteria.")
        print(f"  Available phases: {sorted(set(e.get('phase','?') for e in experiments))}")
        return

    # List mode
    if args.list:
        print(f"\n{'ID':<10} {'Phase':<6} {'Name':<30} {'Epochs':<6} {'Data':<15}")
        print("-" * 75)
        for exp in filtered:
            print(f"{exp['id']:<10} {exp.get('phase','?'):<6} {exp['name']:<30} "
                  f"{exp.get('epochs','?'):<6} {exp.get('data','?'):<15}")
        print(f"\nTotal: {len(filtered)} experiments")
        return

    # Run
    print(f"\n{'='*60}")
    print(f"  BabyLM Batch Runner")
    print(f"  Experiments: {len(filtered)}")
    print(f"  Dry run: {args.dry_run}")
    print(f"  Resume: {args.resume}")
    print(f"{'='*60}")

    completed = 0
    skipped = 0
    failed = 0

    for i, exp in enumerate(filtered):
        if args.resume and is_completed(exp["id"]):
            print(f"  [{i+1}/{len(filtered)}] {exp['id']} — already completed, skipping")
            skipped += 1
            continue

        result = run_experiment(exp, dry_run=args.dry_run)
        save_result(result)

        if result["status"] == "completed":
            completed += 1
        elif result["status"] in ("failed", "error", "timeout"):
            failed += 1
            print(f"  [WARNING] {exp['id']} {result['status']}")

    # ── Print summary table with eval scores ──
    print(f"\n{'='*80}")
    print(f"  Results Summary")
    print(f"{'='*80}")
    all_results = load_results()
    print(f"  {'ID':<10} {'Name':<25} {'Status':<10} {'BLiMP':<8} {'Avg ZS':<8} {'Time':<8}")
    print(f"  {'-'*70}")
    for exp in filtered:
        r = all_results.get(exp["id"], {})
        st = r.get("status", "?")
        blimp = r.get("blimp", "")
        avg_zs = r.get("avg_zero_shot", "")
        elapsed = r.get("elapsed_sec", "")
        blimp_str = f"{blimp:.1f}" if isinstance(blimp, (int, float)) else str(blimp)
        avg_str = f"{avg_zs:.1f}" if isinstance(avg_zs, (int, float)) else str(avg_zs)
        time_str = f"{elapsed}s" if elapsed else ""
        print(f"  {exp['id']:<10} {exp['name']:<25} {st:<10} {blimp_str:<8} {avg_str:<8} {time_str:<8}")
    print(f"\n  Done: {completed} completed, {skipped} skipped, {failed} failed")
    print(f"  Results: {RESULTS_JSON}")
    print(f"{'='*80}")


if __name__ == "__main__":
    main()