#!/usr/bin/env python3
"""Reproduce a deliberately inadequate fit to public NR scattering data.

The design is fixed in calibration-design.json. The baseline is an analytic
comparison model, not a fitted Hamiltonian. No published EOB angle columns
enter this calculation: their relative-angle observable is different.
"""
import argparse
import hashlib
import json
from pathlib import Path
import math
import platform
from urllib.request import Request, urlopen

import mpmath as mp
import numpy as np
import scipy
from scipy.linalg import lstsq

HERE = Path(__file__).resolve().parent
INPUT = HERE / 'ScatteringAngles.dat'
DATA_URL = 'https://arxiv.org/src/2507.08071v1/anc/ScatteringAngles.dat'
DATA_SHA256 = '4c2a8fef5ec23f12c6014b0d5cdaa5f446ec02aadff4191f21bd8d51a9354823'
DESIGN = HERE / 'calibration-design.json'
MODELS = {'constant': 1, 'linear': 4, 'quadratic': 7}


def read_rows():
    design = json.loads(DESIGN.read_text())
    rows = []
    for index, line in enumerate(INPUT.read_text().splitlines()[1:], 1):
        a = list(map(float, line.replace('[', ' ').replace(']', ' ').split()))
        assert len(a) == 19
        total_energy, angular_momentum, q = a[:3]
        nu = q/(1+q)**2
        gamma = 1+(total_energy**2-1)/(2*nu)
        # Only the declared nonspinning rows are fitted. For exclusion diagnostics
        # on aligned-spin rows distinguish orbital from total angular momentum.
        spin_angular_momentum = (a[5]*q*q+a[8])/(1+q)**2
        j = (angular_momentum-spin_angular_momentum)/nu
        rho = 8/(4-3*gamma**2+gamma*math.sqrt(9*gamma**2-8))
        jc = rho/math.sqrt(rho-3)
        q_group = round(q)
        assert abs(q-q_group) < 0.001
        spin = max(map(abs, a[3:9]))
        reasons = []
        if spin > design['spin_tolerance']:
            reasons.append('spinning')
        lo, hi = design['gamma_training_domain']
        if not lo <= gamma <= hi:
            reasons.append('energy outside fitted domain')
        if j <= jc:
            reasons.append('no scattering geodesic baseline')
        partition = 'stress'
        if not reasons:
            for group in ('training', 'validation', 'test'):
                if q_group in design[group+'_q']:
                    partition = group
                    break
        # An arithmetic average propagates an interval by adding its endpoints.
        # The source's estimates are NOT asserted to be Gaussian or rigorous bounds.
        scale_degrees = max((a[10]+a[13])/2, (a[11]+a[14])/2)
        assert scale_degrees > 0
        rows.append(dict(id=index, Gamma=total_energy, gamma=gamma, j=j, nu=nu,
                         q=q, q_group=q_group, spin=spin, j_over_jc=j/jc,
                         target=math.radians((a[9]+a[12])/2),
                         scale=math.radians(scale_degrees),
                         partition=partition, excluded_by=reasons))
    return rows


def baseline(row, precision=45):
    with mp.workdps(precision):
        g, j, nu = (mp.mpf(str(row[key])) for key in ('gamma', 'j', 'nu'))
        p, q = 1/j**2-mp.mpf(1)/12, (3*g*g-2)/(6*j*j)-mp.mpf(1)/108
        phase = mp.acos(3*q/(2*p)*mp.sqrt(-3/p))/3
        a, b, c = sorted(mp.mpf(1)/6+2*mp.sqrt(-p/3)*mp.cos(phase+2*k*mp.pi/3)
                         for k in range(3))
        assert a < 0 < b < c
        k = (b-a)/(c-a)
        geodesic = 2*mp.sqrt(2/(c-a))*(mp.ellipk(k)-mp.ellipf(mp.asin(mp.sqrt(-a/(b-a))), k))-mp.pi
        Gamma = mp.sqrt(1+2*nu*(g-1))
        finite_mass_2pm = 3*mp.pi*(5*g*g-1)*(1/Gamma-1)/(4*j*j)
        return float(geodesic+finite_mass_2pm)


def baseline_quad(row):
    with mp.workdps(60):
        g, j, nu = (mp.mpf(str(row[key])) for key in ('gamma', 'j', 'nu'))
        p, q = 1/j**2-mp.mpf(1)/12, (3*g*g-2)/(6*j*j)-mp.mpf(1)/108
        phase = mp.acos(3*q/(2*p)*mp.sqrt(-3/p))/3
        a, b, c = sorted(mp.mpf(1)/6+2*mp.sqrt(-p/3)*mp.cos(phase+2*k*mp.pi/3)
                         for k in range(3))
        half = mp.quad(lambda s:mp.sqrt(2*b)/mp.sqrt((b*(1-s*s)-a)*(c-b*(1-s*s))),[0,.01,.1,1])
        correction = 3*mp.pi*(5*g*g-1)*(1/mp.sqrt(1+2*nu*(g-1))-1)/(4*j*j)
        return float(2*half-mp.pi+correction)


