#!/usr/bin/env python3
"""Recompute scattering angles from versioned SXS horizon trajectories.

Requires numpy==2.5.3 and h5py==3.16.0. Raw inputs are fetched only with
--fetch, cached locally, and checked against inputs.json on every run.
No published scattering angles or calibrated EOB predictions enter the fits.
"""
import argparse
import hashlib
import importlib.metadata
import json
import math
from pathlib import Path
import platform
import shutil
import subprocess
import tempfile
import urllib.error
import urllib.request

import h5py
import numpy as np

HERE = Path(__file__).resolve().parent
IN_WINDOWS = tuple(range(50, 81, 5))
OUT_WINDOWS = tuple(range(80, 231, 10))
MODELS = ('cubic', 'kepler')


def require(condition, message):
    if not condition:
        raise ValueError(message)


def verify_file(path, entry):
    data = path.read_bytes()
    require(len(data) == entry['bytes'], f'Wrong file size: {path}')
    require(hashlib.sha256(data).hexdigest() == entry['sha256'],
            f'SHA256 mismatch: {path}; existing files are never silently replaced')
    require(hashlib.md5(data).hexdigest() == entry['md5'],
            f'Catalog MD5 mismatch: {path}')


def prepare_inputs(manifest, cache, fetch):
    """Download into a temporary file, verify, then atomically populate cache."""
    for case in manifest['cases']:
        for entry in case['files']:
            path = cache / entry['cache_path']
            if not path.exists():
                if not fetch:
                    raise FileNotFoundError(f'{path}: rerun with --fetch to obtain original input')
                path.parent.mkdir(parents=True, exist_ok=True)
                with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as stream:
                    temporary = Path(stream.name)
                try:
                    request = urllib.request.Request(entry['url'], headers={'User-Agent': 'Mozilla/5.0'})
                    try:
                        with urllib.request.urlopen(request, timeout=60) as response:
                            with temporary.open('wb') as stream:
                                shutil.copyfileobj(response, stream)
                    except (urllib.error.URLError, TimeoutError):
                        # Some public repository redirects reject urllib's TLS client.
                        # curl follows those same public URLs without credentials.
                        if not shutil.which('curl'):
                            raise
                        subprocess.run(['curl', '--fail', '--location', '--silent', '--show-error',
                                        '--connect-timeout', '30', '--max-time', '120',
                                        '--user-agent', 'Mozilla/5.0', entry['url'],
                                        '--output', str(temporary)], check=True)
                    verify_file(temporary, entry)
                    temporary.replace(path)
                finally:
                    temporary.unlink(missing_ok=True)
            verify_file(path, entry)


def fit_asymptote(distance, phi, model, endpoint):
    """Fit a far-separation intercept; angles are unwrapped, in radians."""
    if model == 'cubic':
        z = 100.0 / distance  # Column scaling leaves the polynomial unchanged.
        matrix = np.column_stack([np.ones_like(z), z, z*z, z*z*z])
        coefficients, _, rank, singular = np.linalg.lstsq(matrix, phi, rcond=None)
        require(rank == 4, 'Rank-deficient cubic fit')
        return float(coefficients[0]), {
            'condition': float(singular[0]/singular[-1]),
            'fit_rms_rad': float(np.sqrt(np.mean((phi-matrix@coefficients)**2)))}
    require(model == 'kepler', 'Unknown fitting model')
    matrix = np.column_stack([np.cos(phi), np.sin(phi), np.ones_like(phi)])
    coefficients, _, rank, singular = np.linalg.lstsq(matrix, 1.0/distance, rcond=None)
    require(rank == 3, 'Rank-deficient Kepler fit')
    a, b, c = coefficients
    magnitude = np.hypot(a, b)
    require(abs(c/magnitude) < 1, 'Fitted Kepler branch is not hyperbolic')
    phase = np.arctan2(b, a)
    alpha = np.arccos(-c/magnitude)
    # A cos(phi)+B sin(phi)=sqrt(A^2+B^2) cos(phi-atan2(B,A)).
    roots = []
    for sign in (-1, 1):
        root = phase + sign*alpha
        roots.append(root + 2*np.pi*round((endpoint-root)/(2*np.pi)))
    root = min(roots, key=lambda value: abs(value-endpoint))
    require(abs(a*np.cos(root)+b*np.sin(root)+c) < 1e-12,
            'Kepler asymptote does not solve its fitted equation')
    return float(root), {
        'condition': float(singular[0]/singular[-1]),
        'fit_rms_inverse_M': float(np.sqrt(np.mean((1/distance-matrix@coefficients)**2))),
        'eccentricity': float(magnitude/abs(c))}


