"""Generate Compilertetris_GRPO_synthetic_dataset.ipynb (GRPO + random Toy-IR corpus)."""
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "Compilertetris_GRPO_synthetic_dataset.ipynb"


def cell_md(s: str) -> dict:
    return {
        "cell_type": "markdown",
        "metadata": {},
        "source": [line + "\n" for line in s.strip().split("\n")],
    }


def cell_code(s: str) -> dict:
    return {
        "cell_type": "code",
        "metadata": {},
        "execution_count": None,
        "outputs": [],
        "source": [line + "\n" for line in s.rstrip().split("\n")],
    }

cells: list = []

cells.append(cell_md("""
# Compiler Tetris — GRPO, **synthetic Toy-IR** (Colab **T4 ~8–10 h** preset)

**Runtime (estimate):** after ~20 steps, read `it/s` in the progress bar. Approximate hours ≈ `TRAIN_STEPS / (it/s * 3600)`.
With **`TRAIN_STEPS = 3600`**: at **0.10 it/s** that is about **10 h**; at **0.12 it/s** about **8.3 h**. If you finish much faster/slower, change **`TRAIN_STEPS`** (or the model size / `NUM_GENERATIONS`).

**T4 tips:** 3B + 4bit + below settings fit T4; **7B** or very long `MAX_COMPLETION` can OOM. **High-RAM** runtime in Colab helps. Use **checkpoints** (next cells) in case the session dies.

| What | Value |
|------|--------|
| Space | `greedybeserk95/Compilertetris` |
| Code root | `REPO_DIR` (printed in the “paths” cell) |
| LoRA out | `LORA_DIR` (printed in the “paths” cell) |
| **Checkpoints** | `GRPO_OUTPUT_DIR + "/checkpoint-*"` |

`program_generator` — Toy-IR in `runtime_core` shape; `metahack1` uses a different schema.

**Training evidence (loss + reward plots):** the cell *after* `trainer.train()` saves `training_loss_and_reward.png` under `output_dir` and shows it in the notebook — use it in your README / writeup.
"""))

cells.append(
    cell_md(
        """
## Checkpoints

The GRPO config uses `save_strategy="steps"` and `save_total_limit` so training writes **periodic checkpoints** under `output_dir` (e.g. `$GRPO_OUTPUT_DIR/checkpoint-20`, …) and prunes old ones. After a crash, re-run the setup cells, rebuild `trainer`, then use `resume_from_checkpoint=True` (latest) or a **specific path** (see the cell after training).
"""
    )
)

cells.append(cell_code("""
# --- Central path config (works on Colab *and* local / HF runtimes) ---

import os
import tempfile
from pathlib import Path

def _default_workspace_base() -> str:
    # Colab: /content is writable
    c = "/content"
    if os.path.isdir(c) and os.access(c, os.W_OK):
        return str(Path(c) / "work")

    # Otherwise: a guaranteed-writable temp dir
    return str(Path(tempfile.gettempdir()) / "compilertetris_work")


BASE = os.environ.get("COMPILERTETRIS_BASE", _default_workspace_base())
os.makedirs(BASE, exist_ok=True)

HF_SPACE_REPO = "https://huggingface.co/spaces/greedybeserk95/Compilertetris"
HF_ADAPTER_REPO = "greedybeserk95/Compilertetris-grpo-adapter"
REPO_DIR = str(Path(BASE) / "Compilertetris")
GRPO_OUTPUT_DIR = str(Path(BASE) / "grpo_compilertetris")
LORA_DIR = str(Path(BASE) / "compilertetris_lora")

print("BASE        :", BASE)
print("REPO_DIR    :", REPO_DIR)
print("GRPO_OUTPUT :", GRPO_OUTPUT_DIR)
print("LORA_DIR    :", LORA_DIR)
"""))

cells.append(cell_code("""
import os, subprocess
from pathlib import Path

# Ensure parent is writable/created
Path(REPO_DIR).parent.mkdir(parents=True, exist_ok=True)

if os.path.isdir(REPO_DIR + "/.git"):
    subprocess.check_call(["git", "-C", REPO_DIR, "pull", "--ff-only"], timeout=300)
else:
    subprocess.check_call(["git", "clone", HF_SPACE_REPO, REPO_DIR], timeout=600)
print("Repo:", REPO_DIR)
"""))

cells.append(cell_code("""!nvidia-smi"""))

