"""Subtract official scalar parameters and fit the unresolved high-l tail.

All forces are covariant Schwarzschild components for G=c=M=q_s=1. Parameter
entries already include their complete l dependence.  Higher parameters
have zero sum from l=0 to infinity; they accelerate convergence but do not
change the force.  Fits use the next rational mode-sum bases, not raw powers.
"""
import argparse
import json
import math
import platform
import re
import sys
from pathlib import Path

import mpmath as mp

mp.mp.dps = 70


def number(value):
    if isinstance(value, str):
        value = re.sub(r'`{1,2}[-+0-9.]*', '', value).replace('*^', 'e')
    return mp.mpf(value)


def basis(ell, order):
    """b_k(l)=product_{j=1}^k[(2l+1)^2-4j^2]^{-1}."""
    return 1 / mp.fprod((2 * mp.mpf(ell) + 1)**2 - 4*j*j
                        for j in range(1, order + 1))


def basis_tail(cutoff, order):
    """Exact tail: T_k(N-1)-T_k(N)=b_k(N), with T_k(infinity)=0."""
    n=mp.mpf(cutoff)
    return (n+1)/((2*order-1)*mp.fprod(
        (2*n+2)**2-(2*j-1)**2 for j in range(1,order+1)))


def fit_tail(values, cutoff, start, stop, count=3, first_order=4):
    orders = range(first_order, first_order+count)
    norms = [basis(start, k) for k in orders]
    matrix = mp.matrix([[basis(ell, k)/scale for k, scale in zip(orders, norms)]
                        for ell in range(start, stop+1)])
    rhs = mp.matrix(values[start:stop+1])
    scaled, residual = mp.qr_solve(matrix, rhs)
    coefficients = [c/scale for c, scale in zip(scaled, norms)]
    tail = mp.fsum(c*basis_tail(cutoff, k) for c, k in zip(coefficients, orders))
    return dict(start=start, stop=stop, terms=count, first_order=first_order,
                tail=tail, force=mp.fsum(values[:cutoff+1])+tail,
                fit_residual=mp.sqrt(mp.fsum(
                    (values[ell]-mp.fsum(c*basis(ell,k) for c,k in
                    zip(coefficients,orders)))**2 for ell in range(start,stop+1)) /
                    (stop-start+1)), coefficients=coefficients)


def slope_fit(rows,key,start,stop):
    """Ordinary least-squares slope in log L/log absolute residual."""
    xs=[mp.log(mp.mpf(ell)+mp.mpf('.5')) for ell in range(start,stop+1)]
    ys=[mp.log(abs(rows[ell][key])) for ell in range(start,stop+1)]
    xbar,ybar=mp.fsum(xs)/len(xs),mp.fsum(ys)/len(ys)
    return mp.fsum((x-xbar)*(y-ybar) for x,y in zip(xs,ys))/mp.fsum((x-xbar)**2 for x in xs)


def load_data(raw_path, parameter_path):
    raw = json.loads(raw_path.read_text())
    return regularize(raw, parameter_path)


def regularize(raw, parameter_path):
    if number(raw['radius']) not in (6,10):
        raise ValueError('Stored parameters apply only at exact r/M=6 or 10')
    radius = int(raw['radius'])
    parameters = next(v for v in json.loads(parameter_path.read_text())
                      if v['r_over_M'] == radius)
    terms = [[number(t) for t in row['terms']] for row in parameters['modes']]
    rows = []
    for row in raw['multipoles']:
        ell = row['ell']
        a, b, f2, f4, f6 = terms[ell]
        plus, minus = number(row['force_plus']), number(row['force_minus'])
        average = (plus+minus)/2
        rows.append(dict(ell=ell, plus=plus, minus=minus, average=average,
                         A=a, B=b, F2=f2, F4=f4, F6=f6,
                         AB=average-b, AB2=average-b-f2,
                         AB24=average-b-f2-f4, AB246=average-b-f2-f4-f6,
                         regularized_plus=plus-a-b-f2-f4-f6,
                         regularized_minus=minus+a-b-f2-f4-f6))
    return raw, rows


def serialize(value):
    if isinstance(value, (mp.mpf,mp.mpc)):
        return mp.nstr(value, 45)
    if isinstance(value, dict):
        return {k: serialize(v) for k,v in value.items()}
    if isinstance(value, list):
        return [serialize(v) for v in value]
    return value


