#!/usr/bin/env python3
"""Independent harmonic recovery from actual N-body events; optional NumPy."""
from pathlib import Path
import argparse
import json
import math
import numpy as np
from chopping_model import coefficients

parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--out", type=Path, required=True)
arguments = parser.parse_args()
record = json.loads(arguments.input.read_text())
periods = []
offsets = []
for events in record["events"]:
    epochs = np.arange(len(events)) + 1
    times = np.array([event["time"] for event in events])
    design = np.array([np.ones(len(events)), epochs]).T
    offset, period = np.linalg.lstsq(design, times, rcond=None)[0]
    offsets.append(offset)
    periods.append(period)
times = np.array([event["time"] for event in record["events"][0]])
epochs = np.arange(len(times)) + 1
measured = []
signed_components = []
for harmonics in [4, 8, 12]:
    columns = [np.ones(len(times)), epochs]
    for harmonic in range(1, harmonics + 1):
        phase = 2 * math.pi * harmonic * (times - offsets[1]) / periods[1]
        columns.extend([np.sin(phase), np.cos(phase)])
    fitted = np.linalg.lstsq(np.array(columns).T, times, rcond=None)[0]
    measured.append(float(math.hypot(fitted[2], fitted[3])))
    signed_components.append([float(fitted[2]), float(fitted[3])])
expected = (
    periods[0]
    / (2 * math.pi)
    * (20 / 332946.0487)
    * abs(coefficients(periods[1] / periods[0])[0][0])
)
relative_error = abs(measured[1] / expected - 1)
# At inner transits the relative longitude is minus the outer clock phase.
# Check this sign and quadrature component, not only an unsigned amplitude.
expected_sine = (
    -periods[0]
    / (2 * math.pi)
    * (20 / 332946.0487)
    * coefficients(periods[1] / periods[0])[0][0]
)
assert abs(signed_components[1][0] / expected_sine - 1) < 0.035
assert abs(signed_components[1][1] / expected_sine) < 0.015
assert relative_error < 0.035
assert abs(measured[2] / measured[1] - 1) < 0.005
result = {
    "kind": "Synthetic full N-body timing events compared to independent perturbative chopping amplitude",
    "fitted_harmonic_counts": [4, 8, 12],
    "fitted_first_harmonic_code_time": measured,
    "perturbation_theory_code_time": expected,
    "phase_convention": "Outer fitted-clock phase; inner relative longitude is its negative at the transit sampling instants",
    "fitted_sine_cosine_code_time": signed_components,
    "perturbation_theory_sine_code_time": expected_sine,
    "fractional_difference": relative_error,
    "note": "Finite mass, finite baseline, free/forced eccentricity and truncated expansions remain; not an exact harmonic identity",
}
arguments.out.write_text(json.dumps(result, indent=2) + "\n")
print(
    "PASS: first synodic component differs by",
    relative_error,
    "from independent perturbation theory; harmonic truncation stable",
)