"""Build the judge-friendly Colab notebook for FATHOM.

Designed to run end-to-end on a free Colab CPU runtime in ~3 minutes.
Heavy GPU work is in optional cells that judges can skip.
"""
from __future__ import annotations
import json
from pathlib import Path

NB = {
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {"name": "python3", "display_name": "Python 3"},
        "language_info": {"name": "python"},
        "colab": {"provenance": [], "toc_visible": True},
    },
    "cells": [],
}


def md(src: str) -> dict:
    return {"cell_type": "markdown", "metadata": {}, "source": src.splitlines(keepends=True)}


def code(src: str) -> dict:
    return {
        "cell_type": "code",
        "metadata": {},
        "execution_count": None,
        "outputs": [],
        "source": src.splitlines(keepends=True),
    }


# Cell 1 — title + abstract
NB["cells"].append(md("""\
# FATHOM — Judge Reproducer Notebook

> **The first RL-trained Recursive Language Model.**
> An OpenEnv environment + GRPO training pipeline that teaches Qwen 2.5 Coder 1.5B (4-bit + LoRA) to use a recursive-LM scaffold (Python REPL + recursive `llm()` calls) for long-context QA.
>
> Submission for the **Meta × PyTorch × Hugging Face OpenEnv Hackathon Grand Finale** (Bangalore, April 25–26 2026).

## What this notebook does (3 min on free CPU Colab)

1. Pings the live env Space and confirms it returns `{"status":"ok"}`.
2. Installs ~5 lightweight Python packages (no PyTorch, no bitsandbytes).
3. Downloads the reward code (~30 KB) and runs the **8 adversarial reward probes** locally to prove the verifier blocks each hack.
4. Pulls the actual training plots from the trained-model repo and renders them inline.
5. Embeds the live W&B run with the full GRPO reward trajectory (0.15 → 0.98 over 70 steps).
6. Lists the trained model's adapters + merged checkpoint on the HF Hub.

## What this notebook does NOT do

- Re-train the model. Training the full Qwen 1.5B + LoRA on the FATHOM env requires an A100 (the cell at the bottom of this notebook contains the exact command, but free Colab does not have the GPU). Re-training takes ~50 min and ~$10 on `hf jobs run --flavor=a100-large`.
- Load the merged 1.5B model on CPU. The model is published at `Pratham-math/fathom-1.5b-grpo` for anyone who has GPU compute.

## Submission links

| Resource | URL |
|---|---|
| Environment Space | <https://huggingface.co/spaces/Pratham-math/fathom-env> |
| Trained model | <https://huggingface.co/Pratham-math/fathom-1.5b-grpo> |
| Code repo | <https://huggingface.co/Pratham-math/fathom-code> |
| Mini-blog | <https://huggingface.co/spaces/Pratham-math/fathom-blog> |
| W&B run (v2 — successful) | <https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/y82wmj4x> |
"""))

# Cell 2 — header
NB["cells"].append(md("## 1 · Is the live OpenEnv server actually running?\n"))

# Cell 3 — env health check
NB["cells"].append(code("""\
import urllib.request, json, sys
ENV_URL = "https://Pratham-math-fathom-env.hf.space"

def http_get(path: str, timeout: int = 30) -> tuple[int, str]:
    req = urllib.request.Request(ENV_URL + path, headers={"User-Agent": "fathom-judge-notebook"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, r.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode("utf-8", errors="replace")
    except Exception as e:
        return 0, f"network error: {e}"

for path in ["/healthz", "/openapi.json", "/"]:
    status, body = http_get(path)
    snippet = body[:140].replace("\\n", " ")
    print(f"GET {path:18s}  -> HTTP {status}  | {snippet}")

status, body = http_get("/healthz")
assert status == 200 and "ok" in body, f"Env Space returned {status}: {body[:200]}"
print()
print("PASS — live env Space is healthy.")
"""))

# Cell 4 — header
NB["cells"].append(md("""\
## 2 · Install lightweight deps (~30 s, no GPU)

We only need three things to verify the project: an HTTP client, the HF Hub client to pull artifacts, and Pillow + matplotlib to render training plots inline.

**No PyTorch / Transformers / bitsandbytes / Unsloth in this path** — those are only needed for the optional A100 training cell at the very bottom.
"""))

# Cell 5 — pip install
NB["cells"].append(code("""\
%pip install --quiet "huggingface_hub>=0.28" pillow matplotlib requests
print("deps OK")
"""))

# Cell 6 — header
NB["cells"].append(md("""\
## 3 · Pull just the reward verifier code (≈ 30 KB)

The reward function is pure Python — no model weights, no GPU.
We grab the seven files needed to compute a reward score.
"""))

