#!/usr/bin/env python3
"""Measure secular cycles from frozen IAS15 Cartesian trajectories.
Python3.9+ standard library only;the input is synthetic integration output.
"""
import argparse
import csv
import hashlib
import json
import math
from pathlib import Path


def dot(first, second):
    return sum(a*b for a, b in zip(first, second))


def cross(first, second):
    return [first[1]*second[2]-first[2]*second[1], first[2]*second[0]-first[0]*second[2], first[0]*second[1]-first[1]*second[0]]


def length(vector):
    return math.sqrt(dot(vector, vector))


def elements(row, inclination):
    position, velocity = row[1:4], row[4:7]
    momentum = cross(position, velocity)
    velocity_cross_momentum = cross(velocity, momentum)
    eccentricity = length([value-coordinate/length(position) for value, coordinate in zip(velocity_cross_momentum, position)])
    axis = [0, -math.sin(math.radians(inclination)), math.cos(math.radians(inclination))]
    angle = math.degrees(math.acos(max(-1, min(1, dot(momentum, axis)/length(momentum)))))
    return eccentricity, angle, dot(momentum, axis)


def averaged_flow(state, axis, gr_strength=0):
    momentum, eccentricity = state[:3], state[3:]
    axial_momentum, axial_eccentricity = dot(momentum, axis), dot(eccentricity, axis)
    momentum_cross_axis = cross(momentum, axis)
    eccentricity_cross_axis = cross(eccentricity, axis)
    momentum_cross_eccentricity = cross(momentum, eccentricity)
    return [.75*(axial_momentum*momentum_cross_axis[index]-5*axial_eccentricity*eccentricity_cross_axis[index]) for index in range(3)] + [.75*(2*momentum_cross_eccentricity[index]-5*axial_eccentricity*momentum_cross_axis[index]+axial_momentum*eccentricity_cross_axis[index])+gr_strength*momentum_cross_eccentricity[index]/length(momentum)**3 for index in range(3)]


def periapsis_angle(state, axis):
    momentum, eccentricity = state[:3], state[3:]
    node = cross(axis, momentum)
    transverse = cross(momentum, node)
    return math.atan2(dot(eccentricity, transverse)/length(momentum), dot(eccentricity, node))


def unwrap_angles(angles):
    unwrapped = [angles[0]]
    for previous, current in zip(angles, angles[1:]):
        difference = math.atan2(math.sin(current-previous), math.cos(current-previous))
        unwrapped.append(unwrapped[-1]+difference)
    return [math.degrees(value) for value in unwrapped]


def integrate_average(times, inclination, gr_strength=0):
    axis = [0, -math.sin(math.radians(inclination)), math.cos(math.radians(inclination))]
    state = [0, 0, math.sqrt(1-.05**2), .05, 0, 0]
    tau = 0
    output = []
    for target in times:
        while tau < target-1e-12:
            step = min(.001, target-tau)
            first = averaged_flow(state, axis, gr_strength)
            second = averaged_flow([value+step*slope/2 for value, slope in zip(state, first)], axis, gr_strength)
            third = averaged_flow([value+step*slope/2 for value, slope in zip(state, second)], axis, gr_strength)
            fourth = averaged_flow([value+step*slope for value, slope in zip(state, third)], axis, gr_strength)
            next_state = [value+step*(a+2*b+2*c+d)/6 for value, a, b, c, d in zip(state, first, second, third, fourth)]
            if 1-length(next_state[3:]) <= .03:
                return output
            state = next_state
            tau += step
        output.append((tau, length(state[3:]), math.degrees(math.acos(max(-1, min(1, dot(state[:3], axis)/length(state[:3]))))), periapsis_angle(state, axis)))
    return output


def smoothed_peaks(times, values, outer_period_tau, window_factor=1):
    cadence = times[1]-times[0]
    half_window = max(1, round(window_factor*outer_period_tau/(2*cadence)))
    smooth = [sum(values[max(0,index-half_window):min(len(values),index+half_window+1)])/len(values[max(0,index-half_window):min(len(values),index+half_window+1)]) for index in range(len(values))]
    if max(smooth)-min(smooth)<.2:
        return smooth,[],None
    threshold = min(smooth)+.75*(max(smooth)-min(smooth))
    candidates = [index for index in range(1,len(values)-1) if smooth[index]>smooth[index-1] and smooth[index]>=smooth[index+1] and smooth[index]>threshold]
    groups = []
    for index in candidates:
        if groups and times[index]-times[groups[-1][-1]] < 1:
            groups[-1].append(index)
        else:
            groups.append([index])
    selected = [max(group,key=lambda index:smooth[index]) for group in groups]
    peaks = [times[index] for index in selected]
    periods = [second-first for first,second in zip(peaks,peaks[1:])]
    return smooth, peaks, sum(periods)/len(periods) if periods else None


