#!/usr/bin/env python3
"""Generate a declared synthetic migration trajectory with independent REBOUND IAS15.

The basic public analysis consumes frozen records without REBOUND. Regeneration
optionally needs pinned REBOUND5.2.2 in a separate environment/package directory.
"""
from pathlib import Path
import argparse
import json
import math
import sys
import time
import ctypes
from disk_torques import DEFAULT_DISK, EARTH_SOLAR_MASS_RATIO
from migration_orbits import (
    external_disk_force,
    orbital_elements,
    resonance_angles,
    barycentric_invariants,
)


def generate(
    rebound,
    disk,
    outer_mass_earth,
    initial_ratio,
    orbits,
    output,
    sample_count,
    native_library=None,
    phase_cadence=0.5,
    epsilon=1e-10,
):
    if (
        orbits <= 8
        or sample_count < 16
        or phase_cadence <= 0
        or not 0 < epsilon < 1e-3
        or outer_mass_earth <= 0
        or initial_ratio <= 1
    ):
        raise ValueError(
            "Require positive masses, phase cadence and tolerance; duration above8 orbits, samples at least16, period ratio above1"
        )
    started = time.monotonic()
    simulation = rebound.Simulation()
    simulation.G = 1
    simulation.add(m=1)
    inner_mass = outer_mass_earth * 0.5 * EARTH_SOLAR_MASS_RATIO
    outer_mass = outer_mass_earth * EARTH_SOLAR_MASS_RATIO
    outer_axis = (initial_ratio**2 * (1 + outer_mass) / (1 + inner_mass)) ** (1 / 3)
    simulation.add(
        m=inner_mass, a=1, e=0.001, omega=0.2, M=0, primary=simulation.particles[0]
    )
    simulation.add(
        m=outer_mass,
        a=outer_axis,
        e=0.001,
        omega=1.2,
        M=2.0,
        primary=simulation.particles[0],
    )
    simulation.move_to_com()
    simulation.integrator = "ias15"
    simulation.integrator.epsilon = epsilon
    simulation.force_is_velocity_dependent = 1

    def disk_force(pointer):
        system = pointer.contents
        star = system.particles[0]
        for planet in system.particles[1:]:
            acceleration, _ = external_disk_force(
                planet.m,
                [planet.x - star.x, planet.y - star.y],
                [planet.vx - star.vx, planet.vy - star.vy],
                disk,
            )
            planet.ax += acceleration[0]
            planet.ay += acceleration[1]

    library = None
    if native_library:
        library = ctypes.CDLL(str(native_library))
        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"]),
        )
        library.ac15_acceleration.argtypes = [ctypes.c_double] * 5 + [
            ctypes.POINTER(ctypes.c_double)
        ] * 2
        for radius, eccentricity in [(1, 0), (1.6, 0.01), (0.8, 0.02)]:
            velocity = math.sqrt((1 + outer_mass) / radius) * (1 + eccentricity)
            expected, _ = external_disk_force(
                outer_mass, [radius, 0], [0.001, velocity], disk
            )
            ax, ay = ctypes.c_double(), ctypes.c_double()
            library.ac15_acceleration(
                outer_mass,
                radius,
                0,
                0.001,
                velocity,
                ctypes.byref(ax),
                ctypes.byref(ay),
            )
            assert math.isclose(ax.value, expected[0], rel_tol=2e-12, abs_tol=1e-20)
            assert math.isclose(ay.value, expected[1], rel_tol=2e-12, abs_tol=1e-20)
        simulation.additional_forces = ctypes.cast(
            library.ac15_forces, ctypes.c_void_p
        ).value
    else:
        simulation.additional_forces = disk_force

    domain_coverage = {
        "checks": 0,
        "maximumEccentricityOverAspect": 0.0,
        "maximumThermalMassRatio": 0.0,
        "minimumToomreParameter": math.inf,
        "minimumSeparationOverMutualHill": math.inf,
    }

    def inspect_domain(system):
        domain_coverage["checks"] += 1
        star = system.particles[0]
        axes = []
        for planet in system.particles[1:]:
            position = [planet.x - star.x, planet.y - star.y]
            velocity = [planet.vx - star.vx, planet.vy - star.vy]
            elements = orbital_elements(position, velocity, 1 + planet.m)
            if not 0.35 < elements["semimajorAxis"] < 3:
                return "Semimajor axis outside0.35..3AU"
            axes.append(elements["semimajorAxis"])
            torque = external_disk_force(planet.m, position, velocity, disk)[1]
            domain_coverage["maximumEccentricityOverAspect"] = max(
                domain_coverage["maximumEccentricityOverAspect"],
                elements["eccentricity"] / torque["aspect"],
            )
            domain_coverage["maximumThermalMassRatio"] = max(
                domain_coverage["maximumThermalMassRatio"], torque["thermalMassRatio"]
            )
            domain_coverage["minimumToomreParameter"] = min(
                domain_coverage["minimumToomreParameter"], torque["toomreQ"]
            )
            if elements["eccentricity"] / torque["aspect"] > 0.7:
                return "Eccentricity exceeds0.7 disk aspect ratio"
            if torque["thermalMassRatio"] > 0.2:
                return "Planet exceeds0.2 thermal mass"
            if torque["toomreQ"] < 20:
                return "Local ToomreQ below20"
        hill = ((inner_mass + outer_mass) / 3) ** (1 / 3) * sum(axes) / 2
        separation = math.hypot(
            system.particles[1].x - system.particles[2].x,
            system.particles[1].y - system.particles[2].y,
        )
        domain_coverage["minimumSeparationOverMutualHill"] = min(
            domain_coverage["minimumSeparationOverMutualHill"], separation / hill
        )
        if (
            math.hypot(
                system.particles[1].x - system.particles[2].x,
                system.particles[1].y - system.particles[2].y,
            )
            < 3 * hill
        ):
            return "Separation below3 mutual Hill radii"
        return None

    initial_period = 2 * math.pi / math.sqrt(1 + inner_mass)
    records = []
    stop_reason = None
    previous_angles = None
    unwrapped_angles = [0.0, 0.0]
    largest_phase_step = 0.0
    for sample in range(sample_count + 1):
        target_time = initial_period * orbits * sample / sample_count
        while simulation.t < target_time:
            simulation.integrate(
                min(target_time, simulation.t + initial_period * phase_cadence)
            )
            current_star = simulation.particles[0]
            current_elements = [
                orbital_elements(
                    [p.x - current_star.x, p.y - current_star.y],
                    [p.vx - current_star.vx, p.vy - current_star.vy],
                    1 + p.m,
                )
                for p in simulation.particles[1:]
            ]
            current_angles = resonance_angles(*current_elements)
            if previous_angles is None:
                unwrapped_angles = current_angles[:]
            else:
                for index in range(2):
                    difference = (
                        current_angles[index] - previous_angles[index] + math.pi
                    ) % (2 * math.pi) - math.pi
                    largest_phase_step = max(largest_phase_step, abs(difference))
                    unwrapped_angles[index] += difference
            previous_angles = current_angles
            stop_reason = inspect_domain(simulation)
            if stop_reason:
                break
        if sample == 0:
            current_star = simulation.particles[0]
            previous_angles = resonance_angles(
                *[
                    orbital_elements(
                        [p.x - current_star.x, p.y - current_star.y],
                        [p.vx - current_star.vx, p.vy - current_star.vy],
                        1 + p.m,
                    )
                    for p in simulation.particles[1:]
                ]
            )
            unwrapped_angles = previous_angles[:]
        star = simulation.particles[0]
        elements = []
        for planet in simulation.particles[1:]:
            elements.append(
                orbital_elements(
                    [planet.x - star.x, planet.y - star.y],
                    [planet.vx - star.vx, planet.vy - star.vy],
                    1 + planet.m,
                )
            )
        torques = [
            external_disk_force(
                planet.m,
                [planet.x - star.x, planet.y - star.y],
                [planet.vx - star.vx, planet.vy - star.vy],
                disk,
            )[1]
            for planet in simulation.particles[1:]
        ]
        for current, torque in zip(elements, torques):
            if not 0.35 < current["semimajorAxis"] < 3:
                stop_reason = "Semimajor axis outside0.35..3AU"
            elif current["eccentricity"] / torque["aspect"] > 0.7:
                stop_reason = "Eccentricity exceeds0.7 disk aspect ratio"
            elif torque["thermalMassRatio"] > 0.2:
                stop_reason = "Planet exceeds0.2 thermal mass"
            elif torque["toomreQ"] < 20:
                stop_reason = "Local ToomreQ below20"
        period_ratio = (
            elements[1]["semimajorAxis"] / elements[0]["semimajorAxis"]
        ) ** 1.5 * math.sqrt((1 + inner_mass) / (1 + outer_mass))
        records.append(
            {
                "time": simulation.t,
                "elements": elements,
                "periodRatio": period_ratio,
                "resonanceAngles": resonance_angles(*elements),
                "unwrappedResonanceAngles": unwrapped_angles[:],
                "cartesian": [
                    [planet.x, planet.y, planet.vx, planet.vy]
                    for planet in simulation.particles
                ],
                "torques": torques,
            }
        )
        if stop_reason:
            break
    final = records[-1]
    result = {
        "kind": "Synthetic finite-mass three-body trajectory; not Kepler-223 observations or hydrodynamics",
        "integrator": "REBOUND5.2.2 IAS15 epsilon" + str(epsilon),
        "disk": disk,
        "outerMassEarth": outer_mass_earth,
        "initialPeriodRatio": initial_ratio,
        "requestedInnerOrbits": orbits,
        "masses": [1, inner_mass, outer_mass],
        "units": "G=Mstar=AU=1; one time unit sqrt(AU^3/GMstar)",
        "initialInnerPeriod": initial_period,
        "requestedSamples": sample_count,
        "domainCheckCadenceInnerOrbits": phase_cadence,
        "domainCoverage": domain_coverage,
        "phaseCadenceInnerOrbits": phase_cadence,
        "largestPhaseStep": largest_phase_step,
        "stopped": bool(stop_reason),
        "stopReason": stop_reason,
        "records": records,
    }

    def finite_json(value):
        if isinstance(value, float) and not math.isfinite(value):
            return None
        if isinstance(value, list):
            return [finite_json(item) for item in value]
        if isinstance(value, dict):
            return {key: finite_json(item) for key, item in value.items()}
        return value

    output.write_text(
        json.dumps(finite_json(result), separators=(",", ":"), allow_nan=False) + "\n"
    )
    print(
        json.dumps(
            {
                "seconds": time.monotonic() - started,
                "finalOrbits": final["time"] / initial_period,
                "periodRatio": final["periodRatio"],
                "eccentricities": [
                    element["eccentricity"] for element in final["elements"]
                ],
                "stopReason": stop_reason,
                "samples": len(records),
            }
        ),
        flush=True,
    )
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--rebound-path", type=Path)
    parser.add_argument("--orbits", type=float, default=6000)
    parser.add_argument("--samples", type=int, default=1200)
    parser.add_argument("--outer-mass-earth", type=float, default=4)
    parser.add_argument("--initial-ratio", type=float, default=2.01)
    parser.add_argument("--density", type=float, default=1e-4)
    parser.add_argument("--slope", type=float, default=0.5)
    parser.add_argument("--viscosity", type=float, default=1e-6)
    parser.add_argument("--transition", action="store_true")
    parser.add_argument("--phase-cadence", type=float, default=0.5)
    parser.add_argument("--epsilon", type=float, default=1e-10)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--native-library",
        type=Path,
        help="Optional temporary library compiled from migration-force.c against pinned REBOUND header",
    )
    arguments = parser.parse_args()
    if arguments.rebound_path:
        sys.path.insert(0, str(arguments.rebound_path))
    import rebound

    assert rebound.__version__ == "5.2.2"
    generate(
        rebound,
        {
            **DEFAULT_DISK,
            "surfaceDensityAtOneAu": arguments.density,
            "temperatureSlope": arguments.slope,
            "viscosity": arguments.viscosity,
            "temperatureTransition": arguments.transition,
        },
        arguments.outer_mass_earth,
        arguments.initial_ratio,
        arguments.orbits,
        arguments.output,
        arguments.samples,
        arguments.native_library,
        arguments.phase_cadence,
        arguments.epsilon,
    )


if __name__ == "__main__":
    main()