"""Build the demo notebook: samples from the EP-trained 135M model beside its backprop twin.

The notebook ships with its outputs already filled in, so it can be read without a checkout or a
GPU, and the code in it is real, so it also runs if you have the checkpoints.
"""
import json
import re
from pathlib import Path

RUN = Path('/home/yurenh2/ept/ep_run')
OUT = Path('/home/yurenh2/ept/assets/EPT_135M_samples.ipynb')
PROMPTS = [
    'The main difference between a virus and a bacterium is',
    'To find the area of a circle, you',
    'In 1815, the eruption of Mount Tambora',
    'Photosynthesis is the process by which',
]


def parse(path):
    txt = Path(path).read_text(encoding='utf-8', errors='replace')
    parts = re.split(r'\n--- prompt (\d+), sample (\d+) ---\n', txt)
    out = {}
    for i in range(1, len(parts), 3):
        out[(int(parts[i]), int(parts[i + 1]))] = parts[i + 2].strip()
    return out


def md(source):
    return {'cell_type': 'markdown', 'metadata': {}, 'source': source.splitlines(keepends=True)}


def code(source, stdout=None):
    outputs = []
    if stdout is not None:
        outputs.append({'output_type': 'stream', 'name': 'stdout',
                        'text': stdout.splitlines(keepends=True)})
    return {'cell_type': 'code', 'execution_count': None, 'metadata': {},
            'source': source.splitlines(keepends=True), 'outputs': outputs}


def main():
    ep, bp = parse(RUN / 'runs/gen_rlin.txt'), parse(RUN / 'runs/gen_bp.txt')
    cells = [md("""# A 135M language model trained without backpropagation

Every parameter update in the model below comes from Equilibrium Propagation. There is no backward
pass anywhere in its training, and no layer, block, or output head is trained with a backprop rule.
The model is an ordinary 12 layer transformer, OLMo2 style blocks, 32k vocabulary, trained from
scratch on 2.7B tokens of FineWeb-Edu.

Beside it, for the same prompts and the same sampling seeds, is its backprop twin: identical
architecture, tokenizer, data order, optimizer, step budget, and evaluation, differing only in the
training rule. Final validation cross-entropy is 3.209 for EP and 3.207 for backprop, against a
spread of 0.006 between backprop seeds.

Sampling is an ordinary forward pass. Equilibrium Propagation appears only during training, so
nothing unusual happens here at inference time.
"""),
             code("""from pathlib import Path
import re, subprocess

RUN = Path('ep_run')          # your checkout
CKPT = {'EP': 'runs/fw135m_rlin_s440000.pt',      # trained with Equilibrium Propagation
        'BP': 'runs/fw135m_bp_s440000.pt'}        # the matched backprop twin

def samples(which, n=2, new=110, temp=0.8, topk=40):
    \"\"\"Return {(prompt, sample): text}. Uses the cached generation if it is present.\"\"\"
    cached = RUN / 'runs' / ('gen_rlin.txt' if which == 'EP' else 'gen_bp.txt')
    if cached.exists():
        txt = cached.read_text(encoding='utf-8', errors='replace')
    else:
        txt = subprocess.run([
            'python3', 'casc_eq_train.py', '--gen', str(n), '--gen_new', str(new),
            '--gen_temp', str(temp), '--gen_topk', str(topk), '--tag', f'gen_{which}',
            '--resume', CKPT[which], '--data', 'fineweb_edu', '--wandb', '', '--untie',
            '--L', '12', '--C', '768', '--H', '12', '--T', '256', '--B', '4',
            '--olmo2', '--steps', '440000'], cwd=RUN, capture_output=True, text=True).stdout
    parts = re.split(r'\\n--- prompt (\\d+), sample (\\d+) ---\\n', txt)
    return {(int(parts[i]), int(parts[i+1])): parts[i+2].strip() for i in range(1, len(parts), 3)}

EP, BP = samples('EP'), samples('BP')
print(f'{len(EP)} samples from the EP model, {len(BP)} from the backprop twin')""",
                  f'{len(ep)} samples from the EP model, {len(bp)} from the backprop twin\n'),
             md("""## Side by side

Same prompt, same sampling seed, same temperature and top-k. The left column was trained without
backpropagation.
"""),
             ]

    def render(pi):
        lines = [f'PROMPT: {PROMPTS[pi-1]}', '']
        for si in (1, 2):
            lines += [f'  [ trained with Equilibrium Propagation, sample {si} ]',
                      '  ' + ep.get((pi, si), '').replace('\n', '\n  '), '',
                      f'  [ trained with backpropagation, sample {si} ]',
                      '  ' + bp.get((pi, si), '').replace('\n', '\n  '), '']
        return '\n'.join(lines) + '\n'

    src = """def show(pi):
    print(f'PROMPT: {PROMPTS[pi-1]}\\n')
    for si in (1, 2):
        for name, d in (('Equilibrium Propagation', EP), ('backpropagation', BP)):
            print(f'  [ trained with {name}, sample {si} ]')
            print('  ' + d[(pi, si)].replace('\\n', '\\n  '), '\\n')

PROMPTS = %r
show(%d)"""
    cells.append(code(src % (PROMPTS, 1), render(1)))
    for pi in (2, 3, 4):
        cells.append(code(f'show({pi})', render(pi)))

    cells.append(md("""## What this does and does not show

The two models are close in quality because Equilibrium Propagation is designed to compute the same
gradient as backpropagation in the small nudge limit, and our measurements confirm that it does:
cosine to the backprop gradient stays near 0.9997 through training. The interest is not that the
gradient is different. It is that this update is local, so a physical substrate can perform it
without a global backward pass, and that it survives at this scale under the constraints such a
substrate imposes, including 8 bit weights, injected device noise, and the precision of the
contrast readout.

That last one turned out to matter more than we expected. An earlier version of this run plateaued
around 3.58 because the contrast was recovered by subtracting two large states, which destroys the
part of the nudge that falls below single precision resolution. Reading the contrast from the
stored displacement instead removes the effect at no cost, and is what separates the two EP curves
in the accompanying figure.
"""))

    import uuid
    for c in cells:
        c['id'] = uuid.uuid4().hex[:8]   # nbformat 4.5 requires cell ids
    nb = {'cells': cells, 'metadata': {'kernelspec': {'display_name': 'Python 3',
          'language': 'python', 'name': 'python3'}, 'language_info': {'name': 'python'}},
          'nbformat': 4, 'nbformat_minor': 5}
    OUT.write_text(json.dumps(nb, indent=1, ensure_ascii=False))
    print('wrote', OUT, f'({OUT.stat().st_size/1024:.1f} KB, {len(cells)} cells)')


if __name__ == '__main__':
    main()