/*
 * kepler_funcs.c -- Keplerian propagation and heliocentric type
 *
 * Two-body propagation from classical orbital elements for comets
 * and asteroids.  Handles elliptic (e<1), near-parabolic (e~1),
 * and hyperbolic (e>1) orbits.
 *
 * Also provides the heliocentric type (ecliptic J2000 position in AU)
 * and comet_observe() which computes topocentric coordinates given
 * the comet's orbital elements and Earth's heliocentric position.
 */

#include "postgres.h"
#include "fmgr.h"
#include "funcapi.h"
#include "utils/timestamp.h"
#include "libpq/pqformat.h"
#include "types.h"
#include "astro_math.h"
#include <math.h>

/* Heliocentric type I/O */
PG_FUNCTION_INFO_V1(heliocentric_in);
PG_FUNCTION_INFO_V1(heliocentric_out);
PG_FUNCTION_INFO_V1(heliocentric_recv);
PG_FUNCTION_INFO_V1(heliocentric_send);
PG_FUNCTION_INFO_V1(helio_x);
PG_FUNCTION_INFO_V1(helio_y);
PG_FUNCTION_INFO_V1(helio_z);
PG_FUNCTION_INFO_V1(helio_distance);

/* Propagation functions */
PG_FUNCTION_INFO_V1(kepler_propagate);
PG_FUNCTION_INFO_V1(comet_observe);

/* ================================================================
 * Kepler equation solvers
 * ================================================================
 */

/*
 * Elliptic Kepler equation: M = E - e*sin(E).
 * Newton-Raphson, converges in 3-5 iterations for e < 0.99.
 */
static double
solve_kepler_elliptic(double M, double e)
{
    double E = M;
    int    i;

    /* Better initial guess for high eccentricity */
    if (e > 0.8)
        E = M_PI;

    for (i = 0; i < 30; i++)
    {
        double dE = (E - e * sin(E) - M) / (1.0 - e * cos(E));
        E -= dE;
        if (fabs(dE) < 1.0e-15)
            break;
    }
    return E;
}


/*
 * Hyperbolic Kepler equation: M = e*sinh(H) - H.
 * Newton-Raphson.
 */
static double
solve_kepler_hyperbolic(double M, double e)
{
    /* Initial guess: asinh(M/e) for large M, M for small */
    double H = (fabs(M) > 1.0) ? copysign(log(fabs(M) / e + 1.0), M) : M;
    int    i;

    for (i = 0; i < 30; i++)
    {
        double dH = (e * sinh(H) - H - M) / (e * cosh(H) - 1.0);
        H -= dH;
        if (fabs(dH) < 1.0e-15)
            break;
    }
    return H;
}


/*
 * Near-parabolic: Barker's equation W^3 + 3W - 3M = 0.
 * Returns true anomaly directly.
 */
static double
solve_kepler_parabolic(double M)
{
    double W, f, fp;
    int    i;

    /* Cardano's formula for W^3 + 3W - 3M = 0 gives decent initial guess */
    double disc = sqrt(1.0 + M * M);
    W = cbrt(3.0 * M + 3.0 * disc) - cbrt(-3.0 * M + 3.0 * disc);

    for (i = 0; i < 30; i++)
    {
        f  = W * W * W + 3.0 * W - 3.0 * M;
        fp = 3.0 * W * W + 3.0;
        W -= f / fp;
        if (fabs(f / fp) < 1.0e-15)
            break;
    }

    return 2.0 * atan(W);
}


/* ================================================================
 * Internal: Keplerian position from orbital elements
 *
 * Input:  q (AU), e, inc (rad), omega (rad), Omega (rad),
 *         T_peri (JD), jd (observation JD)
 * Output: pos[3] in AU, ecliptic J2000 frame
 * ================================================================
 */
