/*
 * lambert.c -- Lambert transfer orbit solver (Universal Variable)
 *
 * Solves Lambert's problem using the Universal Variable formulation
 * with Stumpff functions c2(psi) and c3(psi). This handles elliptic,
 * parabolic, and hyperbolic transfers with one unified algorithm.
 *
 * The approach:
 *   1. Compute geometry: r1_mag, r2_mag, delta_theta
 *   2. Compute A from geometry (determines short/long way)
 *   3. Iterate on universal variable z to match time of flight
 *   4. Extract departure/arrival velocities from Lagrange coefficients
 *
 * References:
 *   Curtis, "Orbital Mechanics for Engineering Students" (2014), Ch. 5
 *   Bate, Mueller & White, "Fundamentals of Astrodynamics" (1971), Ch. 5
 */

#include <math.h>
#include "lambert.h"

/* Stumpff function c2(psi) = (1 - cos(sqrt(psi))) / psi */
static double
stumpff_c2(double psi)
{
    double sq;

    if (psi > 1e-6) {
        sq = sqrt(psi);
        return (1.0 - cos(sq)) / psi;
    }
    if (psi < -1e-6) {
        sq = sqrt(-psi);
        return (1.0 - cosh(sq)) / psi;
    }
    /* Taylor series near zero: 1/2 - psi/24 + psi^2/720 */
    return 1.0/2.0 - psi/24.0 + psi*psi/720.0;
}

/* Stumpff function c3(psi) = (sqrt(psi) - sin(sqrt(psi))) / (psi * sqrt(psi)) */
static double
stumpff_c3(double psi)
{
    double sq;

    if (psi > 1e-6) {
        sq = sqrt(psi);
        return (sq - sin(sq)) / (psi * sq);
    }
    if (psi < -1e-6) {
        sq = sqrt(-psi);
        return (sinh(sq) - sq) / (-psi * sq);
    }
    /* Taylor series near zero: 1/6 - psi/120 + psi^2/5040 */
    return 1.0/6.0 - psi/120.0 + psi*psi/5040.0;
}

