"""Finite-mass Cartesian dynamics with explicitly external disk forces.

G=1; planar vectors, all bodies in one inertial barycentric frame.
"""

import math
from disk_torques import torque_components


def orbital_elements(position, velocity, gravity):
    radius = math.hypot(*position)
    speed_squared = sum(value * value for value in velocity)
    radial_product = sum(first * second for first, second in zip(position, velocity))
    energy = speed_squared / 2 - gravity / radius
    angular_momentum = position[0] * velocity[1] - position[1] * velocity[0]
    eccentricity_vector = [
        ((speed_squared - gravity / radius) * coordinate - radial_product * speed)
        / gravity
        for coordinate, speed in zip(position, velocity)
    ]
    eccentricity = math.hypot(*eccentricity_vector)
    semimajor_axis = -gravity / (2 * energy) if energy < 0 else math.inf
    periapsis = math.atan2(eccentricity_vector[1], eccentricity_vector[0])
    if eccentricity > 1e-10 and energy < 0:
        eccentric_anomaly = math.atan2(
            radial_product / math.sqrt(gravity * semimajor_axis),
            1 - radius / semimajor_axis,
        )
        mean_anomaly = eccentric_anomaly - eccentricity * math.sin(eccentric_anomaly)
        longitude = periapsis + mean_anomaly
    else:
        longitude = math.atan2(position[1], position[0])
    return {
        "semimajorAxis": semimajor_axis,
        "eccentricity": eccentricity,
        "periapsis": periapsis,
        "meanLongitude": longitude,
        "specificAngularMomentum": angular_momentum,
        "specificEnergy": energy,
    }


def external_disk_force(mass, position, velocity, disk):
    elements = orbital_elements(position, velocity, 1 + mass)
    if not math.isfinite(elements["semimajorAxis"]):
        return [0.0, 0.0], None
    torque = torque_components(
        mass, elements["semimajorAxis"], elements["eccentricity"], disk
    )
    signed_momentum = mass * elements["specificAngularMomentum"]
    rate = torque["total"] / signed_momentum
    radius_squared = sum(coordinate * coordinate for coordinate in position)
    radial_product = sum(
        coordinate * speed for coordinate, speed in zip(position, velocity)
    )
    damping = (
        0
        if math.isinf(torque["eccentricityTime"])
        else -2 * radial_product / (radius_squared * torque["eccentricityTime"])
    )
    return [
        rate * speed + damping * coordinate
        for coordinate, speed in zip(position, velocity)
    ], torque


def barycentric_invariants(bodies):
    kinetic = sum(
        body["mass"] * sum(speed * speed for speed in body["velocity"]) / 2
        for body in bodies
    )
    potential = 0.0
    for index, first in enumerate(bodies):
        for second in bodies[index + 1 :]:
            distance = math.dist(first["position"], second["position"])
            potential -= first["mass"] * second["mass"] / distance
    angular_momentum = sum(
        body["mass"]
        * (
            body["position"][0] * body["velocity"][1]
            - body["position"][1] * body["velocity"][0]
        )
        for body in bodies
    )
    return {"energy": kinetic + potential, "angularMomentum": angular_momentum}


def resonance_angles(first, second):
    combination = 2 * second["meanLongitude"] - first["meanLongitude"]
    return [
        math.atan2(math.sin(combination - periapsis), math.cos(combination - periapsis))
        for periapsis in [first["periapsis"], second["periapsis"]]
    ]