#!/usr/bin/env python3
"""Write L15 diag configs with SEPARATE promote vs backtrack gates.

Naming: diag_L15_prom{F|CE}{thr}_bt{F|CE|NONE}{thr}
  promF095  = promote on hop-accuracy (frontier) >= 0.95
  btCE090   = backtrack on both-arms score (ce_score) < 0.90
  btNONE    = no backtracking
"""
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1] / "args"

TEMPLATE = """\
# L15 DIAG — SEPARATE gates
#   PROMOTE   (stage i -> i+1): {promote_metric} >= {promote_threshold}
#   BACKTRACK (retrain earlier): {bt_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: 15
pad_latent_to_max: True

accuracy_staging: True
init_stage: 0
promote_metric: {promote_metric}
promote_threshold: {promote_threshold}
promote_on_current_only: {promote_on_current_only}
epochs_per_stage: 25

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

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

backprop_depth: null
train_size: 50000
save_only_improve: False
save_every: 200
uniform_prob: 0.1
model_id: configs/symbol-2layer-8head-768dim-L20.json
load_model_path: ckpts/star-coconut-L15-bfs-stage0-warm/checkpoint_99
seed: 0
resume: 0
bf16: True
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: 3000
lr: !!float "1e-4"
grad_clip: !!float "1.0"
warmup_steps: 200
weight_decay: 0.01
bfs_variant: True
"""

# (basename, promote_metric, promote_thr, bt_on, bt_metric, bt_thr, prom_cur_only)
ARMS = [
    ("diag_L15_promF095_btCE050", "frontier", 0.95, True, "ce_score", 0.50, False),
    ("diag_L15_promF095_btCE090", "frontier", 0.95, True, "ce_score", 0.90, False),
    ("diag_L15_promF095_btCE095", "frontier", 0.95, True, "ce_score", 0.95, False),
    ("diag_L15_promF099_btCE090", "frontier", 0.99, True, "ce_score", 0.90, False),
    ("diag_L15_promF095_btF095", "frontier", 0.95, True, "frontier", 0.95, False),
    ("diag_L15_promCE090_btCE090", "ce_score", 0.90, True, "ce_score", 0.90, False),
    ("diag_L15_promF095_btNONE", "frontier", 0.95, False, "frontier", 0.95, True),
    ("diag_L15_promF099_btCE050", "frontier", 0.99, True, "ce_score", 0.50, False),
]


def main():
    for name, pm, pt, bt, bm, bt_thr, cur_only in ARMS:
        bt_desc = "OFF" if not bt else f"{bm} < {bt_thr} triggers retrain"
        text = TEMPLATE.format(
            name=name,
            promote_metric=pm,
            promote_threshold=pt,
            promote_on_current_only=str(cur_only),
            backtrack=str(bt),
            backtrack_metric=bm,
            backtrack_detect_threshold=bt_thr,
            bt_desc=bt_desc,
        )
        path = ROOT / f"{name}.yaml"
        path.write_text(text)
        print("wrote", path.name)


if __name__ == "__main__":
    main()