def main():
    script_directory = Path(__file__).resolve().parent
    repository = next((parent for parent in script_directory.parents if (parent/'src/features/labs/secular-orbits/data/ias15-provenance.json').is_file()),None)
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--data-dir',type=Path,default=repository/'src/features/labs/secular-orbits/data' if repository else script_directory)
    parser.add_argument('--output-dir',type=Path,default=repository/'docs/astronomy/ac09/analysis' if repository else Path.cwd()/'analysis')
    args=parser.parse_args()
    metadata=json.loads((args.data_dir/'ias15-provenance.json').read_text())
    assert metadata['integrator']=='REBOUND5.2.2 IAS15'
    args.output_dir.mkdir(parents=True,exist_ok=True)
    results=[]
    for case in metadata['cases']:
        content=(args.data_dir/case['file']).read_bytes()
        assert hashlib.sha256(content).hexdigest()==case['sha256']
        rows=[[float(value) for value in row] for row in list(csv.reader(content.decode().splitlines()))[1:]]
        times=[row[0]/case['hierarchyRatio']**3 for row in rows]
        measured=[elements(row,case['inclinationDegrees']) for row in rows]
        eccentricities=[item[0] for item in measured]
        average=integrate_average(times,case['inclinationDegrees'])
        common_count=min(len(average),len(times))
        averaged_angles=unwrap_angles([row[3] for row in average])
        outer_period_tau=2*math.pi*math.sqrt(case['hierarchyRatio']**3/2)/case['hierarchyRatio']**3
        smooth,peaks,period=smoothed_peaks(times,eccentricities,outer_period_tau)
        _,average_peaks,average_period=smoothed_peaks([row[0] for row in average],[row[1] for row in average],outer_period_tau)
        period_variations=[smoothed_peaks(times,eccentricities,outer_period_tau,factor)[2] for factor in [.5,1,2]]
        available=[value for value in period_variations if value is not None]
        residuals=[eccentricities[index]-average[index][1] for index in range(common_count)]
        with (args.output_dir/case['file'].replace('ias15-','derived-')).open('w',newline='') as output:
            writer=csv.writer(output,lineterminator='\n');writer.writerow(['tau','direct_e','direct_inclination_deg','direct_h_dot_outer_axis','smoothed_direct_e','averaged_e','averaged_inclination_deg','direct_minus_averaged_e','averaged_unwrapped_periapsis_angle_deg'])
            for index in range(common_count):writer.writerow([times[index],*measured[index],smooth[index],average[index][1],average[index][2],residuals[index],averaged_angles[index]])
        results.append({'file':case['file'],'forceModel':case.get('forceModel','full'),'provenance':'Derived from saved IAS15 Cartesian positions/velocities and separately integrated quadrupole equations','cadence_tau':times[1]-times[0],'outer_period_tau':outer_period_tau,'direct_max_e':max(eccentricities),'averaged_max_e':max(row[1] for row in average),'averaged_periapsis_angle_advance_degrees':averaged_angles[-1]-averaged_angles[0],'direct_peak_times_tau':peaks,'averaged_peak_times_tau':average_peaks,'direct_period_tau':period,'averaged_period_tau':average_period,'period_range_for_half_to_double_smoothing_window': [min(available),max(available)] if available else None,'raw_e_residual_rms':math.sqrt(sum(value*value for value in residuals)/common_count),'max_axial_momentum_change':max(abs(item[2]-measured[0][2]) for item in measured),'direct_stopped':case['stopped'],'direct_completed_tau':times[-1],'averaged_completed_tau':average[-1][0],'interpretation':('No full secular period is measured from a stopped partial cycle.' if case['stopped'] else 'Small-amplitude stable branch; the large-excursion cycle estimator is not applied.') if period is None else 'Period is measured from grouped peaks after one-outer-orbit smoothing;window variation is a method sensitivity, not a statistical confidence interval.'})
    gr_reference=integrate_average([index*.005 for index in range(2401)],65,.8)
    summary={'schema':'sandboxphysics-ac09-period-analysis-v1','units':'G=M_primary=M_companion=a_inner=1;inner period 2pi; time normalized by tK=hierarchyRatio**3','sampling':'Frozen IAS15 output cadence 0.005 tK; phase-sensitive short-period changes remain in raw Cartesian data','smoothing':'Centered moving average of eccentricity over one outer period;peak groups separated by at least 1 tK;one-half/double window comparison. The large-excursion estimator requires eccentricity range at least 0.2.','scope':'Synthetic dynamical benchmark, not observed exoplanet measurements. Direct references are Newtonian. Conservative GR precession is only added to the averaged model in the separate GR comparison.','cases':results,'controlled_comparisons': [{'hierarchyRatio':hierarchy,'inclinationDegrees':75,'full_newtonian_period_tau':next(item['direct_period_tau'] for item in results if item['file']==full_name),'instantaneous_quadrupole_period_tau':next(item['direct_period_tau'] for item in results if item['file']==quad_name),'doubly_averaged_quadrupole_period_tau':next(item['averaged_period_tau'] for item in results if item['file']==full_name),'interpretation':'Instantaneous quadrupole versus double average changes averaging alone;full Newtonian versus instantaneous quadrupole changes multipole truncation. The two period shifts need not add linearly for long nonlinear evolution.'} for hierarchy,full_name,quad_name in [(6,'ias15-compact.csv','ias15-quadrupole-compact.csv'),(20,'ias15-wide-inclined.csv','ias15-quadrupole-wide.csv')]],'averaged_GR_example':{'inclinationDegrees':65,'initialEccentricity':.05,'epsilon_GR':.8,'max_e':max(row[1] for row in gr_reference),'direct_GR_integration':False},'limitations':['Finite initial eccentricity 0.05 differs from the initially circular analytic limit','Outer-orbit averaging removes real fast oscillations;compare filtered trends and raw records','The direct inner two-body energy is not conserved; circular rotating Hamiltonian is the relevant invariant','Near-primary exclusion stops partial records; never extend a stopped curve to the requested duration','No dissipative tides or orbital migration. Full Newtonian cases retain the full companion force;instantaneous-quadrupole controls and averaged curves both truncate its expansion at quadrupole order.']}
    (args.output_dir/'period-analysis.json').write_text(json.dumps(summary,indent=2)+'\n')
    print(json.dumps(summary,indent=2));print('PASS input hashes,actual-state element extraction,averaged integration and frozen period/failure analysis')


if __name__=='__main__':
    main()