/*
 * transfer_funcs.c -- Interplanetary transfer orbit SQL functions
 *
 * Wraps the Lambert solver for use in PostgreSQL queries.
 * Enables pork chop plots and transfer window analysis as SQL:
 *
 *   SELECT dep_date, arr_date, c3_departure, c3_arrival, delta_v
 *   FROM generate_series(...) dep_date
 *   CROSS JOIN generate_series(...) arr_date
 *   CROSS JOIN LATERAL lambert_transfer(3, 4, dep_date, arr_date) t
 *   WHERE t IS NOT NULL;
 *
 * All functions are IMMUTABLE STRICT PARALLEL SAFE.
 */

#include "postgres.h"
#include "fmgr.h"
#include "funcapi.h"
#include "catalog/pg_type.h"
#include "utils/builtins.h"
#include "utils/timestamp.h"
#include "types.h"
#include "vsop87.h"
#include "lambert.h"
#include <math.h>

PG_FUNCTION_INFO_V1(lambert_transfer);
PG_FUNCTION_INFO_V1(lambert_c3);

/*
 * Compute planet heliocentric velocity via numerical differentiation.
 * Central difference: v = (pos(t+dt) - pos(t-dt)) / (2*dt)
 * dt = 0.01 day ~ 14 minutes, gives ~6 significant digits.
 */
static void
planet_velocity(int vsop_body, double jd, double vel[3])
{
    double pos_fwd[6], pos_bwd[6];
    double dt = 0.01;  /* days */

    GetVsop87Coor(jd + dt, vsop_body, pos_fwd);
    GetVsop87Coor(jd - dt, vsop_body, pos_bwd);

    vel[0] = (pos_fwd[0] - pos_bwd[0]) / (2.0 * dt);
    vel[1] = (pos_fwd[1] - pos_bwd[1]) / (2.0 * dt);
    vel[2] = (pos_fwd[2] - pos_bwd[2]) / (2.0 * dt);
}

/*
 * lambert_transfer(dep_body_id int4, arr_body_id int4,
 *                  dep_time timestamptz, arr_time timestamptz)
 * RETURNS RECORD (c3_departure float8, c3_arrival float8,
 *                 v_inf_dep float8, v_inf_arr float8,
 *                 tof_days float8, sma float8)
 *
 * Solves Lambert's problem for a transfer between two planets.
 * Returns NULL if the solver fails to converge (e.g., TOF too short).
 *
 * C3 = v_infinity^2 (km^2/s^2), the characteristic energy.
 * v_inf = hyperbolic excess velocity at departure/arrival (km/s).
 * tof_days = time of flight in days.
 * sma = transfer orbit semi-major axis (AU).
 */