def pm_angles(gamma, j, Gamma):
    """Full nonspinning angles in radians through first/second PM order."""
    first = 2*(2*gamma**2-1)/(j*math.sqrt(gamma**2-1))
    second_increment = 3*math.pi*(5*gamma**2-1)/(4*Gamma*j*j)
    return first, second_increment


def analytic_checks():
    distance = np.linspace(50., 200., 101)
    phi = .31 + 2/distance - 7/distance**2 + 4/distance**3
    cubic, _ = fit_asymptote(distance, phi, 'cubic', float(phi[-1]))
    require(abs(cubic-.31) < 2e-14, 'Known cubic intercept failed')
    eccentricity, semilatus, phase = 2., 10., 1.17
    alpha = np.arccos(-1/eccentricity)
    roots = []
    for sign in (-1, 1):
        angles = phase + sign*np.linspace(alpha-.2, alpha-.03, 101)
        distance = semilatus/(1+eccentricity*np.cos(angles-phase))
        root, _ = fit_asymptote(distance, angles, 'kepler', float(angles[-1]))
        require(abs(root-(phase+sign*alpha)) < 1e-13, 'Rotated Kepler root failed')
        roots.append(root)
    kepler = roots[1]-roots[0]-np.pi
    require(abs(kepler-2*np.arcsin(1/eccentricity)) < 1e-13,
            'Known Kepler scattering angle failed')
    # With j=b sqrt(gamma^2-1)/Gamma, check an independent b-form of both terms.
    pm_error = 0.
    for gamma in (1.02, 1.1, 2.):
        for nu in (0., .16, .25):
            Gamma = math.sqrt(1+2*nu*(gamma-1))
            b = 100.
            j = b*math.sqrt(gamma*gamma-1)/Gamma
            first, second = pm_angles(gamma, j, Gamma)
            first_b = 2*Gamma*(2*gamma*gamma-1)/(b*(gamma*gamma-1))
            second_b = 3*math.pi*Gamma*(5*gamma*gamma-1)/(4*b*b*(gamma*gamma-1))
            pm_error = max(pm_error, abs(first-first_b), abs(second-second_b))
    require(pm_error < 1e-14, 'PM angular-momentum/impact-parameter conversion failed')
    return {'cubic_intercept_error_rad': abs(cubic-.31),
            'rotated_kepler_angle_error_rad': float(abs(kepler-2*np.arcsin(.5))),
            'pm_impact_parameter_identity_error_rad': pm_error}


def normalized_inputs(metadata):
    m1, m2 = metadata['initial_mass1'], metadata['initial_mass2']
    mass = m1+m2
    nu = m1*m2/mass**2
    Gamma = metadata['initial_ADM_energy']/mass
    gamma = (Gamma*Gamma-1)/(2*nu)+1
    spin_z = (m1*m1*metadata['initial_dimensionless_spin1'][2]
              +m2*m2*metadata['initial_dimensionless_spin2'][2])
    j = (metadata['initial_ADM_angular_momentum'][2]-spin_z)/(m1*m2)
    first, second = pm_angles(gamma, j, Gamma)
    require(gamma > 1 and j > 0, 'This example requires unbound prograde inputs')
    return {
        'M_initial_code_units': mass, 'mass_ratio_m1_over_m2': m1/m2,
        'nu': nu, 'Gamma_ADM_over_M': Gamma, 'gamma': gamma,
        'J_ADM_z_over_M2': metadata['initial_ADM_angular_momentum'][2]/mass**2,
        'spin_z_over_M2': spin_z/mass**2, 'j_orbital': j,
        'initial_ADM_momentum_norm_over_M':
            float(np.linalg.norm(metadata['initial_ADM_linear_momentum'])/mass),
        'maximum_initial_dimensionless_spin_component':
            max(abs(x) for name in ('initial_dimensionless_spin1', 'initial_dimensionless_spin2')
                for x in metadata[name]),
        'chi_1PM_deg': math.degrees(first),
        'chi_2PM_increment_deg': math.degrees(second),
        'chi_through_2PM_deg': math.degrees(first+second)}