cells.append(
    cell_md(
        """
## Install (2 cells)

Unsloth’s dependency graph is finicky. The reliable pattern is:

1) **Install official CUDA `torch` first** (so `flash-attn` / other builds can `import torch` during install).
2) **Install the HF training stack + Unsloth** with `--no-build-isolation` (avoids the common `flash-attn` `egg_info` failure).

After the install cells, **restart the kernel** before importing.
"""
    )
)

cells.append(cell_code("""
# --- Install (1/2): CUDA PyTorch (must be importable before flash-attn / Unsloth extras) ---

!python -m pip uninstall -y torch torchvision torchaudio triton xformers || true
!python -m pip install -U --no-cache-dir pip setuptools wheel packaging ninja
!python -m pip install -U --no-cache-dir --force-reinstall \
  torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

import torch
from torch.utils import _pytree
print("torch:", torch.__version__, torch.__file__)
print("pytree register_constant present:", hasattr(_pytree, "register_constant"))
"""))

cells.append(cell_code("""
# --- Install (2/2): HF stack + unsloth_zoo + Unsloth (A100 uses the Ampere + cu124 + torch2.6 extra) ---

!python -m pip install -U --no-cache-dir --no-build-isolation \
  "transformers==4.55.4" "trl==0.20.0" "datasets==4.3.0" \
  "huggingface-hub==0.30.0" \
  accelerate peft bitsandbytes safetensors sentencepiece \
  matplotlib \
  "unsloth_zoo==2026.4.9" \
  "unsloth[cu124-ampere-torch260]"

import importlib.util
for pkg in ("unsloth", "unsloth_zoo", "trl", "transformers"):
    print(pkg, "OK" if importlib.util.find_spec(pkg) else "MISSING")
"""))

cells.append(cell_code("""
# Keep `import unsloth` on its own line (don't bundle it with `import sys, json`).
import unsloth

import os, sys, json
sys.path.insert(0, REPO_DIR)

import torch
from unsloth import FastLanguageModel
from datasets import Dataset

from runtime_core import (
    CompilerOptimizationEnv,
    Deliverable2_Formatter,
    MockEngine,
    MOCK_PASSES,
)
from program_generator import build_training_program_corpus

print("torch", torch.__version__, "cuda", torch.cuda.is_available())
print("MOCK_PASSES", list(MOCK_PASSES.keys()))
"""))


def _cell_chatml_synth() -> str:
    p0 = """# Qwen2.5 ChatML + synthetic PROGRAMS
IM_END = \"<\" + \"|im_end|>\"
"""
    p1 = r'''
import re

# T4: 200–300 is a good tradeoff (RAM + diversity)
N_TRAIN_PROGRAMS = 250
RANDOM_SEED = 42

SYSTEM_PROMPT = f"""You are a compiler optimization agent.
You will see Toy-IR as pseudo-assembly (Deliverable-2 text).

Output ONLY a valid JSON array of optimization pass names IN ORDER.
Available passes: {", ".join(repr(p) for p in sorted(MOCK_PASSES))}
Rules: JSON array only; max 8 passes; you may repeat passes."""


def build_prompt(program_list: list, program_id: int) -> str:
    obs = Deliverable2_Formatter.translate_state(program_list)
    return (
        f"<|im_start|>system\n{SYSTEM_PROMPT}{IM_END}\n"
        f"<|im_start|>user\n#PROGRAM_ID:{program_id}\n{obs}{IM_END}\n"
        f"<|im_start|>assistant\n"
    )


def program_from_prompt(prompt: str) -> list:
    m = re.search(r"#PROGRAM_ID:(\d+)", prompt)
    if m:
        i = int(m.group(1))
        if 0 <= i < len(PROGRAMS):
            return PROGRAMS[i]
    if "#PROGRAM_JSON" in prompt:
        tail = prompt.split("#PROGRAM_JSON", 1)[1]
        if IM_END in tail:
            tail = tail.split(IM_END, 1)[0]
        raw = tail.strip()
        if raw:
            return json.loads(raw)
    raise ValueError("cannot resolve program from prompt (expect #PROGRAM_ID:N or #PROGRAM_JSON)")


PROGRAMS = build_training_program_corpus(
    n_total=N_TRAIN_PROGRAMS,
    seed=RANDOM_SEED,
    include_builtins=True,
)
print("Corpus size:", len(PROGRAMS))

train_dataset = Dataset.from_dict({
    "prompt": [build_prompt(p, i) for i, p in enumerate(PROGRAMS)],
})
print("Dataset rows:", len(train_dataset))
'''.lstrip("\n")
    return p0 + p1


cells.append(cell_code(_cell_chatml_synth()))

