scripts/tools/generate_experiments.py

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

Generates experiments.csv with ~380 experiments across 22 phases.
Matches the phased plan in notes/experiment_plan.md.

Run: python generate_experiments.py
"""

import csv
import itertools

# ── Output ──
OUTPUT = "experiments.csv"

# ── CSV columns ──
COLUMNS = [
    "id", "name", "phase",
    # Architecture
    "arch", "objective", "hidden_size", "num_layers", "num_heads", "intermediate_size",
    "use_rope", "use_geglu", "use_pre_norm", "use_attention_gate", "use_dwa",
    "use_moe", "moe_num_experts", "moe_top_k", "moe_expert_size", "moe_freq_penalty",
    "use_attn_res", "attn_res_num_blocks",
    "position_bucket_size", "z_loss_weight", "rope_theta",
    "rtd_lambda", "gen_size_ratio",
    # Data
    "data", "tokenizer", "embedding", "embedding_init",
    # Masking
    "masking", "mask_ratio", "mask_ratio_end", "mntp_ratio",
    "amlm_lambda", "amlm_update_interval", "amlm_min_ratio", "amlm_max_ratio",
    # Optimizer
    "optimizer", "lr", "optimizer_betas", "weight_decay", "forgetter",
    # Training
    "epochs", "batch_size", "grad_accum", "seq_len",
    "dropout", "warmup_ratio", "max_grad_norm",
    # KD
    "kd_enabled", "kd_teacher", "kd_temperature", "kd_alpha",
    # Checkpoint
    "ckpt_averaging", "ckpt_avg_last_k",
    # Meta
    "seed", "status", "notes",
]

# ═══════════════════════════════════════════════════════════════════════
# Baseline defaults
# ═══════════════════════════════════════════════════════════════════════
BASELINE = dict(
    arch="gpt_bert", objective="gpt_bert",
    hidden_size=384, num_layers=12, num_heads=6, intermediate_size=1280,
    use_rope=False, use_geglu=True, use_pre_norm=True,
    use_attention_gate=False, use_dwa=False,
    use_moe=False, moe_num_experts=32, moe_top_k=4, moe_expert_size=48, moe_freq_penalty=0.01,
    use_attn_res=False, attn_res_num_blocks=4,
    position_bucket_size=32, z_loss_weight=0.0001, rope_theta=10000.0,
    rtd_lambda=50.0, gen_size_ratio=0.33,
    data="sample_B", tokenizer="bpe", embedding="standard", embedding_init="random",
    masking="standard", mask_ratio=0.30, mask_ratio_end=0.15, mntp_ratio=15,
    amlm_lambda=0.2, amlm_update_interval=200, amlm_min_ratio=0.05, amlm_max_ratio=0.50,
    optimizer="LAMB", lr=0.0141, optimizer_betas="(0.9,0.98)", weight_decay=0.1, forgetter=False,
    epochs=10, batch_size=64, grad_accum=1, seq_len=128,
    dropout=0.1, warmup_ratio=0.06, max_grad_norm=2.0,
    kd_enabled=False, kd_teacher="", kd_temperature=4.0, kd_alpha=0.5,
    ckpt_averaging=False, ckpt_avg_last_k=3,
    seed=42, status="planned", notes="",
)

# ═══════════════════════════════════════════════════════════════════════
# Architecture presets
# ═══════════════════════════════════════════════════════════════════════
ARCH_PRESETS = {
    "gpt_bert":   dict(arch="gpt_bert",       objective="gpt_bert", optimizer="LAMB",  masking="standard", mntp_ratio=15, mask_ratio=0.30, mask_ratio_end=0.15, lr=0.005),
    "gpt2":       dict(arch="gpt2",           objective="clm",      optimizer="AdamW", masking="none",     mntp_ratio=0,  mask_ratio=0,    mask_ratio_end=0, lr=0.0005),
    "modernbert": dict(arch="modernized_bert", objective="mlm",      optimizer="AdamW", masking="standard", mntp_ratio=0,  mask_ratio=0.30, mask_ratio_end=0.15, lr=0.0005),
    "xlstm":      dict(arch="xlstm",          objective="clm",      optimizer="AdamW", masking="none",     mntp_ratio=0,  mask_ratio=0,    mask_ratio_end=0, lr=0.0005),
    "rtd":        dict(arch="rtd",            objective="rtd",      optimizer="LAMB",  masking="standard", mntp_ratio=0,  mask_ratio=0.15, mask_ratio_end=0.15, lr=0.005),
}

# ═══════════════════════════════════════════════════════════════════════
# Sweep values
# ═══════════════════════════════════════════════════════════════════════

# Learning rates
LRS_LAMB  = [0.003, 0.005, 0.007, 0.008, 0.009, 0.010, 0.012, 0.015, 0.018, 0.020, 0.025, 0.030]
LRS_ADAMW = [0.0003, 0.0005, 0.001, 0.002, 0.003, 0.005]
LRS_MUON  = [0.005, 0.008, 0.014, 0.020, 0.030]

# Regularization
DROPOUTS   = [0.00, 0.02, 0.05, 0.08, 0.10, 0.15, 0.20]
WDS        = [0.01, 0.03, 0.05, 0.08, 0.15, 0.20, 1.00]
WARMUPS    = [0.00, 0.03, 0.10, 0.15, 0.20]
GRAD_NORMS = [0.5, 1.0, 3.0, 5.0]

# Masking
MASK_SCHEDULES = [
    (0.15, 0.15), (0.20, 0.15), (0.25, 0.15), (0.35, 0.15),
    (0.40, 0.15), (0.40, 0.20), (0.45, 0.15), (0.30, 0.10),
]
MNTP_RATIOS = [1, 3, 5, 7, 10, 20]

# Batch / Sequence
BATCH_CONFIGS = [  # (batch_size, grad_accum, effective)
    (16, 1, 16), (32, 1, 32), (64, 2, 128), (64, 4, 256),
    (64, 8, 512), (128, 1, 128), (128, 2, 256), (256, 1, 256),
]
SEQ_LENS = [64, 96, 192, 256, 384, 512]

# Model shapes: (hidden, layers, heads, intermediate)
MODEL_SHAPES = [
    (256, 16, 4, 854), (256, 18, 4, 854), (320, 12, 5, 1067),
    (384, 16, 6, 1280), (448, 10, 7, 1494), (512, 8, 8, 1706),
    (512, 10, 8, 1706), (576, 8, 9, 1920),
]

# MoE configs: (num_experts, top_k, expert_size)
MOE_CONFIGS = [
    (16, 2, 96), (16, 4, 96), (32, 2, 48), (32, 8, 48),
    (64, 4, 24), (32, 4, 84),
]
MOE_FREQ_PENALTIES = [0.001, 0.10]

# Data construction methods (id, description)
DATA_CONFIGS = [
    ("champion_replica",   "BabyLM33+FineWeb33+Cosmo34"),
    ("sample_A",           "Strategy A pure quality"),
    ("sample_B",           "Strategy B task quota"),
    ("sample_C",           "Strategy C weighted random"),
    ("sample_D",           "Strategy D embedding similarity"),
    ("B_paraphrase",       "sample_B + Paraphrase ~2M"),
    ("B_variation_sets",   "sample_B + Variation Sets ~1M"),
    ("B_recombitext",      "sample_B + RecombiText ~5M"),
    ("B_cd_synth",         "sample_B + Contrastive Decoding ~3M"),
    ("B_mattr",            "sample_B + MATTR ordering"),
    ("eval_mixed",         "eval data mixed into sampling pool"),
    ("no_eval",            "no eval data"),
    ("full_augment",       "Para + VS + MATTR"),
    ("B_fineweb33",        "sample_B + FineWeb 33%"),
    ("B_fineweb67",        "sample_B + FineWeb 67%"),
    ("B_knowledge",        "sample_B + ConceptNet + GenericsKB"),
    ("B_para_fineweb",     "sample_B + Para25% + FineWeb25%"),
    ("3way_equal",         "1:1:1 official + FineWeb + Cosmo"),
]

# KD teachers
KD_TEACHERS = [
    ("qwen3.5-1.5B", "Qwen/Qwen3.5-1.5B"),
    ("qwen3.5-3B",   "Qwen/Qwen3.5-3B"),
    ("qwen3.5-9B",   "Qwen/Qwen3.5-9B"),
]
KD_TEMPS  = [2.0, 4.0, 6.0, 8.0]
KD_ALPHAS = [0.2, 0.3, 0.5, 0.7, 0.9]

# Misc
BETAS        = ["(0.8,0.95)", "(0.9,0.95)", "(0.9,0.999)", "(0.95,0.999)"]
Z_LOSSES     = [0.0, 0.00005, 0.0005, 0.001]
BUCKET_SIZES = [8, 16, 64, 128]
SEEDS        = [7, 42, 123, 2024, 2026, 9999, 31415, 54321]


# ═══════════════════════════════════════════════════════════════════════
# Experiment generation
# ═══════════════════════════════════════════════════════════════════════
experiments = []
phase_counter = {}


def add(phase: int | str, name: str, notes: str = "", **overrides):
    """Add one experiment with overrides on top of baseline."""
    p = str(phase)
    if p not in phase_counter:
        phase_counter[p] = 0
    phase_counter[p] += 1
    idx = phase_counter[p]

    exp = dict(BASELINE)
    exp.update(overrides)
    exp["id"] = f"P{p}.{idx}"
    exp["name"] = name
    exp["phase"] = p
    exp["notes"] = notes
    experiments.append(exp)


def comment(text):
    experiments.append({"_comment": text})


# ═══════════════════════════════════════════════════════════════════════
# Phase 1: Architecture Comparison (5 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 1: Architecture Comparison (5 exp) ═══")
for arch_name, preset in ARCH_PRESETS.items():
    add(1, f"{arch_name}-baseline", f"Base {arch_name}", **preset)


# ═══════════════════════════════════════════════════════════════════════
# Phase 2: Core Feature Ablations (9 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 2: Core Feature Ablations (9 exp) ═══")
add(2, "gptbert-amlm",         "adaptive masking",         masking="amlm")
add(2, "gptbert-nhot",         "N-hot morphological emb",  embedding="nhot")
add(2, "gptbert-forgetter",    "FORGETTER reset per epoch", forgetter=True)
add(2, "gptbert-fasttext",     "FastText init",            embedding_init="fasttext")
add(2, "gptbert-morfessor",    "Morfessor+BPE tokenizer",  tokenizer="morfessor_bpe")
add(2, "gptbert-moe",          "MoE 32exp k=4 s=48",      use_moe=True)
add(2, "gptbert-attnres",      "AttnRes 4 blocks",         use_attn_res=True)
add(2, "gptbert-moe-attnres",  "MoE + AttnRes combined",   use_moe=True, use_attn_res=True)
add(2, "gptbert-sampleC",      "Strategy C data",          data="sample_C")


# ═══════════════════════════════════════════════════════════════════════
# Phase 3: Learning Rate Sweep (17 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 3: Learning Rate Sweep — 3A quick filter (12×3ep) + 3B promote (5×10ep) ═══")

# 3A: 3-epoch quick filter
for lr in LRS_LAMB:
    add("3A", f"lr-{lr}", f"LAMB lr={lr}", lr=lr, epochs=3)

# 3B: placeholders — top 3 + 2 interpolations promoted to 10ep
for i in range(1, 6):
    add("3B", f"lr-promote{i}", f"Top {i} from 3A → 10ep", epochs=10)


# ═══════════════════════════════════════════════════════════════════════
# Phase 4: Regularization Sweeps (25 experiments, mostly 3ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 4: Regularization Sweeps (25 exp) ═══")

# 4A: Dropout (7×3ep)
comment("── 4A: Dropout ──")
for d in DROPOUTS:
    add("4A", f"drop-{d}", f"dropout={d}", dropout=d, epochs=3)

# 4B: Weight decay (7×3ep)
comment("── 4B: Weight Decay ──")
for wd in WDS:
    add("4B", f"wd-{wd}", f"weight_decay={wd}", weight_decay=wd, epochs=3)

# 4C: Warmup (5×3ep)
comment("── 4C: Warmup Ratio ──")
for w in WARMUPS:
    add("4C", f"warmup-{w}", f"warmup_ratio={w}", warmup_ratio=w, epochs=3)

# 4D: Grad norm (4×3ep)
comment("── 4D: Grad Norm ──")
for gn in GRAD_NORMS:
    add("4D", f"gradnorm-{gn}", f"max_grad_norm={gn}", max_grad_norm=gn, epochs=3)

# 4E: Best combo verification (2×10ep)
comment("── 4E: Best Regularization Combo ──")
add("4E", "reg-best-combo1", "Best dropout+wd+warmup+gradnorm", epochs=10)
add("4E", "reg-best-combo2", "2nd best dropout variant", epochs=10)


# ═══════════════════════════════════════════════════════════════════════
# Phase 5: Feature Combinations (20 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 5: Feature Combinations (20 exp, using TUNED_BASELINE) ═══")

# Features: amlm(A), nhot(N), forgetter(F), moe(M), attnres(R)
feature_combos = [
    ("tuned-baseline",        {}),
    # Singles
    ("tuned-amlm",            dict(masking="amlm")),
    ("tuned-nhot",            dict(embedding="nhot")),
    ("tuned-fgt",             dict(forgetter=True)),
    ("tuned-moe",             dict(use_moe=True)),
    ("tuned-attnres",         dict(use_attn_res=True)),
    # Pairs
    ("tuned-AN",              dict(masking="amlm", embedding="nhot")),
    ("tuned-AF",              dict(masking="amlm", forgetter=True)),
    ("tuned-AM",              dict(masking="amlm", use_moe=True)),
    ("tuned-NF",              dict(embedding="nhot", forgetter=True)),
    ("tuned-NM",              dict(embedding="nhot", use_moe=True)),
    ("tuned-FM",              dict(forgetter=True, use_moe=True)),
    ("tuned-MR",              dict(use_moe=True, use_attn_res=True)),
    # Triples
    ("tuned-ANF",             dict(masking="amlm", embedding="nhot", forgetter=True)),
    ("tuned-ANM",             dict(masking="amlm", embedding="nhot", use_moe=True)),
    ("tuned-AFM",             dict(masking="amlm", forgetter=True, use_moe=True)),
    ("tuned-NFM",             dict(embedding="nhot", forgetter=True, use_moe=True)),
    # Quads
    ("tuned-ANFM",            dict(masking="amlm", embedding="nhot", forgetter=True, use_moe=True)),
    ("tuned-ANFR",            dict(masking="amlm", embedding="nhot", forgetter=True, use_attn_res=True)),
    # All five
    ("tuned-ANFMR",           dict(masking="amlm", embedding="nhot", forgetter=True, use_moe=True, use_attn_res=True)),
]
for combo_name, combo_overrides in feature_combos:
    add(5, combo_name, combo_name, **combo_overrides)


# ═══════════════════════════════════════════════════════════════════════
# Phase 6: Masking & Training Objective Tuning (20 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 6: Masking & Objective Tuning (20 exp) ═══")

# 6A: Mask schedule (8)
comment("── 6A: Mask Rate Schedule ──")
for mr_start, mr_end in MASK_SCHEDULES:
    add("6A", f"mask-{mr_start}to{mr_end}", f"mask {mr_start}->{mr_end}",
        mask_ratio=mr_start, mask_ratio_end=mr_end)

# 6B: MNTP:CLM ratio (6)
comment("── 6B: MNTP:CLM Ratio ──")
for ratio in MNTP_RATIOS:
    add("6B", f"mntp-{ratio}to1", f"MNTP:CLM={ratio}:1", mntp_ratio=ratio)

# 6C: AMLM parameters (6)
comment("── 6C: AMLM Parameters ──")
for lam in [0.05, 0.10, 0.30]:
    add("6C", f"amlm-lam{lam}", f"amlm_lambda={lam}", masking="amlm", amlm_lambda=lam)
add("6C", "amlm-ui50",  "fast update",  masking="amlm", amlm_update_interval=50)
add("6C", "amlm-ui500", "slow update",  masking="amlm", amlm_update_interval=500)
add("6C", "amlm-wide",  "wider range",  masking="amlm", amlm_min_ratio=0.10, amlm_max_ratio=0.60)

# 6D: Frequency masking (3)
comment("── 6D: Frequency Masking ──")
for alpha in [0.2, 0.3, 0.5]:
    add("6D", f"freq-alpha{alpha}", f"frequency masking alpha={alpha}", masking="frequency")


# ═══════════════════════════════════════════════════════════════════════
# Phase 7: Batch Size & Sequence Length (14 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 7: Batch Size & Sequence Length (14 exp) ═══")

# 7A: Effective batch size (8)
comment("── 7A: Batch Size ──")
for bs, ga, eff in BATCH_CONFIGS:
    add("7A", f"batch-{eff}", f"bs={bs} ga={ga} eff={eff}", batch_size=bs, grad_accum=ga)

# 7B: Sequence length (6)
comment("── 7B: Sequence Length ──")
for sl in SEQ_LENS:
    add("7B", f"seqlen-{sl}", f"seq_len={sl}", seq_len=sl)


# ═══════════════════════════════════════════════════════════════════════
# Phase 8: Optimizer Exploration (16 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 8: Optimizer Exploration (16 exp) ═══")

# 8A: AdamW LR sweep (6×3ep)
comment("── 8A: AdamW LR Sweep (3ep) ──")
for lr in LRS_ADAMW:
    add("8A", f"adamw-lr{lr}", f"AdamW lr={lr}", optimizer="AdamW", lr=lr, epochs=3)

# 8B: Muon LR sweep (5×3ep)
comment("── 8B: Muon LR Sweep (3ep) ──")
for lr in LRS_MUON:
    add("8B", f"muon-lr{lr}", f"Muon lr={lr}", optimizer="Muon", lr=lr, epochs=3)

# 8C: Best alternative + FORGETTER (5×10ep)
comment("── 8C: Best Optimizer + FORGETTER (10ep) ──")
add("8C", "adamw-best",     "AdamW best LR from 8A",     optimizer="AdamW", epochs=10)
add("8C", "adamw-best-fgt", "AdamW + FORGETTER",         optimizer="AdamW", forgetter=True, epochs=10)
add("8C", "muon-best",      "Muon best LR from 8B",      optimizer="Muon", epochs=10)
add("8C", "muon-best-fgt",  "Muon + FORGETTER",          optimizer="Muon", forgetter=True, epochs=10)
add("8C", "lamb-fgt-recheck", "LAMB+FGT with tuned cfg", forgetter=True, epochs=10)


# ═══════════════════════════════════════════════════════════════════════
# Phase 9: Architecture Modifications (15 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 9: Architecture Modifications (15 exp) ═══")

# 9A: Toggles (7)
comment("── 9A: Feature Toggles ──")
add("9A", "rope",           "RoPE position encoding",   use_rope=True)
add("9A", "no-geglu",       "Standard GELU FFN",        use_geglu=False)
add("9A", "post-norm",      "Post-LayerNorm",           use_pre_norm=False)
add("9A", "attn-gate",      "GELU gate on V",           use_attention_gate=True)
add("9A", "dwa",            "DenseFormer",              use_dwa=True)
add("9A", "rope-attn-gate", "RoPE + attn gate",         use_rope=True, use_attention_gate=True)
add("9A", "rope-dwa",       "RoPE + DWA",              use_rope=True, use_dwa=True)

# 9B: Model shape (8)
comment("── 9B: Model Shape ──")
for h, l, nh, inter in MODEL_SHAPES:
    add("9B", f"shape-{h}x{l}", f"h={h} l={l} nh={nh} i={inter}",
        hidden_size=h, num_layers=l, num_heads=nh, intermediate_size=inter)


# ═══════════════════════════════════════════════════════════════════════
# Phase 10: MoE Hyperparameter Tuning (8 experiments, conditional)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 10: MoE Tuning (8 exp, conditional on P5) ═══")

for ne, tk, es in MOE_CONFIGS:
    add(10, f"moe-e{ne}k{tk}s{es}", f"{ne} experts top-{tk} size={es}",
        use_moe=True, moe_num_experts=ne, moe_top_k=tk, moe_expert_size=es)
for fp in MOE_FREQ_PENALTIES:
    add(10, f"moe-fp{fp}", f"freq_penalty={fp}",
        use_moe=True, moe_freq_penalty=fp)


# ═══════════════════════════════════════════════════════════════════════
# Phase 11: Tokenizer × Embedding Cross (6 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 11: Tokenizer × Embedding × Vocab Size (9 exp) ═══")
for tok in ["bpe", "morfessor_bpe"]:
    for emb, init in [("standard", "random"), ("nhot", "random"), ("standard", "fasttext")]:
        tok_short = "bpe" if tok == "bpe" else "morf"
        add(11, f"{tok_short}-{emb}-{init}", f"tok={tok} emb={emb} init={init}",
            tokenizer=tok, embedding=emb, embedding_init=init)

# Vocab size sweep (BPE only, standard embedding)
comment("── 11B: Vocab Size Sweep ──")
for vs in [4096, 16384]:
    add("11B", f"vocab-{vs}", f"BPE vocab_size={vs}", tokenizer="bpe")


# ═══════════════════════════════════════════════════════════════════════
# Phase 12: Dataset Construction Methods (18 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 12: Dataset Construction (18 exp) ═══")
for data_id, data_desc in DATA_CONFIGS:
    add(12, f"data-{data_id}", data_desc, data=data_id)


# ═══════════════════════════════════════════════════════════════════════
# Phase 13: Cross-Architecture Verification (15 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 13: Cross-Architecture Verification (15 exp) ═══")

for arch_name in ["gpt2", "modernbert", "xlstm", "rtd"]:
    preset = ARCH_PRESETS[arch_name]
    add(13, f"{arch_name}-tuned", f"{arch_name} with tuned config", **preset)

# + MoE / AttnRes variants for non-RTD
for arch_name in ["gpt2", "modernbert", "xlstm"]:
    preset = ARCH_PRESETS[arch_name]
    add(13, f"{arch_name}-moe", f"{arch_name} + MoE", **preset, use_moe=True)

for arch_name in ["gpt2", "modernbert", "xlstm"]:
    preset = ARCH_PRESETS[arch_name]
    add(13, f"{arch_name}-attnres", f"{arch_name} + AttnRes", **preset, use_attn_res=True)

# RTD + MoE (AttnRes not supported for RTD)
add(13, "rtd-moe", "RTD + MoE", **ARCH_PRESETS["rtd"], use_moe=True)

# Wide variants
for arch_name in ["gpt2", "modernbert", "xlstm"]:
    preset = ARCH_PRESETS[arch_name]
    add(13, f"{arch_name}-wide", f"{arch_name} 512×8",
        **preset, hidden_size=512, num_layers=8, num_heads=8, intermediate_size=1706)


# ═══════════════════════════════════════════════════════════════════════
# Phase 14: LR × Dropout Interaction Grid (27 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 14: LR × Dropout Grid (24×3ep + 3×10ep) ═══")

# 14A: 6 LR × 4 dropout = 24 (3ep)
lr_mults  = [0.7, 0.85, 1.0, 1.15, 1.3, 1.5]
drop_mults = ["0.0", "half", "best", "1.5x"]  # actual values filled at runtime
for lm in lr_mults:
    for dm in drop_mults:
        add("14A", f"grid-lr{lm}x-d{dm}", f"lr×{lm} drop={dm}", epochs=3)

# 14B: Top 3 promoted (10ep)
for i in range(1, 4):
    add("14B", f"grid-promote{i}", f"Top {i} from 14A → 10ep", epochs=10)


# ═══════════════════════════════════════════════════════════════════════
# Phase 15: Miscellaneous Hyperparameters (12 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 15: Misc Hyperparameters (12 exp) ═══")

for zl in Z_LOSSES:
    add(15, f"zloss-{zl}", f"z_loss_weight={zl}", z_loss_weight=zl)

for bs in BUCKET_SIZES:
    add(15, f"bucket-{bs}", f"position_bucket_size={bs}", position_bucket_size=bs)

for b in BETAS:
    b_short = b.replace("(", "").replace(")", "").replace(",", "-")
    add(15, f"betas-{b_short}", f"betas={b}", optimizer_betas=b)


# ═══════════════════════════════════════════════════════════════════════
# Phase 16: Knowledge Distillation (14 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 16: Knowledge Distillation (14 exp) ═══")

default_teacher = KD_TEACHERS[0][1]

# Temperature sweep (4)
for t in KD_TEMPS:
    add(16, f"kd-t{t}", f"KD temp={t}", kd_enabled=True, kd_teacher=default_teacher, kd_temperature=t)

# Alpha sweep (5, fix t=4)
for a in KD_ALPHAS:
    add(16, f"kd-a{a}", f"KD alpha={a}", kd_enabled=True, kd_teacher=default_teacher, kd_alpha=a)

# Larger teachers (2)
for teacher_name, teacher_path in KD_TEACHERS[1:]:
    add(16, f"kd-{teacher_name}", f"teacher={teacher_name}",
        kd_enabled=True, kd_teacher=teacher_path)

# Interaction tests (2)
add(16, "kd-no-fgt",  "KD without FORGETTER",  kd_enabled=True, kd_teacher=default_teacher, forgetter=False)
add(16, "kd-no-amlm", "KD without AMLM",       kd_enabled=True, kd_teacher=default_teacher, masking="standard")

# Control
add(16, "no-kd-control", "No KD (control)")


# ═══════════════════════════════════════════════════════════════════════
# Phase 17: Epoch Count & Checkpoint Strategy (12 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 17: Epoch & Checkpoint (12 exp) ═══")

for ep in [3, 5, 7, 8]:
    add(17, f"epoch-{ep}", f"epochs={ep}", epochs=ep)

# Baseline 10ep already tested; test ckpt averaging
for k in [2, 3, 5, 7]:
    add(17, f"ckptavg-{k}", f"avg last {k}", ckpt_averaging=True, ckpt_avg_last_k=k)

# Early stop + averaging combos
for ep in [5, 7, 8]:
    add(17, f"epoch{ep}-avg3", f"epoch={ep} + avg last 3", epochs=ep, ckpt_averaging=True, ckpt_avg_last_k=3)


# ═══════════════════════════════════════════════════════════════════════
# Phase 18: Seed Variance (8 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 18: Seed Variance (8 exp) ═══")
for s in SEEDS:
    add(18, f"seed-{s}", f"seed={s}", seed=s)


# ═══════════════════════════════════════════════════════════════════════
# Phase 19: Data × Feature Interaction (24 experiments, 10ep)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 19: Data × Feature Interaction (24 exp) ═══")

# Top 3 data × {none, partial, full} features × {base, alt} shapes
feat_configs = {
    "fnone": {},
    "fpart": dict(masking="amlm", embedding="nhot"),
    "ffull": dict(masking="amlm", embedding="nhot", forgetter=True, use_moe=True),
}
shape_configs = {
    "sbase": {},
    "salt":  dict(hidden_size=512, num_layers=8, num_heads=8, intermediate_size=1706),
}

# 3 data × 3 features × 2 shapes = 18
for di in range(1, 4):
    for fname, foverrides in feat_configs.items():
        for sname, soverrides in shape_configs.items():
            add(19, f"d{di}-{fname}-{sname}", f"TOP_DATA_{di} {fname} {sname}",
                **foverrides, **soverrides)

# Extra variants: +Morfessor, +FastText, +KD (6)
for di in range(1, 3):
    add(19, f"d{di}-ffull-morf",    f"TOP_DATA_{di} full+Morfessor",
        masking="amlm", embedding="nhot", forgetter=True, use_moe=True, tokenizer="morfessor_bpe")
    add(19, f"d{di}-ffull-fasttext", f"TOP_DATA_{di} full+FastText",
        masking="amlm", embedding="nhot", forgetter=True, use_moe=True, embedding_init="fasttext")
    add(19, f"d{di}-ffull-kd",      f"TOP_DATA_{di} full+KD",
        masking="amlm", embedding="nhot", forgetter=True, use_moe=True,
        kd_enabled=True, kd_teacher=default_teacher)


# ═══════════════════════════════════════════════════════════════════════
# Phase 20: Final Candidates & Ablation Insurance (20 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 20: Final Candidates & Ablation (20 exp) ═══")

# 20A: Candidate variants (10)
comment("── 20A: Candidates ──")
for i in range(1, 11):
    add("20A", f"final-v{i}", f"Final candidate variant {i}")

# 20B: Ablation (10 — remove one component each)
comment("── 20B: Ablation Insurance ──")
ablations = [
    "no-amlm", "no-nhot", "no-fgt", "no-moe", "no-attnres",
    "no-kd", "no-ckptavg", "baseline-lr", "baseline-drop", "baseline-data",
]
for ab in ablations:
    add("20B", f"ablate-{ab}", f"Remove {ab} from CONFIG_FINAL")


# ═══════════════════════════════════════════════════════════════════════
# Phase 21: Submission Seeds & Full Eval (16 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 21: Submission & Full Eval (16 exp) ═══")

# CONFIG_FINAL × 5 seeds
for s in [42, 7, 123, 2026, 9999]:
    add(21, f"submit-seed{s}", f"FINAL seed={s}", seed=s)

# 2nd/3rd candidates × 3 seeds
for cand in [2, 3]:
    for s in [42, 7, 123]:
        add(21, f"submit-v{cand}-seed{s}", f"Candidate {cand} seed={s}", seed=s)

# Special runs
add(21, "submit-ensemble-avg",  "Cross-seed checkpoint average")
add(21, "submit-ckptavg3",     "Best seed ckpt avg last 3", ckpt_averaging=True, ckpt_avg_last_k=3)
add(21, "submit-ckptavg5",     "Best seed ckpt avg last 5", ckpt_averaging=True, ckpt_avg_last_k=5)
add(21, "submit-glue-ft",      "Full eval + GLUE fine-tuning")


# ═══════════════════════════════════════════════════════════════════════
# Phase 22: Reserve / Exploratory (61 experiments)
# ═══════════════════════════════════════════════════════════════════════
comment("═══ Phase 22: Reserve / Exploratory (61 exp) ═══")

# 22A: Second-architecture deep dive (10)
comment("── 22A: Alt Architecture Deep Dive ──")
for i in range(1, 11):
    add("22A", f"alt-arch-{i}", f"Alt architecture tuning {i}")

# 22B: Augmentation ratio tuning (10)
comment("── 22B: Augmentation Ratio Tuning ──")
for pct in [10, 15, 20, 30, 40, 50]:
    add("22B", f"aug-para{pct}", f"Paraphrase {pct}% mix")
for pct in [25, 50, 75]:
    add("22B", f"aug-recombi{pct}", f"RecombiText {pct}%")
add("22B", "aug-para-recombi", "Para 20% + RecombiText 30%")

# 22C: MoE alternatives (5)
comment("── 22C: MoE Alternatives ──")
for i in range(1, 6):
    add("22C", f"moe-alt-{i}", f"MoE alternative config {i}")

# 22D: Interpolation fine-tuning (10)
comment("── 22D: Parameter Interpolation ──")
for i in range(1, 11):
    add("22D", f"interp-{i}", f"Interpolation between best configs {i}")

# 22E: Debug / re-runs (6)
comment("── 22E: Debug / Re-runs ──")
for i in range(1, 7):
    add("22E", f"debug-{i}", f"Debug slot {i}")

# 22F: Late-breaking ideas (20)
comment("── 22F: Exploratory ──")
for i in range(1, 21):
    add("22F", f"explore-{i}", f"Exploratory slot {i}")


# ═══════════════════════════════════════════════════════════════════════
# Write CSV
# ═══════════════════════════════════════════════════════════════════════

def write_csv():
    with open(OUTPUT, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=COLUMNS, extrasaction="ignore")
        writer.writeheader()
        for exp in experiments:
            if "_comment" in exp:
                row = {col: "" for col in COLUMNS}
                row["id"] = f"## {exp['_comment']}"
                writer.writerow(row)
            else:
                writer.writerow(exp)

    n = sum(1 for e in experiments if "_comment" not in e)
    print(f"Generated {n} experiments across {len(phase_counter)} phases → {OUTPUT}")
    print()
    for phase, count in sorted(phase_counter.items(), key=lambda x: (len(str(x[0])), str(x[0]))):
        print(f"  Phase {phase:>4s}: {count:>3d} experiments")
    print(f"  {'TOTAL':>10s}: {n:>3d} experiments")
    print()
    # Estimate: 3ep ~15min, 10ep ~40min
    quick = sum(1 for e in experiments if "_comment" not in e and e.get("epochs") == 3)
    full  = n - quick
    hours = (quick * 15 + full * 40) / 60
    print(f"  Quick (3ep): {quick} experiments × ~15 min")
    print(f"  Full (10ep): {full} experiments × ~40 min")
    print(f"  Estimated total: ~{hours:.0f} GPU-hours")
    print(f"  With 3 GPUs parallel: ~{hours/3:.0f} hours wall time ({hours/3/24:.1f} days)")


if __name__ == "__main__":
    write_csv()