#!/usr/bin/env python3
"""
Component verification: run each optimization dimension for 3 epochs,
verify loss decreases normally (no nan, no divergence).

Usage:
    python verify_components.py              # run all
    python verify_components.py --group arch # run only architecture group
    python verify_components.py --resume     # skip already passed
"""

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

ROOT = Path(__file__).resolve().parent
RESULTS_FILE = ROOT / "verify_results.json"
CONFIGS_DIR = ROOT / "verify_configs"

# ═══════════════════════════════════════════════════════════════════
# Base config (known working: AdamW + lr=5e-4, verified with full run)
# ═══════════════════════════════════════════════════════════════════
BASE = {
    "name": "verify",
    "data": {"train_file": "data/8_sample_B/train.txt", "tokenizer": "bpe", "max_seq_len": 128, "packing": "concat"},
    "model": {
        "arch": "gpt_bert", "hidden_size": 384, "num_layers": 12, "num_heads": 6,
        "intermediate_size": 1280, "dropout": 0.1, "use_geglu": True, "use_pre_norm": True,
        "z_loss_weight": 0.0001, "use_moe": False, "use_attn_res": False,
    },
    "embedding": {"type": "standard", "init": "random"},
    "training": {
        "objective": "gpt_bert", "epochs": 3, "batch_size": 64,
        "learning_rate": 0.0005, "weight_decay": 0.1, "warmup_ratio": 0.06,
        "max_grad_norm": 2.0, "mntp_ratio": 15, "seed": 42,
    },
    "masking": {"type": "standard", "mask_ratio": 0.30, "mask_ratio_end": 0.15},
    "optimizer": {"type": "adamw", "betas": [0.9, 0.98], "forgetter": False},
    "checkpoint": {"save_every_epoch": False, "save_aoa_checkpoints": False},
}


def deep_copy(d):
    import copy
    return copy.deepcopy(d)


