"""Validation curves at 135M: EP with the readout fix, EP without it, and the backprop twins.

The point of the figure is that the first and third are the same curve, and the second stops
descending partway through. Run after fw135m_rlin completes; it reads the training logs directly.
"""
import re
import statistics as st
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

RUNS = Path('/home/yurenh2/ept/ep_run/runs')
TOTAL = 440000
TOK_PER_STEP = 24 * 256   # effective batch x context


def curve(tag):
    pts = []
    for line in (RUNS / f'{tag}.log').read_bytes().decode(errors='replace').splitlines():
        m = re.match(rf'step (\d+)/{TOTAL} \| train [\d.]+ val ([\d.]+)', line)
        if m:
            pts.append((int(m.group(1)), float(m.group(2))))
    pts.sort()
    return pts


def smooth(pts, win=25):
    xs = [p[0] for p in pts]
    ys = [p[1] for p in pts]
    out = []
    for i in range(len(ys)):
        lo, hi = max(0, i - win // 2), min(len(ys), i + win // 2 + 1)
        out.append(st.median(ys[lo:hi]))
    return xs, out


def main():
    # EP is drawn thick and underneath; the backprop seeds go on top as thin lines, so the reader
    # sees them tracking rather than being hidden by whichever curve was plotted last.
    arms = [
        ('fw135m_rlin',  'Equilibrium Propagation',     '#1f6feb', 3.0, '-',  2),
        ('fw135m_bsign', 'Equilibrium Propagation,\nwithout the readout fix', '#d1442f', 2.6, '-', 2),
        ('fw135m_bp',    'Backprop, seed 1',            '#2b3540', 1.0, '-',  4),
        ('fw135m_bp_s2', 'Backprop, seed 2',            '#7d8794', 1.0, '-',  4),
    ]
    fig, ax = plt.subplots(figsize=(7.6, 4.9))
    finals = {}
    for tag, label, color, lw, ls, zo in arms:
        pts = curve(tag)
        if not pts:
            print(f'skip {tag}: no data')
            continue
        xs, ys = smooth(pts)
        toks = [x * TOK_PER_STEP / 1e9 for x in xs]
        ax.plot(toks, ys, color=color, lw=lw, ls=ls, label=label, zorder=zo)
        tail = [v for s, v in pts if s >= 0.9 * TOTAL]
        finals[tag] = st.mean(tail) if tail else float('nan')
        print(f'{tag:14s} last step {pts[-1][0]:>7}  tail mean {finals[tag]:.4f}')

    ax.set_xlabel('training tokens (billions)')
    ax.set_ylabel('validation cross-entropy')
    ax.set_ylim(3.15, 4.35)
    ax.set_xlim(0, TOTAL * TOK_PER_STEP / 1e9)
    ax.grid(alpha=0.18, lw=0.7)
    for side in ('top', 'right'):
        ax.spines[side].set_visible(False)
    h, l = ax.get_legend_handles_labels()
    order = [2, 3, 0, 1]   # backprop seeds first in the legend, matching how one reads the plot
    leg = ax.legend([h[i] for i in order], [l[i] for i in order],
                    frameon=False, fontsize=9.5, loc='upper right', handlelength=1.6)
    for t in leg.get_texts():
        t.set_va('center')

    if 'fw135m_bsign' in finals and 'fw135m_rlin' in finals:
        ax.annotate('single-precision rounding in the contrast readout;\nthe run stops improving',
                    xy=(1.9, finals['fw135m_bsign'] + 0.005), xytext=(1.35, 3.86),
                    fontsize=9, color='#d1442f',
                    arrowprops=dict(arrowstyle='-', color='#d1442f', lw=0.9, alpha=0.8))
        ax.annotate(f"EP {finals['fw135m_rlin']:.3f}\nBP {st.mean([finals['fw135m_bp'], finals['fw135m_bp_s2']]):.3f}",
                    xy=(2.62, 3.245), fontsize=9.5, color='#2b3540', ha='right')

    ax.set_title('135M-parameter transformer language model, FineWeb-Edu',
                 fontsize=10.5, color='#333', pad=10, loc='left')
    fig.tight_layout()
    for ext in ('png', 'pdf'):
        p = Path(f'/home/yurenh2/ept/assets/figs/fig_135m_curves.{ext}')
        fig.savefig(p, dpi=300, bbox_inches='tight')
        print('wrote', p)


if __name__ == '__main__':
    main()