"""Decomposition + robustness statistics for the committed results (results.md).

Computes, from the saved jlens_v3 artifacts and the base checkpoint:
  1. W_U row-norm decomposition: r(||W_U[k]||, freq), r vs log-frequency,
     Spearman, and per-layer partial correlation of faithful norm with
     frequency after regressing out ||W_U[k]||.
  2. Pearson p-values (normal-approx two-sided) and Spearman per layer.
  3. Synthetic-pair ratio with across-seed mean/SD and bootstrap CI.

Run inside the meru container where artifacts exist:
    python3 src/stats_decomp.py
"""
import os
import pickle
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import numpy as np
import torch

V = 65
TRAIN_BIN = 'data/shakespeare_char/train.bin'
CKPT = 'out-shakespeare-char/ckpt.pt'


def rankdata(a):
    """Average ranks (ties averaged) — minimal scipy-free rankdata."""
    sorter = np.argsort(a, kind='mergesort')
    ranks = np.empty(len(a), dtype=float)
    ranks[sorter] = np.arange(1, len(a) + 1)
    i = 0
    while i < len(a):
        j = i
        while j + 1 < len(a) and a[sorter[j + 1]] == a[sorter[i]]:
            j += 1
        if j > i:
            ranks[sorter[i:j + 1]] = (i + j) / 2.0 + 1
        i = j + 1
    return ranks


def spearman(x, y):
    return np.corrcoef(rankdata(x), rankdata(y))[0, 1]


def pearson_p(r, n):
    """Two-sided p for Pearson r via normal approx on Fisher's z (fine for n=65)."""
    import math
    from math import erf, sqrt
    z = abs(r) * math.sqrt((n - 2) / max(1 - r * r, 1e-12))
    return 2 * (1 - 0.5 * (1 + erf(z / sqrt(2))))


def load_layer(l):
    for p in (f'outputs/jlens_v3/layer{l}.pt', f'outputs/jlens_v3_layer{l}.pt'):
        if os.path.exists(p):
            return torch.load(p, map_location='cpu')
    raise FileNotFoundError(f'no jlens artifact for layer {l}')


def main():
    train = np.memmap(TRAIN_BIN, dtype=np.uint16, mode='r')
    counts = np.bincount(train, minlength=V).astype(float)
    freq = counts / counts.sum() * 100
    logf = np.log10(freq)

    ckpt = torch.load(CKPT, map_location='cpu')
    wu = ckpt['model']['lm_head.weight'].float().numpy()
    wu_norms = np.linalg.norm(wu, axis=1)

    print("=== W_U row-norm decomposition ===")
    print(f"r(||W_U[k]||, freq)        = {np.corrcoef(wu_norms, freq)[0, 1]:+.3f}")
    print(f"r(||W_U[k]||, log10 freq)  = {np.corrcoef(wu_norms, logf)[0, 1]:+.3f}")
    print(f"Spearman(||W_U[k]||, freq) = {spearman(wu_norms, freq):+.3f}")

    print("\n=== per-layer decomposition ===")
    print(f"{'layer':<6}{'r(faith,freq)':>14}{'p':>10}{'r(faith,WU)':>13}"
          f"{'Spear(faith)':>14}{'partial|WU':>13}")
    for l in range(6):
        d = load_layer(l)
        fn = np.array([d['faithful_norms'][k] for k in range(V)])
        rf = np.corrcoef(fn, freq)[0, 1]
        rfw = np.corrcoef(fn, wu_norms)[0, 1]
        sf = spearman(fn, freq)
        A = np.vstack([wu_norms, np.ones(V)]).T
        resid = fn - A @ np.linalg.lstsq(A, fn, rcond=None)[0]
        r_part = np.corrcoef(resid, freq)[0, 1]
        print(f"L{l:<5}{rf:>+14.3f}{pearson_p(rf, V):>10.1e}{rfw:>+13.3f}"
              f"{sf:>+14.3f}{r_part:>+13.3f}")

    print("\n=== synthetic pair ratio ===")
    with open('data/synth_pair/meta.pkl', 'rb') as f:
        meta = pickle.load(f)
    iid_s = meta['stoi']['@']
    iid_n = meta['stoi']['#']
    ratios = []
    for s in range(3):
        vals_s, vals_n = [], []
        for l in (2, 3, 4):
            d = torch.load(f'outputs/synth_pair/seed{s}/layer{l}.pt',
                           map_location='cpu')
            fn = d['faithful_norms']
            vals_s.append(fn[iid_s])
            vals_n.append(fn[iid_n])
        ratios.append(np.mean(vals_s) / max(np.mean(vals_n), 1e-9))
    r_arr = np.array(ratios)
    rng = np.random.RandomState(0)
    boot = [np.mean(rng.choice(r_arr, 3, replace=True)) for _ in range(10000)]
    print(f"ratio per seed: {[f'{r:.3f}' for r in ratios]}")
    print(f"mean +/- SD: {r_arr.mean():.3f} +/- {r_arr.std(ddof=1):.3f}")
    print(f"bootstrap 95% CI: [{np.percentile(boot, 2.5):.3f}, "
          f"{np.percentile(boot, 97.5):.3f}]")


if __name__ == '__main__':
    main()