def features(row, name):
    x, t, nu = 4/row['j'], (row['gamma']-1)/0.1, row['nu']
    return nu*x**3*np.array([1, x, t, nu, x*x, x*t, t*t][:MODELS[name]])


def fit(rows, name):
    X = np.array([features(r, name)/r['scale'] for r in rows])
    y = np.array([(r['target']-r['baseline'])/r['scale'] for r in rows])
    norms = np.linalg.norm(X, axis=0)
    A = X/norms
    b_svd, _, rank, singular = lstsq(A, y, lapack_driver='gelsd')
    b_qr = lstsq(A, y, lapack_driver='gelsy')[0]
    assert rank == MODELS[name]
    assert np.max(np.abs(A@(b_svd-b_qr))) < 1e-7
    coefficient = b_svd/norms
    _, covariance_singular, vh = np.linalg.svd(A, full_matrices=False)
    covariance_scaled = (vh.T/covariance_singular**2)@vh
    # Check the whitened design directly: forming A.T@A would square its
    # condition number and make an identity-product check lose extra digits.
    whitened = (A@vh.T)/covariance_singular
    assert np.max(np.abs(whitened.T@whitened-np.eye(rank))) < 1e-9
    covariance = covariance_scaled/norms[:,None]/norms[None,:]
    return coefficient, dict(rank=int(rank), scaled_design_condition=float(singular[0]/singular[-1]),
                            scaled_hessian_condition=float((singular[0]/singular[-1])**2),
                            formal_gaussian_scale_covariance_rad_squared=covariance.tolist(),
                            covariance_scope='Conditional on independent Gaussian errors having the chosen scales and an exact linear model; not asserted for the NR estimates or this misspecified model',
                            qr_svd_max_prediction_difference_rad=float(np.max(np.abs(np.array([features(r,name) for r in rows])@((b_svd-b_qr)/norms)))))


def metrics(rows, name, coefficient):
    residual = np.array([r['baseline']+features(r, name)@coefficient-r['target'] for r in rows])
    raw_baseline = np.array([r['baseline']-r['target'] for r in rows])
    scales = np.array([r['scale'] for r in rows])
    return dict(n=len(rows), source_rows=[r['id'] for r in rows],
                inputs=[{key:r[key] for key in ('q','nu','Gamma','gamma','j')} for r in rows],
                prediction_degrees=((residual+np.array([r['target'] for r in rows]))*180/math.pi).tolist(),
                rms_degrees=float(np.sqrt(np.mean(residual**2))*180/math.pi),
                max_absolute_degrees=float(np.max(np.abs(residual))*180/math.pi),
                rms_in_source_error_scales=float(np.sqrt(np.mean((residual/scales)**2))),
                baseline_rms_degrees=float(np.sqrt(np.mean(raw_baseline**2))*180/math.pi),
                max_absolute_fractional_error=float(np.max(np.abs(residual/np.array([r['target'] for r in rows])))),
                residual_degrees=(residual*180/math.pi).tolist())