def analyze_case(case, cache, full_fits):
    directory = cache / case['cache_directory']
    result = {'id': case['id'], 'doi': case['doi'], 'resolutions': []}
    all_positions = []
    for lev in (1, 2, 3):
        metadata = json.loads((directory/f'Lev{lev}-metadata.json').read_text())
        require(metadata['spells_revision'] == case['spells_revision'],
                'Unexpected source simulation revision')
        inputs = normalized_inputs(metadata)
        mass = inputs['M_initial_code_units']
        with h5py.File(directory/f'Lev{lev}-Horizons.h5') as stream:
            positions = [stream[f'Ah{hole}.dir/CoordCenterInertial.dat'][:] for hole in 'AB']
        require(all(np.all(np.isfinite(x)) for x in positions), 'Nonfinite trajectory input')
        require(np.array_equal(positions[0][:, 0], positions[1][:, 0]),
                'Horizon coordinate times differ')
        all_positions.append(positions[0])
        time = positions[0][:, 0]
        require(np.all(np.diff(time)>0), 'Times are not strictly increasing')
        distance = np.linalg.norm(positions[0][:, 1:]-positions[1][:, 1:], axis=1)/mass
        peri = np.argmin(distance)
        require(np.all(np.diff(distance[:peri+1])<0) and np.all(np.diff(distance[peri:])>0),
                'This example requires one incoming and one outgoing leg')
        require(distance[0]>200 and distance[-1]>350, 'Trajectory does not cover the fitting windows')
        resolution = {
            'lev': lev, 'initial_inputs_and_PM': inputs,
            'reference_time_code_units': metadata['reference_time'],
            'time_range_code_units': [float(time[0]), float(time[-1])],
            'points': len(time), 'D_initial_min_final_over_M':
                [float(distance[0]), float(distance[peri]), float(distance[-1])],
            'metadata_pbj': metadata.get('pbj_info'), 'holes': []}
        paired_values = []
        for hole, xyz in zip('AB', positions):
            phi = np.unwrap(np.arctan2(xyz[:, 2], xyz[:, 1]))
            require(np.all(np.diff(phi)>0), 'This example expects prograde nearly planar trajectories')
            fits = {'in': {}, 'out': {}}
            legs = [('in', IN_WINDOWS, 200, np.arange(len(time))<peri),
                    ('out', OUT_WINDOWS, 350, np.arange(len(time))>peri)]
            for leg, lowers, upper, side in legs:
                for model in MODELS:
                    fits[leg][model] = []
                    for lower in lowers:
                        mask = side & (distance>=lower) & (distance<=upper)
                        d, angles = distance[mask], phi[mask]
                        require(len(d)>30, 'Too few samples in fitting window')
                        angle, diagnostics = fit_asymptote(d, angles, model, float(angles[np.argmax(d)]))
                        fits[leg][model].append({'lower': lower, 'upper': upper, 'samples': len(d),
                                                'asymptotic_phi_rad': angle, **diagnostics})
            combinations = []
            for model in MODELS:
                for ingoing in fits['in'][model]:
                    for outgoing in fits['out'][model]:
                        theta = math.degrees(outgoing['asymptotic_phi_rad']-ingoing['asymptotic_phi_rad']-np.pi)
                        combinations.append([model, ingoing['lower'], outgoing['lower'], theta])
            values = np.array([x[3] for x in combinations])
            paired_values.append(values)
            coarse = np.array([x[3] for x in combinations if x[1]%10==0 and (x[2]-80)%30==0])
            sampling = []
            for stride in (2, 4):
                sampled = []
                for model in MODELS:
                    angles_by_leg = {}
                    for leg, lowers, upper, side in legs:
                        angles_by_leg[leg] = []
                        for lower in lowers:
                            indices = np.flatnonzero(side & (distance>=lower) & (distance<=upper))[::stride]
                            d, angles = distance[indices], phi[indices]
                            angle, _ = fit_asymptote(d, angles, model, float(angles[np.argmax(d)]))
                            angles_by_leg[leg].append(angle)
                    sampled.extend(math.degrees(out-inc-np.pi) for inc in angles_by_leg['in']
                                   for out in angles_by_leg['out'])
                sampled = np.array(sampled)
                sampling.append({'stride': stride, 'mean_deg': float(sampled.mean()),
                                 'mean_difference_from_full_deg': float(sampled.mean()-values.mean()),
                                 'maximum_same_window_difference_deg': float(np.max(abs(sampled-values)))})
            summary = {
                'hole': hole, 'phi_initial_peri_final_rad':
                    [float(phi[0]), float(phi[peri]), float(phi[-1])],
                'z_variation_over_M': float(np.ptp(xyz[:, 3])/mass),
                'window_model_mean_deg': float(values.mean()),
                'range_deg': [float(values.min()), float(values.max())],
                'model_means_deg': {model: float(np.mean([x[3] for x in combinations if x[0]==model]))
                                    for model in MODELS},
                'coarse_window_mean_deg': float(coarse.mean()),
                'window_grid_mean_difference_deg': float(values.mean()-coarse.mean()),
                'sampling_checks': sampling,
                'angle_columns': ['model', 'incoming_lower_D_over_M', 'outgoing_lower_D_over_M', 'theta_deg'],
                'angles': combinations}
            if full_fits:
                summary['leg_fit_diagnostics'] = fits
            resolution['holes'].append(summary)
        paired = np.mean(paired_values, axis=0)
        resolution['average_angle_deg'] = float(paired.mean())
        resolution['average_angle_range_deg'] = [float(paired.min()), float(paired.max())]
        for label in ('chi_1PM_deg', 'chi_through_2PM_deg'):
            resolution[label+'_minus_extracted_average_deg'] = inputs[label]-resolution['average_angle_deg']
        result['resolutions'].append(resolution)
    values = [r['average_angle_deg'] for r in result['resolutions']]
    result['successive_resolution_differences_deg'] = [values[1]-values[0], values[2]-values[1]]
    result['shared_prefix_checks'] = []
    for index in (0, 1):
        length = min(len(all_positions[index]), len(all_positions[2]))
        require(np.array_equal(all_positions[index][:length, 0], all_positions[2][:length, 0]),
                'Cannot compare resolution prefixes on different time grids')
        differences = np.max(abs(all_positions[index][:length, 1:]-all_positions[2][:length, 1:]), axis=1)
        changed = np.flatnonzero(differences>0)
        require(len(changed)>0, 'Distinct resolution files have identical coordinates')
        first = int(changed[0])
        result['shared_prefix_checks'].append({
            'levels': [index+1, 3], 'first_coordinate_difference_time': float(all_positions[index][first, 0]),
            'bitwise_identical_prefix_samples': first})
    return result


