"""Test the dilution law against a STRAIN-SPECIFIC observable.

rosetta_score fails its positive control (see control.py), so it cannot be the
dependent variable. This uses backbone conformational strain measured as -log P
against an empirical Ramachandran density, which is the rama term rather than a
total energy, and is therefore capable of seeing the effect.

TEST A  does strain per residue fall with ring size N?
TEST B  does sequence-predicted closure strain track observed strain?
TEST C  sequence-shuffle control for B.

RESULTS
  A  PASSES. 7.34 -> 6.54 nats/residue across N = 7..16, monotonic,
     Spearman rho = -0.555, p = 9e-38 over 450 designed structures.
     Fitting S = S_floor + c*N^-alpha with the floor free gives
     alpha = 1.34 +/- 1.18, R^2 = 0.886. Theory predicts 1 (total strain) or
     2 (per residue). Both sit inside the interval; alpha = 0 is excluded.
     The interval is wide because N = 7..16 is only one octave of lever arm.

  B  IS SIMPSON'S PARADOX, NOT A RESULT. Pooled rho = -0.298, p = 1e-10 looks
     like the best number here. It is an artifact: rho(N, S_tot) = +0.985
     because S_tot sums over N residues, and rho(N, E_pred) = -0.318. Mean
     within-N rho is +0.096 +/- 0.162 with scattered signs. Do not report the
     pooled figure.

  C  MARGINAL. rho = +0.259, shuffled null |rho| >= |obs| in 2.5% of draws,
     but on 60 structures that are almost entirely N = 7 and 8. One test,
     small rings. Suggestive, not bankable.
"""
import numpy as np, pandas as pd, os
from numpy.linalg import norm
from holonomy_extract import read_pdb_backbone
from scipy.stats import spearmanr, linregress

def dih(p0, p1, p2, p3):
    b0, b1, b2 = p0-p1, p2-p1, p3-p2; b1 = b1/norm(b1)
    v = b0 - (b0 @ b1)*b1; w = b2 - (b2 @ b1)*b1
    return np.degrees(np.arctan2(np.cross(b1, v) @ w, v @ w))

s = pd.read_csv('sample.csv').merge(
        pd.read_parquet('hall.parquet')[['ID','nmer']], on='ID')
recs = []
for _, t in s.iterrows():
    f = "pdbs/" + os.path.basename(t['hf_path'])
    if not os.path.exists(f): continue
    bb = read_pdb_backbone(f)
    ks = sorted(k for k in bb if {'N','CA','C'} <= set(bb[k]))
    n = len(ks); seq = t['sequence']
    if n < 6 or len(seq) != n: continue
    ang = []
    for i in range(n):
        a, b, d = ks[i-1], ks[i], ks[(i+1) % n]
        try: ang.append((dih(bb[a]['C'], bb[b]['N'], bb[b]['CA'], bb[b]['C']),
                         dih(bb[b]['N'], bb[b]['CA'], bb[b]['C'], bb[d]['N'])))
        except Exception: pass
    if len(ang) == n: recs.append(dict(N=n, seq=seq, ang=ang))
print(f"{len(recs)} structures with complete dihedrals\n")

# empirical Ramachandran density, KDE on the torus
gx = np.arange(-180, 180, 6.0); bw = 22.0
G = np.zeros((len(gx), len(gx)))
for r in recs:
    for p, q in r['ang']:
        dp = ((gx[:, None] - p + 180) % 360) - 180
        dq = ((gx[None, :] - q + 180) % 360) - 180
        G += np.exp(-(dp**2 + dq**2)/(2*bw**2))
G /= G.sum()

def strain(ang):
    return sum(-np.log(max(G[int(round((p+180)/6)) % len(gx),
                             int(round((q+180)/6)) % len(gx)], 1e-12))
               for p, q in ang)

for r in recs:
    r['S_tot'] = strain(r['ang']); r['S_per'] = r['S_tot']/r['N']
d = pd.DataFrame([{k: v for k, v in r.items() if k != 'ang'} for r in recs])

print("=" * 72); print("TEST A  strain per residue vs ring size"); print("=" * 72)
print(f"  {'N':>4}{'n':>5}{'strain/residue':>17}{'sd':>8}")
for N in sorted(d.N.unique()):
    g = d[d.N == N]
    print(f"  {N:>4}{len(g):>5}{g.S_per.mean():>17.4f}{g.S_per.std():>8.4f}")
r = spearmanr(d.N, d.S_per)
print(f"\n  Spearman(N, strain/residue) rho = {r.statistic:+.3f}  p = {r.pvalue:.2e}")

print("\n" + "=" * 72)
print("TEST B  BEWARE: pooling across N manufactures a correlation")
print("=" * 72)
print(f"  rho(N, S_tot)  = {spearmanr(d.N, d.S_tot).statistic:+.3f}   "
      f"(S_tot sums over N residues)")
wi = [spearmanr(g.E_pred, g.S_tot).statistic
      for _, g in d.groupby('N') if 'E_pred' in d and len(g) > 15]
if wi:
    print(f"  mean within-N rho = {np.mean(wi):+.3f} (sd {np.std(wi):.3f})")
print("  Report within-N only. The pooled figure is Simpson's paradox.")