# ═══════════════════════════════════════════════════════════════════
# Test configurations: each is (name, group, overrides)
# ═══════════════════════════════════════════════════════════════════
TESTS = [
    # ── D: Model Architecture ──
    ("D0-gpt2",        "arch", {"model": {"arch": "gpt2"}, "training": {"objective": "clm"}, "masking": {"type": "standard", "mask_ratio": 0.0, "mask_ratio_end": 0.0}}),
    ("D1-gptbert",     "arch", {}),  # base config IS gpt-bert
    ("D2-modernbert",  "arch", {"model": {"arch": "modernized_bert"}, "training": {"objective": "mlm"}}),
    ("D3-xlstm",       "arch", {"model": {"arch": "xlstm"}, "training": {"objective": "clm"}, "masking": {"type": "standard", "mask_ratio": 0.0, "mask_ratio_end": 0.0}}),
    ("D4-rtd",         "arch", {"model": {"arch": "rtd"}, "training": {"objective": "rtd"}}),
    ("D5-moe-gptbert", "arch", {"model": {"use_moe": True}}),
    ("D5-moe-gpt2",    "arch", {"model": {"arch": "gpt2", "use_moe": True}, "training": {"objective": "clm"}, "masking": {"type": "standard", "mask_ratio": 0.0, "mask_ratio_end": 0.0}}),
    ("D5-moe-rtd",     "arch", {"model": {"arch": "rtd", "use_moe": True}, "training": {"objective": "rtd"}}),
    ("D6-attnres-gptbert", "arch", {"model": {"use_attn_res": True}}),
    ("D6-attnres-gpt2",   "arch", {"model": {"arch": "gpt2", "use_attn_res": True}, "training": {"objective": "clm"}, "masking": {"type": "standard", "mask_ratio": 0.0, "mask_ratio_end": 0.0}}),
    ("D5D6-moe-attnres",  "arch", {"model": {"use_moe": True, "use_attn_res": True}}),

    # ── C: Embedding ──
    ("C0-standard",    "embed", {}),  # base
    ("C1-nhot",        "embed", {"embedding": {"type": "nhot"}}),
    # C2 FastText needs trained model, skip for now

    # ── B: Tokenizer ──
    ("B0-bpe",         "tok",   {}),  # base
    ("B1-morfessor",   "tok",   {"data": {"tokenizer": "morfessor_bpe"}}),

    # ── E: Masking ──
    ("E0-standard",    "mask",  {}),  # base (30%->15% decay)
    ("E1-amlm",        "mask",  {"masking": {"type": "amlm"}}),
    ("E3-frequency",   "mask",  {"masking": {"type": "frequency"}}),

    # ── F: Optimizer ──
    ("F0-adam",        "optim", {"optimizer": {"type": "adam"}}),
    ("F1-adamw",       "optim", {}),  # base
    ("F2-lamb-lr5e4",  "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.0005}}),
    ("F2-lamb-lr3e3",  "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.003}}),
    ("F2-lamb-lr5e3",  "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.005}}),
    ("F2-lamb-lr8e3",  "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.008}}),
    ("F2-lamb-lr1e2",  "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.01}}),
    ("F2-lamb-lr14e3", "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.0141}}),
    ("F2-lamb-nofp16", "optim", {"optimizer": {"type": "lamb"}, "training": {"learning_rate": 0.0141, "fp16": False}}),
    ("F3-forgetter",   "optim", {"optimizer": {"forgetter": True}, "training": {"weight_decay": 1.0}}),
    ("F4-muon",        "optim", {"optimizer": {"type": "muon"}, "training": {"learning_rate": 0.01}}),

    # ── G: Hyperparams ──
    ("G6-sentence",    "hyper", {"data": {"packing": "sentence"}}),

    # ── Combos (known good from smoke test, verify loss quality) ──
    ("combo-amlm-nhot-fgt",     "combo", {"masking": {"type": "amlm"}, "embedding": {"type": "nhot"}, "optimizer": {"forgetter": True}, "training": {"weight_decay": 1.0}}),
    ("combo-moe-attnres-amlm",  "combo", {"model": {"use_moe": True, "use_attn_res": True}, "masking": {"type": "amlm"}}),
]


def merge_config(base, overrides):
    """Deep merge overrides into base config."""
    result = deep_copy(base)
    for key, val in overrides.items():
        if isinstance(val, dict) and key in result and isinstance(result[key], dict):
            result[key].update(val)
        else:
            result[key] = val
    return result


def run_test(name, overrides):
    """Run one verification test. Returns (status, final_loss, time_sec)."""
    cfg = merge_config(BASE, overrides)
    cfg["name"] = f"verify_{name}"

    CONFIGS_DIR.mkdir(parents=True, exist_ok=True)
    yaml_path = CONFIGS_DIR / f"{name}.yaml"
    with open(yaml_path, "w") as f:
        yaml.dump(cfg, f, default_flow_style=False)

    print(f"  [{name}]", end=" ", flush=True)

    start = time.time()
    try:
        result = subprocess.run(
            [sys.executable, "-u", "-m", "scripts.03_training.train",
             "--config", str(yaml_path), "--skip-eval"],
            cwd=str(ROOT), capture_output=True, text=True,
            timeout=1800, env={**__import__('os').environ, "PYTHONUNBUFFERED": "1"},
        )
        elapsed = time.time() - start
        output = result.stdout + "\n" + result.stderr

        # Parse final loss
        import re
        losses = re.findall(r"Loss ([\d.]+|nan|inf)", output)
        final_loss = losses[-1] if losses else "?"

        if result.returncode != 0:
            err = result.stderr.strip().split("\n")[-1][:100]
            print(f"FAIL ({elapsed:.0f}s) — {err}")
            return "fail", final_loss, elapsed, err

        if final_loss in ("nan", "inf"):
            print(f"NAN ({elapsed:.0f}s) — loss diverged")
            return "nan", final_loss, elapsed, "loss diverged"

        # Check loss is reasonable (< 10 for 3 epochs)
        try:
            fl = float(final_loss)
            if fl > 10:
                print(f"HIGH ({elapsed:.0f}s) — loss={fl:.4f}")
                return "high_loss", final_loss, elapsed, f"loss={fl}"
            print(f"OK ({elapsed:.0f}s) loss={fl:.4f}")
            return "pass", final_loss, elapsed, ""
        except ValueError:
            print(f"OK ({elapsed:.0f}s) loss={final_loss}")
            return "pass", final_loss, elapsed, ""

    except subprocess.TimeoutExpired:
        print(f"TIMEOUT")
        return "timeout", "?", 1800, "timeout"
    except Exception as e:
        print(f"ERROR — {e}")
        return "error", "?", 0, str(e)


def load_results():
    if RESULTS_FILE.exists():
        with open(RESULTS_FILE) as f:
            return json.load(f)
    return {}


def save_results(results):
    with open(RESULTS_FILE, "w") as f:
        json.dump(results, f, indent=2)


def main():
    parser = argparse.ArgumentParser(description="Verify all components work correctly")
    parser.add_argument("--group", choices=["arch", "embed", "tok", "mask", "optim", "hyper", "combo"],
                        help="Run only this group")
    parser.add_argument("--resume", action="store_true", help="Skip already passed tests")
    parser.add_argument("--test", help="Run a specific test by name")
    args = parser.parse_args()

    tests = TESTS
    if args.group:
        tests = [(n, g, o) for n, g, o in tests if g == args.group]
    if args.test:
        tests = [(n, g, o) for n, g, o in tests if n == args.test]

    results = load_results() if args.resume else {}
    passed = failed = skipped = 0

    print(f"\n{'='*60}")
    print(f"  Component Verification: {len(tests)} tests, 3 epochs each")
    print(f"{'='*60}\n")

    current_group = None
    for name, group, overrides in tests:
        if group != current_group:
            current_group = group
            print(f"\n── {group.upper()} ──")

        if args.resume and results.get(name, {}).get("status") == "pass":
            skipped += 1
            continue

        status, loss, elapsed, err = run_test(name, overrides)
        results[name] = {"status": status, "loss": loss, "time": round(elapsed), "error": err}
        save_results(results)

        if status == "pass":
            passed += 1
        else:
            failed += 1

    # Summary
    print(f"\n{'='*60}")
    print(f"  {'PASS':<8} {'FAIL/NAN':<10} {'SKIP':<8}")
    print(f"  {passed:<8} {failed:<10} {skipped:<8}")

    if failed > 0:
        print(f"\n  Issues found:")
        for name, r in results.items():
            if r.get("status") not in ("pass", None):
                print(f"    {name:<30} {r['status']:<8} loss={r.get('loss','?')} {r.get('error','')[:60]}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()