int
lambert_solve_uv(const double r1[3], const double r2[3],
                 double tof_days, double mu, int prograde,
                 lambert_result *result)
{
    double r1_mag, r2_mag;
    double cos_dtheta, sin_dtheta, dtheta;
    double A, z, z_lo, z_hi;
    double c2, c3, y, x, t, dt_dz;
    double f, g, gdot;
    int i;
    int max_iter = 50;
    double tol = 1e-10;
    double cross_z;

    result->converged = 0;

    if (tof_days <= 0.0)
        return 0;

    /* Magnitudes */
    r1_mag = sqrt(r1[0]*r1[0] + r1[1]*r1[1] + r1[2]*r1[2]);
    r2_mag = sqrt(r2[0]*r2[0] + r2[1]*r2[1] + r2[2]*r2[2]);

    if (r1_mag < 1e-12 || r2_mag < 1e-12)
        return 0;

    /* Transfer angle from cross product */
    cos_dtheta = (r1[0]*r2[0] + r1[1]*r2[1] + r1[2]*r2[2]) / (r1_mag * r2_mag);

    /* Clamp for numerical safety */
    if (cos_dtheta > 1.0) cos_dtheta = 1.0;
    if (cos_dtheta < -1.0) cos_dtheta = -1.0;

    /* Cross product z-component determines orbit direction */
    cross_z = r1[0]*r2[1] - r1[1]*r2[0];

    if (prograde) {
        if (cross_z < 0.0)
            sin_dtheta = -sqrt(1.0 - cos_dtheta*cos_dtheta);
        else
            sin_dtheta = sqrt(1.0 - cos_dtheta*cos_dtheta);
    } else {
        if (cross_z >= 0.0)
            sin_dtheta = -sqrt(1.0 - cos_dtheta*cos_dtheta);
        else
            sin_dtheta = sqrt(1.0 - cos_dtheta*cos_dtheta);
    }

    dtheta = atan2(sin_dtheta, cos_dtheta);
    if (dtheta < 0.0)
        dtheta += 2.0 * M_PI;

    /* A parameter (Curtis eq. 5.35) */
    A = sin_dtheta * sqrt(r1_mag * r2_mag / (1.0 - cos_dtheta));

    if (fabs(A) < 1e-14)
        return 0;  /* Degenerate: 0 or 180 deg transfer */

    /*
     * Newton-Raphson iteration on z (universal variable).
     * z > 0: elliptic, z = 0: parabolic, z < 0: hyperbolic.
     * Initial bracket: z_lo = -4*pi^2 (max hyperbolic), z_hi from geometry.
     */
    z_lo = -4.0 * M_PI * M_PI;
    z_hi = 4.0 * M_PI * M_PI;
    z = 0.0;  /* Start at parabolic */

    for (i = 0; i < max_iter; i++) {
        c2 = stumpff_c2(z);
        c3 = stumpff_c3(z);

        y = r1_mag + r2_mag + A * (z * c3 - 1.0) / sqrt(c2);

        if (y < 0.0) {
            /* y must be positive; adjust z upward */
            z_lo = z;
            z = (z + z_hi) * 0.5;
            continue;
        }

        x = sqrt(y / c2);

        /* Time of flight for this z */
        t = (x*x*x * c3 + A * sqrt(y)) / sqrt(mu);

        /* Check convergence */
        if (fabs(t - tof_days) < tol * fabs(tof_days) + 1e-12) {
            result->converged = 1;
            break;
        }

        /* Derivative dt/dz for Newton step */
        if (fabs(z) > 1e-6) {
            dt_dz = (x*x*x * (stumpff_c2(z) - 3.0*c3/(2.0*c2)) / (2.0*c2)
                     + (3.0*c3*sqrt(y)/c2 + A/sqrt(y)*(1.0 - z*c3/c2)) * A / (2.0*c2*sqrt(c2)))
                    / sqrt(mu);
            /* Simplified: use bisection if Newton overshoots */
            if (fabs(dt_dz) < 1e-20) {
                /* Fall back to bisection */
                if (t < tof_days)
                    z_lo = z;
                else
                    z_hi = z;
                z = (z_lo + z_hi) * 0.5;
                continue;
            }
            z = z - (t - tof_days) / dt_dz;
        } else {
            /* Near parabolic: bisection */
            if (t < tof_days)
                z_lo = z;
            else
                z_hi = z;
            z = (z_lo + z_hi) * 0.5;
        }

        /* Keep z in bounds */
        if (z < z_lo) z = z_lo + 0.1 * (z_hi - z_lo);
        if (z > z_hi) z = z_hi - 0.1 * (z_hi - z_lo);
    }

    if (!result->converged) {
        /*
         * Newton didn't converge; try pure bisection as fallback.
         * Reset bounds and iterate.
         */
        z_lo = -4.0 * M_PI * M_PI;
        z_hi = 4.0 * M_PI * M_PI;

        for (i = 0; i < 100; i++) {
            z = (z_lo + z_hi) * 0.5;
            c2 = stumpff_c2(z);
            c3 = stumpff_c3(z);

            y = r1_mag + r2_mag + A * (z * c3 - 1.0) / sqrt(c2);
            if (y < 0.0) {
                z_lo = z;
                continue;
            }

            x = sqrt(y / c2);
            t = (x*x*x * c3 + A * sqrt(y)) / sqrt(mu);

            if (fabs(t - tof_days) < tol * fabs(tof_days) + 1e-12) {
                result->converged = 1;
                break;
            }

            if (t < tof_days)
                z_lo = z;
            else
                z_hi = z;

            if (fabs(z_hi - z_lo) < 1e-14)
                break;
        }
    }

    if (!result->converged)
        return 0;

    /* Lagrange coefficients (Curtis eqs. 5.46) */
    c2 = stumpff_c2(z);
    c3 = stumpff_c3(z);
    y = r1_mag + r2_mag + A * (z * c3 - 1.0) / sqrt(c2);

    f = 1.0 - y / r1_mag;
    g = A * sqrt(y / mu);
    gdot = 1.0 - y / r2_mag;

    /* Departure and arrival velocity vectors */
    {
        int k;
        for (k = 0; k < 3; k++) {
            result->v1[k] = (r2[k] - f * r1[k]) / g;
            result->v2[k] = (gdot * r2[k] - r1[k]) / g;
        }
    }

    /* Semi-major axis */
    if (fabs(c2) > 1e-14)
        result->sma = y / c2 / (2.0);   /* a = y / (2 * c2(z)) ... but simpler: */
    else
        result->sma = 1e30;  /* Near-parabolic */

    /* Correct SMA from z */
    if (fabs(z) > 1e-10)
        result->sma = y / (2.0 * c2);

    return 1;
}