def check_finite(value):
    if isinstance(value, dict):
        for item in value.values():
            check_finite(item)
    elif isinstance(value, (list, tuple)):
        for item in value:
            check_finite(item)
    elif isinstance(value, float):
        require(math.isfinite(value), 'Nonfinite generated result')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--fetch', action='store_true', help='Download missing original inputs')
    parser.add_argument('--cache', type=Path, default=Path('.sxs-scattering-cache'))
    parser.add_argument('--output', type=Path, default=Path('sxs-scattering-results.json'))
    parser.add_argument('--full-fits', action='store_true', help='Include individual leg-fit residuals and conditions')
    args = parser.parse_args()
    manifest_path = HERE/'inputs.json'
    manifest = json.loads(manifest_path.read_text())
    require(tuple(manifest['method']['incoming_lower_D_over_M']) == IN_WINDOWS
            and tuple(manifest['method']['outgoing_lower_D_over_M']) == OUT_WINDOWS
            and manifest['method']['incoming_upper_D_over_M'] == 200
            and manifest['method']['outgoing_upper_D_over_M'] == 350,
            'Manifest fitting windows disagree with this script')
    prepare_inputs(manifest, args.cache, args.fetch)
    report = {
        'schema_version': 1,
        'script_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        'inputs_manifest_sha256': hashlib.sha256(manifest_path.read_bytes()).hexdigest(),
        'environment': {'python': platform.python_version(),
                        **{name: importlib.metadata.version(name) for name in ('numpy', 'h5py')}},
        'method': manifest['method'], 'scope': manifest['scope'],
        'analytic_checks': analytic_checks(),
        'cases': [analyze_case(case, args.cache, args.full_fits) for case in manifest['cases']]}
    check_finite(report)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2, allow_nan=False)+'\n')
    for case in report['cases']:
        print(case['id'], 'mean angles by resolution (degrees):',
              [r['average_angle_deg'] for r in case['resolutions']])
    print('Wrote', args.output)


if __name__ == '__main__':
    main()