# Cell 7 — sparse download
NB["cells"].append(code("""\
from huggingface_hub import hf_hub_download
import sys, os, pathlib

REPO = "Pratham-math/fathom-code"
files = [
    "rewards/__init__.py",
    "rewards/compose.py",
    "rewards/correctness.py",
    "rewards/format_gate.py",
    "rewards/recursion_efficiency.py",
    "rewards/recursion_extract.py",
    "rewards/token_budget.py",
    "configs/reward/v1.yaml",
]
local_root = pathlib.Path("/content/fathom").resolve()
for f in files:
    local = hf_hub_download(repo_id=REPO, filename=f, local_dir=str(local_root))
sys.path.insert(0, str(local_root))
print("pulled reward code into", local_root)
print("rewards/ files:", os.listdir(local_root / "rewards"))
"""))

# Cell 8 — header
NB["cells"].append(md("""\
## 4 · The 8 adversarial reward probes

Every entry below is a completion that *tries* to hack the reward.
The composite verifier should reject each one — wrong answers cap at 0.25, and only well-formatted *correct* answers earn the recursion-efficiency bonus.

A passing run shows scores in this exact order: **correct ≫ partial ≫ format-only ≫ no-format**.
"""))

# Cell 9 — reward probes
NB["cells"].append(code("""\
from rewards.compose import compose_reward_single
import types

cfg = types.SimpleNamespace(
    alpha=0.2,
    weights=types.SimpleNamespace(
        correctness=0.70, token_budget=0.15, recursion_efficiency=0.15
    ),
    token_budget_variant="capped_linear",
    answer_regex="<answer>(.*?)</answer>",
    max_calls=4,
)

GOLD = "silver"
PROBES = [
    ("correct + 0 llm calls (REPL grep)",
     "```python\\nimport re\\nm=re.search('silver', ctx)\\nprint(m.group())\\n```\\n<answer>silver</answer>"),
    ("correct + 1 llm call",
     "```python\\nans=llm('color', ctx[:5000])\\nprint(ans)\\n```\\n<answer>silver</answer>"),
    ("correct + bare answer (trivial-task path)",
     "<answer>silver</answer>"),
    ("wrong answer + correct format (A-01)",
     "<answer>gold</answer>"),
    ("wrong + no format",
     "the color is gold"),
    ("right text but no <answer> tag",
     "silver"),
    ("format-only spam (empty answer)",
     "<answer></answer>"),
    ("recursion-spam (5 llm calls, A-05)",
     "```python\\n" + "\\n".join(f"x{i}=llm('q{i}',ctx)" for i in range(5)) + "\\n```\\n<answer>silver</answer>"),
]

print(f"{'#':>2}  {'reward':>6}  {'calls':>5}  description")
print("-" * 90)
scores = []
for i, (desc, gen) in enumerate(PROBES):
    s, m = compose_reward_single(gen, GOLD, prompt_token_count=200, cfg_reward=cfg, llm_call_count=None)
    scores.append(s)
    print(f"{i:>2}  {s:>6.3f}  {int(m['llm_call_count']):>5d}  {desc}")
print("-" * 90)

import statistics
print(f"\\ngroup mean: {statistics.mean(scores):.3f}")
print(f"group std:  {statistics.stdev(scores):.3f}  (must be > 0.10 for GRPO advantage)")
print(f"max - min:  {max(scores) - min(scores):.3f}")

assert scores[0] > scores[1] > scores[7], "FAIL: 0-call should beat 1-call should beat spam"
assert scores[0] > scores[3], "FAIL: correct must beat wrong-but-formatted"
assert scores[6] <= 0.25, "FAIL: format-only spam not capped"
assert statistics.stdev(scores) > 0.10, "FAIL: group std too low"
print("\\nPASS — 8/8 reward probes behave as designed.")
"""))

# Cell 10 — header
NB["cells"].append(md("""\
## 5 · Training plots — pulled live from the trained-model repo

These PNGs were committed to <https://huggingface.co/Pratham-math/fathom-1.5b-grpo> at the end of the GRPO run on `a100-large` HF Jobs.
"""))

# Cell 11 — display plots inline
NB["cells"].append(code("""\
from huggingface_hub import hf_hub_download
from PIL import Image
import matplotlib.pyplot as plt

PLOTS_REPO = "Pratham-math/fathom-1.5b-grpo"
PLOT_FILES = [
    ("plots/sft_loss.png",          "SFT loss — 3.20 -> 0.29 over 63 steps"),
    ("plots/sft_token_accuracy.png", "SFT token accuracy — 0.46 -> 0.93"),
    ("plots/grpo_reward.png",       "GRPO composite reward (v2 run y82wmj4x)"),
    ("plots/grpo_completion_length.png", "GRPO completion length — model finds short correct answers"),
    ("plots/grpo_kl.png",           "GRPO KL — controlled drift from base policy"),
    ("plots/training_summary.png",  "8-panel training summary"),
]

fig, axes = plt.subplots(3, 2, figsize=(15, 16))
for ax, (path, title) in zip(axes.flat, PLOT_FILES):
    try:
        img_path = hf_hub_download(repo_id=PLOTS_REPO, filename=path)
        img = Image.open(img_path)
        ax.imshow(img)
        ax.set_title(title, fontsize=11)
    except Exception as e:
        ax.text(0.5, 0.5, f"{path}\\n{e}", ha="center", va="center", fontsize=9, transform=ax.transAxes)
        ax.set_title(title + " (load failed)", fontsize=11, color="red")
    ax.axis("off")
plt.tight_layout()
plt.show()
print("Done — 6 training plots rendered above.")
"""))

