"""Recompute one spin slice of the published Kerr scalar amplification table and compare.

Usage:
    python3 reproduce_amplification_slice.py AUTHOR_TABLE.dat slice-inputs.json slice-outputs.json

AUTHOR_TABLE.dat is sigma_VS_om_a_s0l1m1.dat from superradiance_spin_1.zip, distributed by the
CENTRA/GRIT amplification-factor page (Brito, Cardoso and Pani). It is not redistributed here; the inputs
file records the archive SHA-256 so a reader can confirm the same download.
Columns of the table: a/M, M omega, Z = dE_out/dE_in - 1. Units G = c = M = 1.
"""
import hashlib
import json
import platform
import sys
import numpy as np
import scipy
from kerr_scalar_scattering import amplification_factor, horizon_quantities, wronskian_along_path
from low_frequency_amplification import low_frequency_amplification


def classify(computed, published, relative_tolerance, roundoff_floor):
    if abs(published) < roundoff_floor:
        return 'roundoff-floor'
    return 'pass' if abs(computed - published) <= relative_tolerance * abs(published) else 'fail'


def main(table_path, inputs_path, output_path):
    with open(inputs_path) as handle:
        inputs = json.load(handle)
    with open(table_path, 'rb') as handle:
        table_bytes = handle.read()
    table = np.loadtxt(table_path)
    spin = inputs['spin']
    rows = table[np.abs(table[:, 0] - spin) < 1e-9]
    if rows.size == 0:
        raise SystemExit(f'No rows with a = {spin} in {table_path}')
    tolerance, floor = inputs['relativeTolerance'], inputs['roundoffFloor']
    r_plus, r_minus, omega_h = horizon_quantities(spin)
    records, failures = [], []
    for _, omega, published in rows:
        details = amplification_factor(omega, spin, inputs['l'], inputs['m'], return_details=True)
        status = classify(details['Z'], published, tolerance, floor)
        # Normalized by the incoming flux, so the comparison stays finite at the threshold omega = m Omega_H.
        flux_mismatch = abs(details['infinityWronskian'] - details['horizonWronskian']) / (omega * abs(details['I']) ** 2)
        record = {'omega': omega, 'publishedZ': published, 'computedZ': details['Z'],
                  'lowFrequencyZ': low_frequency_amplification(omega, spin, inputs['l'], inputs['m']),
                  'status': status, 'fluxBalanceRelativeMismatch': flux_mismatch}
        records.append(record)
        if status == 'fail':
            failures.append(record)
    compared = [r for r in records if r['status'] != 'roundoff-floor']
    relative = np.array([abs(r['computedZ'] - r['publishedZ']) / abs(r['publishedZ']) for r in compared])
    radius_checks = []
    for omega in inputs['outerRadiusCheckFrequencies']:
        base = amplification_factor(omega, spin, inputs['l'], inputs['m'])
        doubled = amplification_factor(omega, spin, inputs['l'], inputs['m'], outer_radius=2 * max(250.0, 80.0 / omega))
        radius_checks.append({'omega': omega, 'Z': base, 'ZWithDoubledOuterRadius': doubled,
                              'relativeChange': abs(doubled / base - 1)})
    wronskian_radii, wronskian_values = wronskian_along_path(inputs['wronskianFrequency'], spin, inputs['l'], inputs['m'])
    zero_spin = [{'omega': w, 'Z': amplification_factor(w, 0.0, inputs['l'], inputs['m'])} for w in inputs['zeroSpinFrequencies']]
    low_frequency = [r for r in records if r['omega'] <= inputs['lowFrequencyCompareBelow']]
    output = {
        'software': {'python': platform.python_version(), 'numpy': np.__version__, 'scipy': scipy.__version__},
        'tableSha256': hashlib.sha256(table_bytes).hexdigest(),
        'tableSha256MatchesInputs': hashlib.sha256(table_bytes).hexdigest() == inputs['authorArchive']['fileSha256'],
        'spin': spin, 'l': inputs['l'], 'm': inputs['m'],
        'horizonAngularVelocity': omega_h, 'superradiantThreshold': inputs['m'] * omega_h,
        'rowsCompared': len(compared), 'rowsAtRoundoffFloor': len(records) - len(compared),
        'medianRelativeDifference': float(np.median(relative)), 'maximumRelativeDifference': float(relative.max()),
        'failures': failures,
        'signAgreement': all(np.sign(r['computedZ']) == np.sign(r['publishedZ']) for r in compared),
        'maximumFluxBalanceMismatch': max(r['fluxBalanceRelativeMismatch'] for r in records),
        'wronskianAlongPath': {'radii': wronskian_radii.tolist(), 'values': wronskian_values.tolist(),
                               'relativeSpread': float(np.ptp(wronskian_values) / abs(wronskian_values.mean()))},
        'outerRadiusChecks': radius_checks,
        'zeroSpinLimit': zero_spin,
        'lowFrequencyComparison': [{'omega': r['omega'], 'publishedZ': r['publishedZ'], 'lowFrequencyZ': r['lowFrequencyZ'],
                                    'ratio': r['lowFrequencyZ'] / r['publishedZ']} for r in low_frequency],
        'rows': records,
    }
    with open(output_path, 'w') as handle:
        json.dump(output, handle, indent=1)
        handle.write('\n')
    print(f"compared {len(compared)} rows, median relative difference {output['medianRelativeDifference']:.2e}, "
          f"max {output['maximumRelativeDifference']:.2e}, failures {len(failures)}")
    return 1 if failures else 0


if __name__ == '__main__':
    sys.exit(main(*sys.argv[1:4]))