static void
kepler_position(double q, double e, double inc, double omega, double Omega,
                double T_peri, double jd, double pos[3])
{
    double dt = jd - T_peri;
    double v = 0.0, r = 0.0;
    double x_orb, y_orb;
    double cos_om, sin_om, cos_Om, sin_Om, cos_i, sin_i;
    double Px, Py, Pz, Qx, Qy, Qz;

    if (e < 0.99)
    {
        double a = q / (1.0 - e);
        double n = GAUSS_K / (a * sqrt(a));
        double M = fmod(n * dt, 2.0 * M_PI);
        double E;

        if (M < 0.0)
            M += 2.0 * M_PI;

        E = solve_kepler_elliptic(M, e);
        v = 2.0 * atan2(sqrt(1.0 + e) * sin(E / 2.0),
                         sqrt(1.0 - e) * cos(E / 2.0));
        r = a * (1.0 - e * cos(E));
    }
    else if (e > 1.01)
    {
        double a = q / (e - 1.0);
        double n = GAUSS_K / (a * sqrt(a));
        double M = n * dt;
        double H = solve_kepler_hyperbolic(M, e);

        v = 2.0 * atan2(sqrt(e + 1.0) * tanh(H / 2.0),
                         sqrt(e - 1.0));
        r = a * (e * cosh(H) - 1.0);
    }
    else
    {
        double n = GAUSS_K * sqrt(1.0 / (2.0 * q * q * q));
        double M = n * dt;

        v = solve_kepler_parabolic(M);
        r = q * (1.0 + tan(v / 2.0) * tan(v / 2.0));
    }

    x_orb = r * cos(v);
    y_orb = r * sin(v);

    /* Perifocal-to-ecliptic rotation (P, Q vectors) */
    cos_om = cos(omega);
    sin_om = sin(omega);
    cos_Om = cos(Omega);
    sin_Om = sin(Omega);
    cos_i  = cos(inc);
    sin_i  = sin(inc);

    Px = cos_Om * cos_om - sin_Om * sin_om * cos_i;
    Py = sin_Om * cos_om + cos_Om * sin_om * cos_i;
    Pz = sin_om * sin_i;

    Qx = -cos_Om * sin_om - sin_Om * cos_om * cos_i;
    Qy = -sin_Om * sin_om + cos_Om * cos_om * cos_i;
    Qz = cos_om * sin_i;

    pos[0] = Px * x_orb + Qx * y_orb;
    pos[1] = Py * x_orb + Qy * y_orb;
    pos[2] = Pz * x_orb + Qz * y_orb;
}


/* ================================================================
 * Heliocentric type I/O
 *
 * Text format: (x_au, y_au, z_au)
 * ================================================================
 */

Datum
heliocentric_in(PG_FUNCTION_ARGS)
{
    char             *str = PG_GETARG_CSTRING(0);
    pg_heliocentric  *result;
    double            x, y, z;
    int               nfields;

    result = (pg_heliocentric *) palloc(sizeof(pg_heliocentric));

    nfields = sscanf(str, " ( %lf , %lf , %lf )", &x, &y, &z);

    if (nfields != 3)
        ereport(ERROR,
                (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
                 errmsg("invalid input syntax for type heliocentric: \"%s\"", str),
                 errhint("Expected (x_au,y_au,z_au).")));

    result->x = x;
    result->y = y;
    result->z = z;

    PG_RETURN_POINTER(result);
}

Datum
heliocentric_out(PG_FUNCTION_ARGS)
{
    pg_heliocentric *h = (pg_heliocentric *) PG_GETARG_POINTER(0);

    PG_RETURN_CSTRING(psprintf("(%.10f,%.10f,%.10f)", h->x, h->y, h->z));
}

Datum
heliocentric_recv(PG_FUNCTION_ARGS)
{
    StringInfo        buf = (StringInfo) PG_GETARG_POINTER(0);
    pg_heliocentric  *result;

    result = (pg_heliocentric *) palloc(sizeof(pg_heliocentric));
    result->x = pq_getmsgfloat8(buf);
    result->y = pq_getmsgfloat8(buf);
    result->z = pq_getmsgfloat8(buf);

    PG_RETURN_POINTER(result);
}

Datum
heliocentric_send(PG_FUNCTION_ARGS)
{
    pg_heliocentric *h = (pg_heliocentric *) PG_GETARG_POINTER(0);
    StringInfoData   buf;

    pq_begintypsend(&buf);
    pq_sendfloat8(&buf, h->x);
    pq_sendfloat8(&buf, h->y);
    pq_sendfloat8(&buf, h->z);

    PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}

/* --- heliocentric accessors --- */

Datum
helio_x(PG_FUNCTION_ARGS)
{
    pg_heliocentric *h = (pg_heliocentric *) PG_GETARG_POINTER(0);
    PG_RETURN_FLOAT8(h->x);
}

Datum
helio_y(PG_FUNCTION_ARGS)
{
    pg_heliocentric *h = (pg_heliocentric *) PG_GETARG_POINTER(0);
    PG_RETURN_FLOAT8(h->y);
}

Datum
helio_z(PG_FUNCTION_ARGS)
{
    pg_heliocentric *h = (pg_heliocentric *) PG_GETARG_POINTER(0);
    PG_RETURN_FLOAT8(h->z);
}

Datum
helio_distance(PG_FUNCTION_ARGS)
{
    pg_heliocentric *h = (pg_heliocentric *) PG_GETARG_POINTER(0);
    PG_RETURN_FLOAT8(sqrt(h->x * h->x + h->y * h->y + h->z * h->z));
}


/* ================================================================
 * kepler_propagate(q, e, i_deg, omega_deg, Omega_deg, T_peri_jd,
 *                  timestamptz) -> heliocentric
 *
 * Propagate a body from classical orbital elements via two-body
 * Keplerian dynamics.  Returns heliocentric ecliptic J2000 position.
 * ================================================================
 */