cells.append(cell_code("""
def env_reward_for_completion(prompt: str, completion: str, max_env_steps: int = 8) -> float:
    try:
        program = program_from_prompt(prompt)
    except Exception:
        return -8.0
    try:
        actions = Deliverable2_Formatter.extract_action_array(completion)
    except Exception:
        return -5.0
    actions = [str(a).strip() for a in actions][: max_env_steps]
    if not actions:
        return -4.0
    env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=max_env_steps)
    env.reset(program)
    total = 0.0
    for a in actions:
        step = env.step(a)
        total += float(step.reward)
        if step.done:
            break
    return float(total)


def make_reward_function(max_env_steps: int = 10):
    def reward_func(prompts: list, completions: list, **kwargs) -> list:
        return [env_reward_for_completion(p, c, max_env_steps) for p, c in zip(prompts, completions)]
    return reward_func
"""))

cells.append(cell_code("""
from trl import GRPOConfig, GRPOTrainer

# --- T4-friendly (3B 4-bit). For A100+ you can try Qwen2.5-7B and MAX_COMPLETION=384. ---
MODEL_NAME = "unsloth/Qwen2.5-3B-Instruct"
MAX_SEQ_LEN = 1024
MAX_COMPLETION = 256

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=MAX_SEQ_LEN,
    dtype=None,
    load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
    model,
    r=32, lora_alpha=32, lora_dropout=0.0, bias="none",
    use_gradient_checkpointing="unsloth", random_state=0,
)
print("Model:", MODEL_NAME, "| max_seq", MAX_SEQ_LEN, "| completion cap", MAX_COMPLETION)
"""))

# ---- T4 / ~8–10 h wall time: tune TRAIN_STEPS after you see it/s in the first minutes ----
cells.append(
    cell_code(
        """
# 4 rollouts per prompt: good for GRPO on T4; 6–8 is heavier (slower, more VRAM)
NUM_GENERATIONS = 4
LEARNING_RATE = 1.5e-5
# Target ~8–10 h on T4 when it/s is ~0.10–0.12 (typical for this stack). Re-tune if your it/s differs.
#  hours ≈ TRAIN_STEPS / (it/s * 3600)
TRAIN_STEPS = 3600
GRAD_ACCUM = 2

# GRPO_OUTPUT_DIR is set in the "Central path config" cell
CHECKPOINT_EVERY = 200
KEEP_LAST_N_CHECKPOINTS = 3

grpo_config = GRPOConfig(
    output_dir=GRPO_OUTPUT_DIR,
    learning_rate=LEARNING_RATE,
    per_device_train_batch_size=NUM_GENERATIONS,
    gradient_accumulation_steps=GRAD_ACCUM,
    num_generations=NUM_GENERATIONS,
    max_completion_length=MAX_COMPLETION,
    max_prompt_length=MAX_SEQ_LEN,
    remove_unused_columns=False,
    temperature=0.7,
    max_steps=TRAIN_STEPS,
    logging_steps=20,
    save_strategy="steps",
    save_steps=CHECKPOINT_EVERY,
    save_total_limit=KEEP_LAST_N_CHECKPOINTS,
    seed=0,
    report_to="none",
    use_vllm=False,
)

reward_fn = make_reward_function()

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[reward_fn],
    args=grpo_config,
    train_dataset=train_dataset,
)
"""
    )
)

cells.append(cell_code("""
print("Starting GRPO (synthetic corpus)…")
trainer.train()
print("Done.")
"""))

