import marimo

__generated_with = "0.24.0"
app = marimo.App(width="medium", auto_download=["html"])


@app.cell
def _():
    # ===================================================================== #
    #  DragonCode LLM Family — production notebook (thin wrapper)
    #  (c) 2026 Dragon Limited. All rights reserved.
    #
    #  This notebook is a THIN LAUNCHER over scripts/run_dragoncode.py.
    #  The single source of truth for training logic lives in
    #  ~/DragonCode/scripts/*.py  (NOT duplicated in cells).
    #
    #  Industrial-standard guarantees enforced by the scripts:
    #   * 100% code-domain data — codeparrot/codeparrot-clean +
    #     open-r1/codeforces-cots (permissive licenses only). NO generic web.
    #   * 4 training bugs fixed: grad-accum dead-loop (local_steps counter),
    #     lr=0 (token-progress schedule), bytes JSON serialization (base64),
    #     allow_duplicate_filename (removed for hf_hub 1.24.0).
    #   * resume-aware + idempotent stage markers (never restart from zero).
    # ===================================================================== #
    import os, sys, subprocess, json, time

    SCRIPT_DIR = os.path.expanduser("~/DragonCode/scripts")
    CONFIG_DIR = os.path.expanduser("~/DragonCode/configs")
    LOG_DIR = os.path.expanduser("~/DragonCode/logs")

    # Auth: the scripts force the "dragonlimited" account token internally, so
    # we only need to make sure the box has *some* HF_TOKEN exported.
    os.environ.setdefault("HF_TOKEN", os.environ.get("HF_TOKEN", ""))
    os.environ["HF_HOME"] = os.path.expanduser("~/.cache/huggingface")

    # 5-model tier order (Chinchilla-optimal, 20 tokens/param).
    TIER_ORDER = ["150M", "387M", "787M", "1.2B", "2.4B"]
    CHINCHILLA = {
        "150M": 3_000_000_000,
        "387M": 7_740_000_000,
        "787M": 15_740_000_000,
        "1.2B": 24_000_000_000,
        "2.4B": 48_000_000_000,
    }

    return (SCRIPT_DIR, CONFIG_DIR, LOG_DIR, TIER_ORDER, CHINCHILLA, os, subprocess, time)


@app.cell
def _(SCRIPT_DIR, CONFIG_DIR, LOG_DIR, os, subprocess, time):
    # ===================================================================== #
    #  Stage runner — delegates every stage to scripts/run_dragoncode.py
    #  (the single source of truth). Stages per tier (domain-only 12-step):
    #    pretrain → sft → dpo → golf → merge → verify → gguf
    #  150M/387M: pretrain only. 2.4B: no DPO this cycle.
    # ===================================================================== #
    TIER_STAGES = {
        "150M": ["pretrain"],
        "387M": ["pretrain"],
        "787M": ["pretrain", "sft", "dpo", "golf", "merge", "verify", "gguf"],
        "1.2B": ["pretrain", "sft", "dpo", "golf", "merge", "verify", "gguf"],
        "2.4B": ["pretrain", "sft", "golf", "merge", "verify", "gguf"],
    }

    def _run(args, logfile):
        """Run a CLI stage, streaming stdout to its own log file (tail-friendly)."""
        os.makedirs(LOG_DIR, exist_ok=True)
        with open(logfile, "a") as lf:
            lf.write(f"\n=== {time.strftime('%Y-%m-%dT%H:%M:%SZ')} {' '.join(args)} ===\n")
            lf.flush()
            proc = subprocess.Popen(
                args, cwd=SCRIPT_DIR,
                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                text=True, bufsize=1,
            )
            assert proc.stdout is not None
            for line in proc.stdout:
                lf.write(line); lf.flush()
                print(line, end="", flush=True)
            rc = proc.wait()
        return rc

    def run_stage(tier, stage):
        script = f"dragoncode_{stage}.py"
        config = os.path.join(CONFIG_DIR, f"DragonCode-{tier}.yaml")
        cmd = [sys.executable, script, "--tier", tier, "--config", config]
        logfile = os.path.join(LOG_DIR, f"DragonCode-{tier}-{stage}.log")
        print(f"\n[DRIVE] {tier}/{stage} -> {' '.join(cmd)}", flush=True)
        rc = _run(cmd, logfile)
        if rc != 0:
            print(f"[DRIVE] {tier}/{stage} FAILED rc={rc} (see {logfile})", flush=True)
            return False
        print(f"[DRIVE] {tier}/{stage} OK", flush=True)
        return True

    return (TIER_STAGES, run_stage)


@app.cell
def _(TIER_ORDER, TIER_STAGES, run_stage, time):
    # ===================================================================== #
    #  Full pipeline — sequential, resume-safe, single epoch per tier.
    #  Unique stop condition: user interrupt. No early-exit, no auto-exit.
    # ===================================================================== #
    def run_all():
        for tier in TIER_ORDER:
            for stage in TIER_STAGES[tier]:
                ok = run_stage(tier, stage)
                attempt = 0
                while not ok and attempt < 3:
                    attempt += 1
                    time.sleep(8)
                    print(f"[DRIVE] {tier}/{stage} retry {attempt}/3", flush=True)
                    ok = run_stage(tier, stage)
                if not ok:
                    print(f"[DRIVE] {tier}/{stage} failed after retries — stopping pipeline", flush=True)
                    return
        print("[DRIVE] All 5 models complete.", flush=True)


    return (run_all,)


@app.cell
def _(run_all):
    # ===================================================================== #
    #  LAUNCH — press Run (▶) on this cell to train all 5 models end to end.
    # ===================================================================== #
    run_all()

    return


if __name__ == "__main__":
    app.run()