#!/usr/bin/env python3
"""Pack a full IAS15 reference, preserving unwrapped phase before downsampling.

Run after generate-reference.py. Optional REBOUND is used only for the final
same-time eight-orbit clip. These are synthetic trajectories, not observations.
"""
from pathlib import Path
import argparse, ctypes, importlib.util, json, math, sys


def pack(data, rebound, library_path, identifier):
    disk = data["disk"]
    library = ctypes.CDLL(str(library_path))
    library.ac15_configure.argtypes = [ctypes.c_double] * 4 + [ctypes.c_int]
    library.ac15_configure(
        disk["surfaceDensityAtOneAu"],
        disk["aspectRatioAtOneAu"],
        disk["temperatureSlope"],
        disk["viscosity"],
        int(disk["temperatureTransition"]),
    )
    history = []
    for record in data["records"]:
        history.append(
            [
                record["time"] / data["initialInnerPeriod"],
                *[x["semimajorAxis"] for x in record["elements"]],
                *[x["eccentricity"] for x in record["elements"]],
                record["periodRatio"],
                *[v * 180 / math.pi for v in record["unwrappedResonanceAngles"]],
            ]
        )
    final = data["records"][-1]
    start_time = final["time"] - 8 * data["initialInnerPeriod"]
    checkpoint = next(
        row for row in reversed(data["records"]) if row["time"] <= start_time
    )
    simulation = rebound.Simulation()
    simulation.G = 1
    for mass, point in zip(data["masses"], checkpoint["cartesian"]):
        simulation.add(m=mass, x=point[0], y=point[1], vx=point[2], vy=point[3])
    simulation.integrator = "ias15"
    simulation.integrator.epsilon = 1e-10
    simulation.force_is_velocity_dependent = 1
    simulation.additional_forces = ctypes.cast(
        library.ac15_forces, ctypes.c_void_p
    ).value
    clip = []
    for index in range(1025):
        target = start_time + 8 * data["initialInnerPeriod"] * index / 1024
        simulation.integrate(target - checkpoint["time"])
        clip.append(
            [
                target / data["initialInnerPeriod"],
                *[
                    value
                    for planet in simulation.particles
                    for value in [planet.x, planet.y, planet.vx, planet.vy]
                ],
            ]
        )
    wrapped = identifier not in ["capture", "fast", "gravity"]
    if wrapped:
        for row in history:
            row[6] = (row[6] + 180) % 360 - 180
            row[7] = (row[7] + 180) % 360 - 180
    late = history[len(history) // 2 :]
    spans = [max(row[k] for row in late) - min(row[k] for row in late) for k in [6, 7]]
    return {
        "id": identifier,
        "phaseMode": (
            "unwrapped" if identifier in ["capture", "fast", "gravity"] else "wrapped"
        ),
        "kind": data["kind"],
        "integrator": data["integrator"],
        "settings": {
            "outerMassEarth": data["outerMassEarth"],
            "aspectRatioAtOneAu": disk["aspectRatioAtOneAu"],
            "logViscosity": math.log10(disk["viscosity"]),
            "temperatureSlope": disk["temperatureSlope"],
            "initialPeriodRatio": data["initialPeriodRatio"],
            "surfaceDensityAtOneAu": disk["surfaceDensityAtOneAu"],
            "temperatureTransition": disk["temperatureTransition"],
        },
        "masses": data["masses"],
        "initialInnerPeriod": data["initialInnerPeriod"],
        "historyColumns": [
            "innerOrbits",
            "innerA",
            "outerA",
            "innerE",
            "outerE",
            "periodRatio",
            "wrappedInnerAngleDegrees" if wrapped else "unwrappedInnerAngleDegrees",
            "wrappedOuterAngleDegrees" if wrapped else "unwrappedOuterAngleDegrees",
        ],
        "history": history,
        "clipColumns": [
            "innerOrbits",
            "starX",
            "starY",
            "starVx",
            "starVy",
            "innerX",
            "innerY",
            "innerVx",
            "innerVy",
            "outerX",
            "outerY",
            "outerVx",
            "outerVy",
        ],
        "clip": clip,
        "stopped": data["stopped"],
        "stopReason": data["stopReason"],
        "diagnosis": {
            "lateHalfSpansDegrees": None if wrapped else spans,
            "largestSampleAngleStepDegrees": data["largestPhaseStep"] * 180 / math.pi,
            "phaseSamplingInnerOrbits": 0.5,
            "savedHistoryIntervalInnerOrbits": data["requestedInnerOrbits"]
            / data["requestedSamples"],
            "note": "Unwrap every half initial inner orbit before downsampling. Finite window, not an indefinite stability proof.",
        },
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--id", required=True)
    parser.add_argument("--rebound-path", type=Path)
    parser.add_argument("--native-library", type=Path, required=True)
    args = parser.parse_args()
    if args.rebound_path:
        sys.path.insert(0, str(args.rebound_path))
    import rebound

    assert rebound.__version__ == "5.2.2"
    args.output.write_text(
        json.dumps(
            pack(
                json.loads(args.input.read_text()),
                rebound,
                args.native_library,
                args.id,
            ),
            separators=(",", ":"),
            allow_nan=False,
        )
        + "\n"
    )


if __name__ == "__main__":
    main()