def main():
    global INPUT
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', choices=('train', 'test', 'reproduce'), default='reproduce',
                        help='Reproduce evaluates the already-published test; it is not a new blind experiment')
    parser.add_argument('--input', type=Path, default=INPUT)
    parser.add_argument('--fetch', action='store_true', help='Download the immutable source table if absent')
    parser.add_argument('--output-dir', type=Path, default=Path('scattering-calibration-results'))
    args = parser.parse_args()
    INPUT = args.input.resolve()
    if args.fetch and not INPUT.exists():
        request = Request(DATA_URL, headers={'User-Agent': 'blackhole.org reproducible teaching example'})
        with urlopen(request, timeout=90) as response:
            contents = response.read()
        assert hashlib.sha256(contents).hexdigest() == DATA_SHA256
        INPUT.parent.mkdir(parents=True, exist_ok=True)
        INPUT.write_bytes(contents)
    assert hashlib.sha256(INPUT.read_bytes()).hexdigest() == DATA_SHA256, 'Source table checksum mismatch'
    args.output_dir.mkdir(parents=True, exist_ok=True)
    rows = read_rows()
    # Domain and all partitions depend only on initial-condition metadata.
    for row in rows:
        if row['partition'] != 'stress':
            row['baseline'] = baseline(row)
            assert abs(row['baseline']-baseline_quad(row)) < 1e-13
    parts = {key: [r for r in rows if r['partition']==key]
             for key in ('training', 'validation', 'test', 'stress')}
    hashes = {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in (INPUT, DESIGN, Path(__file__))}
    summary = dict(hashes=hashes, environment=dict(python=platform.python_version(),
                   numpy=np.__version__, scipy=scipy.__version__, mpmath=mp.__version__),
                   counts={k:len(v) for k,v in parts.items()},
                   partitions={k:[r['id'] for r in v] for k,v in parts.items()},
                   exclusions=[{'id':r['id'],'reason':r['excluded_by']} for r in parts['stress']])
    build = getattr(scipy.__config__, 'CONFIG', {})
    summary['environment']['blas_lapack'] = {
        name: {key: value for key, value in build.get('Build Dependencies', {}).get(name, {}).items()
               if key in ('name', 'version', 'has ilp64')}
        for name in ('blas', 'lapack')}
    summary['environment']['scipy_compilers'] = {
        name: {key: value for key, value in data.items() if key in ('name', 'version')}
        for name, data in build.get('Compilers', {}).items()}
    summary['environment']['platform'] = platform.platform()
    fits = {}
    for name in MODELS:
        coef, check = fit(parts['training'], name)
        fits[name] = dict(coefficients=coef.tolist(), check=check,
                         training=metrics(parts['training'],name,coef),
                         validation=metrics(parts['validation'],name,coef))
    best = min(MODELS, key=lambda name: fits[name]['validation']['rms_in_source_error_scales'])
    minimum = fits[best]['validation']['rms_in_source_error_scales']
    best = min((name for name in MODELS if fits[name]['validation']['rms_in_source_error_scales'] <= minimum*1.01),
               key=lambda name:MODELS[name])
    summary['candidates'] = fits
    summary['selected'] = best
    summary['selected_passes_training_validation_one_percent'] = all(
        fits[best][part]['max_absolute_fractional_error'] <= .01
        for part in ('training','validation'))
    # Nested training-size and weak-ridge variations use validation only.
    equal = [r for r in parts['training'] if r['q_group']==1]
    unequal = [r for r in parts['training'] if r['q_group']!=1]
    summary['training_size_checks'] = []
    for count in (10,13,len(equal)):
        subset = equal[:count]+unequal
        coefficient, check = fit(subset,best)
        summary['training_size_checks'].append(dict(n=len(subset), coefficients=coefficient.tolist(),
                 condition=check['scaled_design_condition'],validation=metrics(parts['validation'],best,coefficient)))
    X = np.array([features(r,best)/r['scale'] for r in parts['training']])
    y = np.array([(r['target']-r['baseline'])/r['scale'] for r in parts['training']])
    norms = np.linalg.norm(X,axis=0)
    summary['ridge_sensitivity'] = []
    for lam in (1e-8,1e-4):
        augmented = np.vstack((X/norms,np.sqrt(lam)*np.eye(MODELS[best])))
        coefficient = lstsq(augmented,np.r_[y,np.zeros(MODELS[best])])[0]/norms
        summary['ridge_sensitivity'].append(dict(lambda_scaled=lam,
                 validation=metrics(parts['validation'],best,coefficient)))
    # Test evaluation requires the saved, matching training/validation decision.
    locked = args.output_dir/'calibration-training.json'
    if args.mode in ('train', 'reproduce'):
        locked.write_text(json.dumps(summary, indent=2)+'\n')
    if args.mode in ('test', 'reproduce'):
        previous = json.loads(locked.read_text())
        assert previous['hashes'] == hashes and previous['selected'] == best
        assert previous['candidates'] == fits
        summary['test'] = metrics(parts['test'], best, np.array(fits[best]['coefficients']))
        summary['accepted'] = summary['selected_passes_training_validation_one_percent'] and summary['test']['max_absolute_fractional_error'] <= .01
        stress_defined = [r for r in parts['stress'] if r['j_over_jc'] > 1 and r['spin'] <= 1e-6]
        for row in stress_defined:
            row['baseline'] = baseline(row)
            assert abs(row['baseline']-baseline_quad(row)) < 1e-13
        summary['stress'] = metrics(stress_defined, best, np.array(fits[best]['coefficients']))
        summary['stress']['scope'] = 'Nonspinning rows outside the declared energy domain; diagnostic failure, not validation or retuning input'
        summary['stress']['unsupported_spin_rows'] = [r['id'] for r in parts['stress'] if r['spin'] > 1e-6]
        summary['stress']['undefined_baseline_rows'] = [r['id'] for r in parts['stress'] if r['j_over_jc'] <= 1 and r['spin'] <= 1e-6]
        summary['stress']['undefined_baseline_scope'] = 'A limitation of the Schwarzschild comparison function, not a comparable-mass or spinning-binary capture criterion'
    output = args.output_dir/('calibration-test.json' if args.mode != 'train' else 'calibration-training.json')
    output.write_text(json.dumps(summary, indent=2)+'\n')
    print(json.dumps({k:summary[k] for k in ('counts','selected')},indent=2))
    for name, report in fits.items():
        print(name, 'condition', report['check']['scaled_design_condition'],
              'train deg', report['training']['rms_degrees'],
              'validation deg', report['validation']['rms_degrees'])
    if args.mode != 'train':
        print('LOCKED TEST', json.dumps(summary['test'],indent=2))


if __name__ == '__main__':
    main()