cells.append(
    cell_code(
        r"""
# --- Loss + reward plots (submission): from trainer.state.log_history after train() ---

!pip install -q matplotlib

import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

PLOT_DIR = GRPO_OUTPUT_DIR
os.makedirs(PLOT_DIR, exist_ok=True)
PNG_PATH = os.path.join(PLOT_DIR, "training_loss_and_reward.png")

def _reward_from_log_row(h):
    v = h.get("reward")
    if isinstance(v, (int, float)):
        return float(v), "reward"
    best_k, best_v = None, None
    for k, v in h.items():
        if not isinstance(v, (int, float)):
            continue
        klow = k.replace("-", "/").lower()
        if "reward" not in klow or "std" in klow or "clip" in klow or "per_token" in klow:
            continue
        if "mean" in klow or k == "reward":
            return float(v), k
        if "mean" not in klow and best_k is None:
            best_k, best_v = k, float(v)
    if best_k is not None:
        return best_v, best_k
    return None, None

def extract_series(history):
    sl, vl, sr, vr = [], [], [], []
    rlabel = None
    for h in history:
        s = h.get("step")
        if s is None:
            continue
        lo = h.get("loss")
        if isinstance(lo, (int, float)):
            sl.append(s)
            vl.append(float(lo))
        r_val, rk = _reward_from_log_row(h)
        if r_val is not None and rk:
            if rlabel is None:
                rlabel = rk
            if rk == rlabel:
                sr.append(s)
                vr.append(r_val)
    return (sl, vl, "loss"), (sr, vr, rlabel or "reward")

(loss_s, loss_v, _lk), (rew_s, rew_v, rew_lab) = extract_series(trainer.state.log_history)
print("Points — loss:", len(loss_v), "| reward:", len(rew_v))
if trainer.state.log_history:
    print("Last log row keys (sample):", list(trainer.state.log_history[-1].keys())[:25])

fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
if loss_v:
    ax0.plot(loss_s, loss_v, "b.-", label="loss", linewidth=1, markersize=2)
    ax0.set_ylabel("training loss")
    ax0.set_title("GRPO (this Colab run)")
    ax0.grid(True, alpha=0.3)
    ax0.legend()
else:
    ax0.text(0.5, 0.5, "No 'loss' in log_history", ha="center", transform=ax0.transAxes)

if rew_v:
    ax1.plot(rew_s, rew_v, "g.-", label=rew_lab, linewidth=1, markersize=2)
    ax1.set_ylabel("mean reward" if "mean" in (rew_lab or "") else "reward")
    ax1.set_xlabel("global step")
    ax1.grid(True, alpha=0.3)
    ax1.legend()
else:
    ax1.text(0.5, 0.5, "No reward column found — see keys above", ha="center", transform=ax1.transAxes)
    ax1.set_xlabel("global step")

plt.tight_layout()
plt.savefig(PNG_PATH, dpi=150, bbox_inches="tight")
print("Saved:", PNG_PATH)
try:
    from IPython.display import Image, display
    display(Image(PNG_PATH))
except Exception as e:
    print("Display:", e)
finally:
    plt.close("all")
print("If keys differ:  print(trainer.state.log_history[-1])")
"""
    )
)

cells.append(
    cell_code(
        """
# List on-disk checkpoints (for resume or manual export)
import glob, os

ckpts = sorted(
    glob.glob(os.path.join(GRPO_OUTPUT_DIR, "checkpoint-*")),
    key=lambda p: int(p.split("checkpoint-")[-1]) if p.split("checkpoint-")[-1].isdigit() else 0,
)
print("Checkpoints in", GRPO_OUTPUT_DIR, ":", len(ckpts))
for c in ckpts:
    print(" ", c)
if ckpts:
    print("Latest:", ckpts[-1])
"""
    )
)

cells.append(
    cell_md(
        """
**Resume after disconnect / crash** — re-run: clone, pip, imports, dataset, reward, **model load**, and **trainer** cells. Then run **one** of:

- `trainer.train(resume_from_checkpoint=True)` — continues from the latest `checkpoint-*` in `output_dir`
- `trainer.train(resume_from_checkpoint="<your GRPO_OUTPUT_DIR>/checkpoint-300")` — example (use an existing `checkpoint-*` folder; print `GRPO_OUTPUT_DIR` from the path cell)
"""
    )
)

cells.append(
    cell_code(
        """
# Uncomment to resume from the latest checkpoint (run after re-creating `trainer` in a new session)
# trainer.train(resume_from_checkpoint=True)
"""
    )
)

cells.append(
    cell_code(
        """
SAVE_DIR = LORA_DIR
trainer.model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)
print("Saved to", SAVE_DIR)
"""
    )
)

cells.append(
    cell_code(
        """
from huggingface_hub import login, HfApi

login()
HfApi().create_repo(HF_ADAPTER_REPO, exist_ok=True, repo_type="model")
trainer.model.push_to_hub(HF_ADAPTER_REPO, private=True)
tokenizer.push_to_hub(HF_ADAPTER_REPO, private=True)
print("Pushed to https://huggingface.co/" + HF_ADAPTER_REPO)
"""
    )
)

if __name__ == "__main__":
    nb = {
        "nbformat": 4,
        "nbformat_minor": 5,
        "metadata": {
            "colab": {"provenance": [], "gpuType": "T4"},
            "kernelspec": {
                "display_name": "Python 3",
                "language": "python",
                "name": "python3",
            },
            "language_info": {"name": "python"},
        },
        "cells": cells,
    }
    OUT.write_text(json.dumps(nb, indent=2), encoding="utf-8")
    print("Wrote", OUT)