Datum
kepler_propagate(PG_FUNCTION_ARGS)
{
    double   q_au      = PG_GETARG_FLOAT8(0);
    double   ecc       = PG_GETARG_FLOAT8(1);
    double   inc_deg   = PG_GETARG_FLOAT8(2);
    double   omega_deg = PG_GETARG_FLOAT8(3);
    double   Omega_deg = PG_GETARG_FLOAT8(4);
    double   T_peri_jd = PG_GETARG_FLOAT8(5);
    int64    ts        = PG_GETARG_INT64(6);

    double           jd;
    double           pos[3];
    pg_heliocentric *result;

    if (q_au <= 0.0)
        ereport(ERROR,
                (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
                 errmsg("perihelion distance must be positive: %.6f", q_au)));

    if (ecc < 0.0)
        ereport(ERROR,
                (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
                 errmsg("eccentricity must be non-negative: %.6f", ecc)));

    jd = timestamptz_to_jd(ts);

    kepler_position(q_au, ecc,
                    inc_deg * DEG_TO_RAD,
                    omega_deg * DEG_TO_RAD,
                    Omega_deg * DEG_TO_RAD,
                    T_peri_jd, jd, pos);

    result = (pg_heliocentric *) palloc(sizeof(pg_heliocentric));
    result->x = pos[0];
    result->y = pos[1];
    result->z = pos[2];

    PG_RETURN_POINTER(result);
}


/* ================================================================
 * comet_observe(q, e, i, omega, Omega, T_peri,
 *               earth_x_au, earth_y_au, earth_z_au,
 *               observer, timestamptz) -> topocentric
 *
 * Full pipeline: Keplerian propagation -> geocentric position ->
 * equatorial J2000 -> precession to date -> topocentric az/el.
 *
 * Earth's heliocentric ecliptic J2000 position must be provided.
 * In Phase 1, Skyfield or similar supplies this.
 * In Phase 2, planet_heliocentric(BODY_EARTH, ts) provides it.
 * ================================================================
 */
Datum
comet_observe(PG_FUNCTION_ARGS)
{
    double       q_au      = PG_GETARG_FLOAT8(0);
    double       ecc       = PG_GETARG_FLOAT8(1);
    double       inc_deg   = PG_GETARG_FLOAT8(2);
    double       omega_deg = PG_GETARG_FLOAT8(3);
    double       Omega_deg = PG_GETARG_FLOAT8(4);
    double       T_peri_jd = PG_GETARG_FLOAT8(5);
    double       earth_x   = PG_GETARG_FLOAT8(6);
    double       earth_y   = PG_GETARG_FLOAT8(7);
    double       earth_z   = PG_GETARG_FLOAT8(8);
    pg_observer *obs       = (pg_observer *) PG_GETARG_POINTER(9);
    int64        ts        = PG_GETARG_INT64(10);

    double          jd;
    double          comet_helio[3];
    double          geo_ecl[3];
    double          geo_equ[3];
    double          ra_j2000, dec_j2000, geo_dist;
    double          ra_date, dec_date;
    double          gmst_val, lst, ha;
    double          az, el;
    pg_topocentric *result;

    if (q_au <= 0.0)
        ereport(ERROR,
                (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
                 errmsg("perihelion distance must be positive: %.6f", q_au)));

    if (ecc < 0.0)
        ereport(ERROR,
                (errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
                 errmsg("eccentricity must be non-negative: %.6f", ecc)));

    jd = timestamptz_to_jd(ts);

    /* Comet heliocentric ecliptic J2000 position */
    kepler_position(q_au, ecc,
                    inc_deg * DEG_TO_RAD,
                    omega_deg * DEG_TO_RAD,
                    Omega_deg * DEG_TO_RAD,
                    T_peri_jd, jd, comet_helio);

    /* Geocentric ecliptic position = comet_helio - earth_helio */
    geo_ecl[0] = comet_helio[0] - earth_x;
    geo_ecl[1] = comet_helio[1] - earth_y;
    geo_ecl[2] = comet_helio[2] - earth_z;

    /* Ecliptic J2000 -> equatorial J2000 */
    ecliptic_to_equatorial(geo_ecl, geo_equ);

    /* Cartesian -> spherical (RA, Dec, distance) */
    cartesian_to_spherical(geo_equ, &ra_j2000, &dec_j2000, &geo_dist);

    /* Precess J2000 RA/Dec to date */
    precess_j2000_to_date(jd, ra_j2000, dec_j2000, &ra_date, &dec_date);

    /* Hour angle and az/el */
    gmst_val = gmst_from_jd(jd);
    lst = gmst_val + obs->lon;
    ha  = lst - ra_date;

    equatorial_to_horizontal(ha, dec_date, obs->lat, &az, &el);

    result = (pg_topocentric *) palloc(sizeof(pg_topocentric));
    result->azimuth    = az;
    result->elevation  = el;
    result->range_km   = geo_dist * AU_KM;
    result->range_rate = 0.0;  /* no velocity computation in Phase 1 */

    PG_RETURN_POINTER(result);
}