Datum
lambert_transfer(PG_FUNCTION_ARGS)
{
    int32       dep_body = PG_GETARG_INT32(0);
    int32       arr_body = PG_GETARG_INT32(1);
    int64       dep_ts   = PG_GETARG_INT64(2);
    int64       arr_ts   = PG_GETARG_INT64(3);
    double      dep_jd, arr_jd, tof_days;
    double      r1[6], r2[6];
    double      v_planet_dep[3], v_planet_arr[3];
    double      v_inf_dep[3], v_inf_arr[3];
    double      v_inf_dep_mag, v_inf_arr_mag;
    double      c3_dep, c3_arr;
    int         dep_vsop, arr_vsop;
    lambert_result lr;
    TupleDesc   tupdesc;
    Datum       values[6];
    bool        nulls[6];
    HeapTuple   tuple;
    double      au_per_day_to_km_per_s;
    int         k;

    /* Validate body IDs */
    if (dep_body < 1 || dep_body > 8)
        ereport(ERROR,
                (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                 errmsg("lambert_transfer: dep_body_id %d must be 1-8", dep_body)));
    if (arr_body < 1 || arr_body > 8)
        ereport(ERROR,
                (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                 errmsg("lambert_transfer: arr_body_id %d must be 1-8", arr_body)));
    if (dep_body == arr_body)
        ereport(ERROR,
                (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                 errmsg("lambert_transfer: departure and arrival bodies must differ")));

    dep_jd = timestamptz_to_jd(dep_ts);
    arr_jd = timestamptz_to_jd(arr_ts);
    tof_days = arr_jd - dep_jd;

    if (tof_days <= 0.0)
        ereport(ERROR,
                (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                 errmsg("lambert_transfer: arrival must be after departure")));

    /* VSOP87 uses 0-based: Mercury=0, ..., Neptune=7 */
    dep_vsop = dep_body - 1;
    arr_vsop = arr_body - 1;

    /* Get planet positions */
    GetVsop87Coor(dep_jd, dep_vsop, r1);
    GetVsop87Coor(arr_jd, arr_vsop, r2);

    /* Solve Lambert's problem (prograde, short-way) */
    if (!lambert_solve_uv(r1, r2, tof_days, GAUSS_K2, 1, &lr))
        PG_RETURN_NULL();

    /* Planet velocities via numerical differentiation */
    planet_velocity(dep_vsop, dep_jd, v_planet_dep);
    planet_velocity(arr_vsop, arr_jd, v_planet_arr);

    /* Hyperbolic excess velocity = transfer velocity - planet velocity */
    /* Convert AU/day to km/s: 1 AU/day = 149597870.7 / 86400 km/s */
    au_per_day_to_km_per_s = AU_KM / 86400.0;

    for (k = 0; k < 3; k++) {
        v_inf_dep[k] = (lr.v1[k] - v_planet_dep[k]) * au_per_day_to_km_per_s;
        v_inf_arr[k] = (lr.v2[k] - v_planet_arr[k]) * au_per_day_to_km_per_s;
    }

    v_inf_dep_mag = sqrt(v_inf_dep[0]*v_inf_dep[0] +
                         v_inf_dep[1]*v_inf_dep[1] +
                         v_inf_dep[2]*v_inf_dep[2]);
    v_inf_arr_mag = sqrt(v_inf_arr[0]*v_inf_arr[0] +
                         v_inf_arr[1]*v_inf_arr[1] +
                         v_inf_arr[2]*v_inf_arr[2]);

    /* C3 = v_infinity^2 */
    c3_dep = v_inf_dep_mag * v_inf_dep_mag;
    c3_arr = v_inf_arr_mag * v_inf_arr_mag;

    /* Build composite return */
    if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
        ereport(ERROR,
                (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
                 errmsg("function returning record called in context "
                        "that cannot accept type record")));

    tupdesc = BlessTupleDesc(tupdesc);

    memset(nulls, 0, sizeof(nulls));
    values[0] = Float8GetDatum(c3_dep);
    values[1] = Float8GetDatum(c3_arr);
    values[2] = Float8GetDatum(v_inf_dep_mag);
    values[3] = Float8GetDatum(v_inf_arr_mag);
    values[4] = Float8GetDatum(tof_days);
    values[5] = Float8GetDatum(lr.sma);

    tuple = heap_form_tuple(tupdesc, values, nulls);
    PG_RETURN_DATUM(HeapTupleGetDatum(tuple));
}


/*
 * lambert_c3(dep_body_id int4, arr_body_id int4,
 *            dep_time timestamptz, arr_time timestamptz)
 * RETURNS float8
 *
 * Convenience function: returns departure C3 only (km^2/s^2).
 * NULL if solver doesn't converge.
 * Useful for pork chop plot coloring.
 */
Datum
lambert_c3(PG_FUNCTION_ARGS)
{
    int32       dep_body = PG_GETARG_INT32(0);
    int32       arr_body = PG_GETARG_INT32(1);
    int64       dep_ts   = PG_GETARG_INT64(2);
    int64       arr_ts   = PG_GETARG_INT64(3);
    double      dep_jd, arr_jd, tof_days;
    double      r1[6], r2[6];
    double      v_planet_dep[3];
    double      v_inf_dep[3];
    double      c3_dep;
    int         dep_vsop, arr_vsop;
    lambert_result lr;
    double      au_per_day_to_km_per_s;
    int         k;

    if (dep_body < 1 || dep_body > 8 || arr_body < 1 || arr_body > 8)
        PG_RETURN_NULL();
    if (dep_body == arr_body)
        PG_RETURN_NULL();

    dep_jd = timestamptz_to_jd(dep_ts);
    arr_jd = timestamptz_to_jd(arr_ts);
    tof_days = arr_jd - dep_jd;

    if (tof_days <= 0.0)
        PG_RETURN_NULL();

    dep_vsop = dep_body - 1;
    arr_vsop = arr_body - 1;

    GetVsop87Coor(dep_jd, dep_vsop, r1);
    GetVsop87Coor(arr_jd, arr_vsop, r2);

    if (!lambert_solve_uv(r1, r2, tof_days, GAUSS_K2, 1, &lr))
        PG_RETURN_NULL();

    planet_velocity(dep_vsop, dep_jd, v_planet_dep);

    au_per_day_to_km_per_s = AU_KM / 86400.0;

    for (k = 0; k < 3; k++)
        v_inf_dep[k] = (lr.v1[k] - v_planet_dep[k]) * au_per_day_to_km_per_s;

    c3_dep = v_inf_dep[0]*v_inf_dep[0] +
             v_inf_dep[1]*v_inf_dep[1] +
             v_inf_dep[2]*v_inf_dep[2];

    PG_RETURN_FLOAT8(c3_dep);
}