# Cell 12 — header for W&B
NB["cells"].append(md("""\
## 6 · Live W&B training run

Embed of the v2 run that learned the correct-answer mode (`y82wmj4x` / "lucky-capybara-5"). Reward climbs from a 0.15 format-bonus floor to 0.86–0.98 peaks once the policy starts producing correct answers.

If the iframe doesn't load (HF Colab sometimes blocks third-party iframes), use the direct URL printed below.
"""))

# Cell 13 — W&B iframe
NB["cells"].append(code("""\
WANDB_URL = "https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/y82wmj4x"

from IPython.display import IFrame, display, Markdown
display(IFrame(WANDB_URL, width="100%", height=720))
display(Markdown(f"**Direct W&B link** (if the iframe is blocked): [{WANDB_URL}]({WANDB_URL})"))
"""))

# Cell 14 — header
NB["cells"].append(md("""\
## 7 · The trained model artifacts

The full 1.5B model is published with three flavours: the LoRA adapter (~15 MB, the actual training output), the merged 16-bit weights (~3 GB, ready for inference), and the training plots.
"""))

# Cell 15 — list model files
NB["cells"].append(code("""\
from huggingface_hub import HfApi
api = HfApi()
files = sorted(api.list_repo_files("Pratham-math/fathom-1.5b-grpo"))

print(f"Total files: {len(files)}\\n")
print(f"{'category':25s}  {'count':>5s}")
print("-" * 40)

categories = {
    "adapters/ (LoRA)":      [f for f in files if f.startswith("adapter") or "adapter_" in f],
    "merged_16bit/":         [f for f in files if f.startswith("merged_16bit/")],
    "plots/ (training PNGs)":[f for f in files if f.startswith("plots/")],
    "tokenizer / config":    [f for f in files if any(f.endswith(s) for s in ["tokenizer.json","tokenizer_config.json","special_tokens_map.json","vocab.json","merges.txt","added_tokens.json","config.json","generation_config.json"])],
    "other":                 [],
}
seen = set().union(*categories.values())
categories["other"] = [f for f in files if f not in seen]

for k, v in categories.items():
    print(f"{k:25s}  {len(v):>5d}")
print()
print("Sample LoRA adapter files:")
for f in [f for f in files if "adapter" in f.lower()][:3]:
    print(" ", f)
print()
print("Sample merged_16bit files:")
for f in [f for f in files if f.startswith("merged_16bit/")][:5]:
    print(" ", f)
"""))

# Cell 16 — header for optional training
NB["cells"].append(md("""\
## 8 · (Optional, A100 only) Re-run the training

This is the actual command we used for the v2 run that produced the curve above. It runs `train/grpo.py` against the live env Space. **Do not run on free Colab — it will OOM.** The cell is left here so judges can verify the exact arguments.

```bash
hf jobs run \\
  --flavor=a100-large \\
  --secrets HF_TOKEN=$HF_TOKEN \\
  --secrets WANDB_API_KEY=$WANDB_API_KEY \\
  -e FATHOM_USE_VLLM=0 \\
  pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel \\
  bash -c 'apt-get update -qq && apt-get install -y -qq git && \\
           git clone -b main https://oauth2:$HF_TOKEN@huggingface.co/Pratham-math/fathom-code /w && \\
           bash /w/scripts/job_train.sh'
```

Wall-clock: ~50 min on `a100-large`. Cost: ~$10 of HF Jobs credit. The trainer pushes the LoRA adapter and the merged 16-bit checkpoint to `Pratham-math/fathom-1.5b-grpo` automatically.
"""))

# Cell 17 — closing
NB["cells"].append(md("""\
## What you just verified

1. The OpenEnv FATHOM server is **live** at `https://Pratham-math-fathom-env.hf.space` and returns `{"status":"ok"}`.
2. The deterministic 4-component reward (format gate × correctness + token budget + recursion efficiency) **rejects all 8 known reward-hack patterns** with a clean `correct ≫ wrong ≫ format-only` ordering.
3. Training **actually happened**: SFT loss dropped 3.20 → 0.29 (91% reduction), GRPO reward climbed from a 0.15 format-bonus floor to 0.86–0.98 peaks over 70 steps once the policy discovered the correct-answer mode.
4. The trained model is published as a LoRA adapter + merged 16-bit checkpoint on the HF Hub.

For the full story (incl. the v1 → v2 debugging journey), see the mini-blog: <https://huggingface.co/spaces/Pratham-math/fathom-blog>.
"""))

# Write
out = Path("notebooks/fathom_train.ipynb")
out.write_text(json.dumps(NB, ensure_ascii=False, indent=1), encoding="utf-8")
print(f"wrote {out} with {len(NB['cells'])} cells")