# SCRIPT: RS03B_LAYER_RESOLVED
# W2 of RS-03 repeated with the layer resolved (breakpoints inside it). Kill: PREREGISTRATION_RS03.md addendum.
import json, numpy as np
import rs03_push_lib as L
a1, a2, b1, b2 = 4.0, 1.0, 1.0, 0.5
rows = {}; ok = True; rich = 0.0
for kind in ['tanh', 'erf', 'ramp']:
    rows[kind] = []
    for th in [10, 25, 29, 31, 60]:
        rs, errs = [], []
        for ow in [0.08, 0.04, 0.02, 0.01]:
            r, C = L.refl_resolved(kind, a1, b1, a2, b2, ow/10, th)
            rs.append(r); errs.append(abs(r - L.gamma_pred(a1, a2, C)))
        mono = all(e2 < e1 for e1, e2 in zip(errs, errs[1:]))
        q = np.log2(abs(rs[1] - rs[2])/abs(rs[2] - rs[3]))
        r0 = rs[3] + (rs[3] - rs[2])/(2**q - 1)
        rich = max(rich, abs(r0 - L.gamma_pred(a1, a2, C)))
        ok &= mono
        rows[kind].append(dict(theta=th, errs=errs, order=q))
    print(f"W2b {kind:4s}: " + "  ".join(f"th{r_['theta']}: {r_['errs'][-1]:.1e} (order {r_['order']:.2f})" for r_ in rows[kind]), flush=True)
ok = bool(ok and rich < 1e-5)
print(f"W2b Richardson to w = 0: max |r0 - Gamma| = {rich:.1e}  {'PASS' if ok else 'KILL'}")
json.dump(dict(rows=rows, richardson=rich, passed=ok), open('rs03b_results.json', 'w'), indent=1, default=float)