#!/usr/bin/env python3
"""Generate the L15 "make it work" sweep.

Diagnosis these configs are built from (see logs/diag_L15_*.log):

  * Stage 0 is NOT the bottleneck. hop-1 sits at frontier 1.00 / ce_score 0.94-0.98
    for the entire run and never degrades. The pinned stage-0 run saturates at
    ce_score 0.945 by epoch 45 and only reaches 0.961 by epoch 150.
  * The bottleneck is that ce_score decays monotonically with depth
    (0.96 -> 0.44 by hop 15) and nothing in the live arms pushed back:
    btNONE had no backtracking, btCE050 fires below 0.50 which only hop 15
    ever reaches. Promotion on frontier@0.95 keeps advancing while balance rots.
  * The L10 arm that worked used ce_score on BOTH gates at 0.95 and ended flat
    (0.96-0.98 at every one of its 10 hops). No L15 arm has used that recipe.

Three differences from the L10 winner are therefore corrected here: both gates on
ce_score@0.95, bf16 off (the balance score compares near-equal log-probabilities,
where bf16's ~8-bit mantissa is a real noise source), and the stage-0 warm start
points at checkpoint_150 instead of the stale checkpoint_99.

The 4-layer arms warm-start by layer expansion: a strict=False load puts the
2-layer checkpoint's embeddings, lm_head, ln_f and blocks h.0/h.1 into the
4-layer model and leaves h.2/h.3 random. The only existing 4-layer checkpoint
(L12 stage-0) is unusable -- it never left chance at hop-1 frontier 0.027.
"""
import os

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ARGS = os.path.join(BASE, "args")

STAGE0_CKPT = "ckpts/star-coconut-L15-bfs-stage0-warm/checkpoint_150"
M2L = "configs/symbol-2layer-8head-768dim-L20.json"
M4L = "configs/symbol-4layer-8head-768dim-L20.json"

TEMPLATE = """# {desc}
project: coconut
save_path: ckpts
name: "{name}"

only_eval: False
coconut: True
cot: False
no_thoughts: False
no_cot: False

c_thought: 1
max_latent_stage: {max_stage}
pad_latent_to_max: True

accuracy_staging: True
init_stage: {init_stage}
promote_metric: ce_score
promote_threshold: {thr}
promote_on_current_only: False
epochs_per_stage: 25

backtrack: {backtrack}
backtrack_metric: ce_score
backtrack_detect_threshold: {thr}
remember_rate: 0.3
revert_staging: False

eval_every: 5
log_every: 5
perhop_val_samples: 256
perhop_train_samples: 64
eval_print_full: False

backprop_depth: null
train_size: {train_size}
save_only_improve: False
save_every: 200
uniform_prob: 0.1
model_id: {model_id}
load_model_path: {ckpt}
seed: 0
resume: 0
bf16: False
train_path: data/star_2arm_L15_train_fo_bfs.json
val_path: data/star_2arm_L15_valid_fo_bfs.json
reset_optimizer: False
batch_size_training: 128
debug: False
gradient_accumulation_steps: 1
num_epochs: 4000
lr: !!float "1e-4"
grad_clip: !!float "1.0"
warmup_steps: 200
weight_decay: 0.01
bfs_variant: True
"""

# (name, model, ce threshold, train_size, desc)
ARMS = [
    ("L15_push_2L_ce95_50k", M2L, 0.95, 50000,
     "Exact L10-winner recipe at L15: 2-layer, ce_score@0.95 on both gates, fp32, 50k."),
    ("L15_push_2L_ce95_100k", M2L, 0.95, 0,
     "Data lever: same as the mirror arm but all 100k graphs."),
    ("L15_push_4L_ce95_50k", M4L, 0.95, 50000,
     "Capacity lever: 4-layer via layer-expansion warm start, ce_score@0.95, 50k."),
    ("L15_push_4L_ce90_50k", M4L, 0.90, 50000,
     "Capacity + looser gate: 4-layer, ce_score@0.90, 50k."),
]

written = []
for name, model_id, thr, train_size, desc in ARMS:
    body = TEMPLATE.format(desc=desc, name=name, max_stage=15, init_stage=1,
                           thr=thr, backtrack="True", train_size=train_size,
                           model_id=model_id, ckpt=STAGE0_CKPT)
    path = os.path.join(ARGS, name + ".yaml")
    open(path, "w").write(body)
    written.append(path)

# 4-layer stage-0 builder: insurance in case layer expansion does not hold.
# Pinned at stage 0 (max_latent_stage == init_stage == 0 so promote can never fire).
s0 = TEMPLATE.format(
    desc="4-layer L15 stage-0, pinned. Produces a native 4-layer stage-0 checkpoint "
         "in case layer-expansion warm start degrades hop-1.",
    name="L15_push_s0_4L", max_stage=0, init_stage=0, thr=0.95,
    backtrack="False", train_size=0, model_id=M4L, ckpt=STAGE0_CKPT)
path = os.path.join(ARGS, "L15_push_s0_4L.yaml")
open(path, "w").write(s0)
written.append(path)

for p in written:
    print("wrote", os.path.relpath(p, BASE))