def comparison_suite(directory, parameter_path):
    """Recompute the published checks from two precision runs and fast modes.

    The allowances describe observed convergence, not total validated errors;
    stable or shared solver bias is not bounded. Failing a refinement check
    must not silently widen them. Both precision runs use the same 41 modes.
    """
    archive_path=directory/'checks.json'
    archive=json.loads(archive_path.read_text()) if archive_path.exists() else None
    output = dict(
        normalization='G=c=M=q_s=1; covariant Schwarzschild F_r=q_s partial_r Phi_R',
        scope='Massless minimally coupled scalar; circular geodesics; retarded boundaries',
        selected_tail=dict(cutoff=40, start=25, stop=40, first_order=4, terms=4),
        uncertainty_method='The estimated_abs_error fields are empirical convergence allowances exceeding observed precision, cutoff, fit-window and fit-order changes, not total validated errors or interval bounds. Stable or shared solver bias is not bounded. At r/M=6, the independent thirty-mode reference is rounded at 1e-14 and the cited full-force reference has seven significant figures; neither independently validates the 2e-15 allowance as a total error estimate.',
        recorded_solver_environment=dict(python='3.12.14', platform='macOS ARM64',
                         pybhpt='1.0.0', numpy='2.2.6', scipy='1.18.1', mpmath='1.3.0',
                         wolfram='15.0.0'),
        analysis_environment=dict(python=sys.version,platform=platform.platform(),
                                  mpmath=mp.__version__),
        source_revisions={
            'Teukolsky':'39ee4fb7266e780e0b6d4ae0c4a0556c91323264',
            'KerrGeodesics':'c7c57a936236ddcce98963481440bca0c433b3c3',
            'SpinWeightedSpheroidalHarmonics':'48cd35591f25ee5910445abad5f18d2b33c140ad',
            'RegularizationParameters':'ad85c6f0088bc47f3a0a2ab22c8d1fd5b36b8681',
            'pybhpt_source_inspection_only':'5e853f9b653e1054af8e36eafe7d179db0b233dd'},
        runtime_wheel=dict(name='pybhpt-1.0.0-cp312-cp312-macosx_14_0_arm64.whl',
            sha256='2c588dabd7f7ed2d678969425d4c3b19d8d974cc88d552df706b62507e996dd1',
            note='The separately inspected Git source is not asserted to be this wheel build commit.'),
        radii={})
    for radius, allowance, published in [(6,'2e-15','0.00016772830804'),
                                          (10,'2e-18','0.000013784482234')]:
        high, rows = load_data(directory/f'modes-r{radius}-wp55.json', parameter_path)
        low_path=directory/f'modes-r{radius}-wp35.json'
        if low_path.exists():
            low, low_rows = load_data(low_path, parameter_path)
        else:
            # Archived comparison input: recompute all subtractions from the
            # original one-sided numbers, not from saved residuals or fits.
            saved=archive['radii'][str(radius)]['precision35_multipoles']
            low,low_rows=regularize(dict(radius=radius,base_precision=35,
                multipoles=[dict(ell=r['ell'],force_plus=r['plus'],force_minus=r['minus'])
                            for r in saved]),parameter_path)
        assert high['base_precision']==55 and low['base_precision']==35
        assert len(rows)>=41 and len(low_rows)>=41
        rows, low_rows = rows[:41], low_rows[:41]
        values = [r['AB246'] for r in rows]
        selected = fit_tail(values,40,25,40,4)
        low_fit = fit_tail([r['AB246'] for r in low_rows],40,25,40,4)
        fits = [fit_tail(values,40,a,b,n) for a,b,n in
                [(25,40,3),(25,40,4),(25,40,5),(20,35,4),(30,40,4)]]
        cutoff_fit = fit_tail(values,35,20,35,4)
        variations = dict(
            precision=abs(selected['force']-low_fit['force']),
            cutoff_35_to_40=abs(fits[3]['force']-cutoff_fit['force']),
            fit_terms=max(abs(f['force']-selected['force']) for f in fits[:3]),
            fit_windows=max(abs(f['force']-selected['force']) for f in fits[3:]))
        assert max(variations.values()) < mp.mpf(allowance), variations
        sides = {key:mp.fsum(r[key] for r in rows)+selected['tail']
                 for key in ['regularized_plus','regularized_minus']}
        assert abs(sides['regularized_plus']-sides['regularized_minus']) < mp.mpf('1e-45')
        slopes = {f'{a}:{b}':{key:slope_fit(rows,key,a,b)
                  for key in ['AB','AB2','AB24','AB246']}
                  for a,b in [(20,35),(25,40)]}
        for window in slopes.values():
            assert abs(window['AB']+2)<mp.mpf('.15')
            assert abs(window['AB2']+4)<mp.mpf('.15')
        # A wrong outer-side sign must leave a growing, not decaying, mode.
        wrong_a = rows[-1]['plus']+rows[-1]['A']-rows[-1]['B']
        assert abs(wrong_a)>abs(rows[-1]['A'])
        fast_path = directory/f'modes-r{radius}-flux.json'
        if not fast_path.exists(): fast_path=directory/f'modes-r{radius}-auto.json'
        if fast_path.exists():
            fast, fast_rows = load_data(fast_path,parameter_path)
        else:
            saved=archive['radii'][str(radius)]['double_precision']
            fast,fast_rows=regularize(dict(radius=radius,
                multipoles=saved['raw_multipoles'],flux_check=saved['flux_check'],
                modes=[dict(wronskian_relative_drift=saved['maximum_Wronskian_relative_drift'])]),
                parameter_path)
        assert len(fast_rows)>=81
        flux=fast['flux_check']
        flux=dict(flux)
        flux['infinity']=math.fsum(m['infinity'] for m in flux['modes'])
        flux['horizon']=math.fsum(m['horizon'] for m in flux['modes'])
        flux['local_Ft_over_ut']=math.sqrt(1-3/radius)*math.fsum(
            r['force_time'] for r in fast['multipoles'][:flux['lmax']+1])
        flux['difference']=flux['local_Ft_over_ut']-flux['infinity']-flux['horizon']
        flux['relative_difference']=flux['difference']/(flux['infinity']+flux['horizon'])
        assert flux['lmax']==30 and abs(flux['relative_difference'])<1e-10
        wronskian=max(m['wronskian_relative_drift'] for m in fast['modes'])
        assert wronskian<1e-10
        double_sums={str(n):mp.fsum(r['AB246'] for r in fast_rows[:n+1])
                     for n in [40,60,80]}
        partial29=mp.fsum(values[:30])
        assert abs(partial29-number(published))<mp.mpf('5e-15')
        jumps=[]
        for mode in high['modes']:
            ell,emm=mode['ell'],mode['m']
            if ell>40: continue
            angular=(2*ell+1)*mp.binomial(ell+emm,(ell+emm)//2)*mp.binomial(
                ell-emm,(ell-emm)//2)/mp.mpf(2)**(2*ell)
            plus=mp.mpc(*map(number,mode['force_plus']))
            minus=mp.mpc(*map(number,mode['force_minus']))
            jumps.append(abs(plus-minus+angular*mp.sqrt(1-mp.mpf(3)/radius)/
                             (radius*(radius-2))))
        assert max(jumps)<mp.mpf('1e-45')
        result=dict(radius=radius,lmax=40,accepted_force=selected['force'],
            estimated_abs_error=mp.mpf(allowance), selected_fit=selected,
            refinement_changes=variations, alternative_fits=fits,
            cutoff35_fit=cutoff_fit, precision35_fit=low_fit,
            maximum_precision_mode_change=max(abs(a['average']-b['average'])
                                             for a,b in zip(rows,low_rows)),
            maximum_high_precision_source_jump_residual=max(jumps),
            regularized_sides=sides, regression_slopes=slopes,
            wrong_A_negative_control_at_lmax=wrong_a,
            published_30_mode_comparison=dict(reference='Warburton and Wardell 2014, Table I, p.10',
                cutoff=29,computed=partial29,published=published,difference=partial29-number(published)),
            double_precision=dict(lmax=80, AB246_partial_sums=double_sums,
                raw_multipoles=fast['multipoles'],
                maximum_Wronskian_relative_drift=wronskian,
                maximum_source_jump_residual=max(abs(m['jump_residual']) for m in fast['multipoles']),
                flux_check=flux,
                note='Double precision resolves the general convergence trend but is not the source of the high-precision error allowance.'),
            multipoles=rows, precision35_multipoles=low_rows,
            high_precision_method_counts={method:sum(m['method']==method and m['ell']<=40
                    for m in high['modes']) for method in ['Static','MST','NumericalIntegration']})
        if radius==10:
            ref=number('0.000013784482575667959')
            result['published_infinite_sum_comparison']=dict(
                reference='Heffernan, Ottewill and Wardell 2012, section V F, p.27',
                published=ref,absolute_difference=abs(selected['force']-ref),
                note='Their parenthetical tail-only uncertainty is not a total numerical error bound.')
            assert abs(selected['force']-ref)<mp.mpf(allowance)
        output['radii'][str(radius)]=result
    return output


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('raw', type=Path, nargs='?')
    parser.add_argument('--compare-directory', type=Path,
                        help='Compare paired precision and fast runs; write the complete checks.json')
    parser.add_argument('--parameters', type=Path,
                        default=Path(__file__).with_name('scalar-parameters-60digits.json'))
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    if args.compare_directory:
        result=comparison_suite(args.compare_directory,args.parameters)
        args.output.write_text(json.dumps(serialize(result),indent=2)+'\n')
        for radius,entry in result['radii'].items():
            print(radius,mp.nstr(entry['accepted_force'],25),
                  'empirical convergence allowance',entry['estimated_abs_error'])
        return
    if args.raw is None: parser.error('Supply a raw mode file or --compare-directory')
    raw, rows = load_data(args.raw, args.parameters)
    lmax = len(rows)-1
    # Independently verify the exact tail formula against direct summation.
    for k in range(1,7):
        direct = mp.fsum(basis(ell,k) for ell in range(11,1001))
        assert abs(direct-(basis_tail(10,k)-basis_tail(1000,k))) < mp.mpf('1e-65')
        assert abs(basis_tail(-1,k)) < mp.mpf('1e-65')
    # The notebook entries include the exact rational basis, not just its
    # leading inverse power. Check that identification before tail fitting.
    for order,key in enumerate(['F2','F4','F6'],1):
        coefficient = rows[0][key]/basis(0,order)
        assert max(abs(row[key]/basis(row['ell'],order)/coefficient-1)
                   for row in rows) < mp.mpf('1e-55')
    sums = {key: mp.fsum(row[key] for row in rows)
            for key in ['AB','AB2','AB24','AB246','regularized_plus','regularized_minus']}
    fits=[]
    values=[row['AB246'] for row in rows]
    for start,stop in [(15,25),(20,30),(20,35),(25,35),(25,40),(30,40),(35,45),(40,50),(50,60),
                       (40,60),(45,65),(50,70),(55,75),(60,80)]:
        if stop>lmax: continue
        for count in [3,4,5]:
            fits.append(fit_tail(values,lmax,start,stop,count))
    cutoff_sums = {cutoff: {key: mp.fsum(row[key] for row in rows[:cutoff+1])
                          for key in ['AB','AB2','AB24','AB246']}
                   for cutoff in [10,20,29,30,35,40,60,80] if cutoff<=lmax}
    slopes = {f'{start}:{stop}': {
                key: mp.log(abs(rows[stop][key]/rows[start][key])) /
                mp.log((mp.mpf(stop)+mp.mpf('.5'))/(mp.mpf(start)+mp.mpf('.5')))
                for key in ['AB','AB2','AB24','AB246']}
              for start,stop in [(15,25),(20,30),(25,35),(30,40)] if stop<=lmax}
    regression_slopes={f'{start}:{stop}': {key:slope_fit(rows,key,start,stop)
                       for key in ['AB','AB2','AB24','AB246']}
                       for start,stop in [(20,35),(25,40),(20,40),(26,40),
                                          (51,70),(56,75),(61,80),(56,80),(66,80)]
                       if stop<=lmax}
    cutoff_tail_checks=[]
    for cutoff in [30,35,40,60,80]:
        if cutoff>lmax:continue
        for count in [2,3,4]:
            cutoff_tail_checks.append(dict(cutoff=cutoff,**fit_tail(
                values,cutoff,cutoff-15,cutoff,count)))
    wrong_a = rows[-1]['plus']+rows[-1]['A']-rows[-1]['B']
    assert abs(wrong_a)>abs(rows[-1]['A'])
    result=dict(raw_file=str(args.raw), radius=raw['radius'], lmax=lmax,
                normalization='G=c=M=q_s=1; covariant Schwarzschild Fr',
                partial_sums=sums, cutoff_sums=cutoff_sums, slopes=slopes,
                regression_slopes=regression_slopes, cutoff_tail_checks=cutoff_tail_checks,
                wrong_A_negative_control_at_lmax=wrong_a,
                fits=fits, multipoles=rows)
    args.output.write_text(json.dumps(serialize(result),indent=2)+'\n')
    print('r=',raw['radius'],'lmax=',lmax,'partial=',mp.nstr(sums['AB246'],24))
    print('side_difference=',mp.nstr(sums['regularized_plus']-sums['regularized_minus'],8))
    for f in fits:
        print(f['start'],f['stop'],f['terms'],mp.nstr(f['force'],25),
              'tail',mp.nstr(f['tail'],7),'fit_rms',mp.nstr(f['fit_residual'],5))
    if 'force_time' in raw['multipoles'][0]:
        print('Ft=',mp.nstr(mp.fsum(number(row['force_time']) for row in raw['multipoles']),22))


if __name__=='__main__':
    main()