# SPDX-License-Identifier: MIT
"""YS-EN-0002 · Navigating by parallax.

Where am I, given the directions to a few nearby stars whose positions I
already know (from Gaia)? Each observed direction is a line in space that
passes through the star; the ship sits where those lines cross.

  star i at heliocentric position s_i, ship at r
  observed unit vector  v_i = (s_i - r) / |s_i - r|

Least squares: with P_i = I - v_i v_iᵀ (projector perpendicular to the
line of sight), minimise Σ |P_i (r - s_i)|² / w_i², giving

  (Σ P_i / w_i²) r = Σ P_i s_i / w_i²,     w_i = σ_i · |s_i|

(|s_i| stands in for |s_i - r|; the difference moves the fix by ~5e-6 au.)

and the position covariance (Σ P_i / w_i²)⁻¹. One star constrains the two
directions perpendicular to its line of sight with error σ·d; the third
direction needs a second star well separated on the sky.

Usage: python parallax_fix.py   (writes results.json; requires numpy)
"""
from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
AU_PER_PC = 648_000 / math.pi          # IAU 2015 B2: 1 pc = 648000/π au
MAS = math.pi / (180 * 3600 * 1000)    # radians per milliarcsecond


def unit(ra_deg: float, dec_deg: float) -> np.ndarray:
    a, d = math.radians(ra_deg), math.radians(dec_deg)
    return np.array([math.cos(d) * math.cos(a), math.cos(d) * math.sin(a), math.sin(d)])


def star_position(ra, dec, parallax_mas) -> np.ndarray:
    """Heliocentric position in au from Gaia RA/Dec (deg) and parallax (mas)."""
    d_au = AU_PER_PC / (parallax_mas / 1000)
    return d_au * unit(ra, dec)


def observe(stars: np.ndarray, r: np.ndarray) -> np.ndarray:
    v = stars - r
    return v / np.linalg.norm(v, axis=1, keepdims=True)


def fix(stars: np.ndarray, dirs: np.ndarray, sigma_rad) -> tuple[np.ndarray, np.ndarray]:
    """Weighted least-squares intersection of lines of sight.

    Weights use a first-pass range estimate (|s_i| is fine: the ship is
    tiny compared with the star distances), so the solve is linear.
    """
    sigma = np.broadcast_to(np.asarray(sigma_rad, float), (len(stars),))
    A = np.zeros((3, 3))
    b = np.zeros(3)
    for s, v, sg in zip(stars, dirs, sigma):
        P = np.eye(3) - np.outer(v, v)
        w2 = (sg * np.linalg.norm(s)) ** 2
        A += P / w2
        b += P @ s / w2
    cov = np.linalg.inv(A)
    return np.linalg.solve(A, b), cov


def perturb(dirs: np.ndarray, sigma_rad: float, rng) -> np.ndarray:
    """Add isotropic angular noise of sigma per axis to each unit vector."""
    out = []
    for v in dirs:
        # two axes perpendicular to v
        a = np.cross(v, [0, 0, 1.0])
        if np.linalg.norm(a) < 1e-8:
            a = np.cross(v, [0, 1.0, 0])
        a /= np.linalg.norm(a)
        c = np.cross(v, a)
        e1, e2 = rng.normal(0, sigma_rad, 2)
        w = v + e1 * a + e2 * c
        out.append(w / np.linalg.norm(w))
    return np.array(out)


def separation_deg(u1, u2) -> float:
    return math.degrees(math.acos(np.clip(np.dot(u1, u2), -1, 1)))


# ---------------------------------------------------------------- inputs
# Gaia DR3 (VizieR I/355/gaiadr3), epoch J2016.0: RA, Dec (deg), parallax (mas)
GAIA = {
    "Proxima Centauri": (217.39232147201, -62.67607511677, 768.0665),
    "Barnard's Star": (269.44850252544, 4.73942005111, 546.9759),
    "Wolf 359": (164.10319030756, 7.00272694098, 415.1794),
    "Lalande 21185": (165.83095967578, 35.94865303266, 392.7529),
    "Ross 154": (282.45878901752, -23.83709744873, 336.0266),
}
# Lauer et al. 2025 (AJ 170, 22), Table 2: Gaia positions propagated to the
# New Horizons epoch (2020 Apr 22-23), and the directions LORRI observed.
NH_GAIA_EPOCH = {"Proxima Centauri": (217.38246430, -62.67515412), "Wolf 359": (164.09852751, 6.99949593)}
NH_OBSERVED = {"Proxima Centauri": (217.36315, -62.676324), "Wolf 359": (164.0943006, 7.001008)}
SIGMA_NH = 0.18  # arcsec per star after six images (Lauer et al. §3.2)
# JPL (DSN) mean position of New Horizons over the imaging window, Table 5, au
JPL_NH = np.array([13.5495, -42.0195, -16.4573])
ARCSEC = math.pi / 648_000
LY_AU = 63_241.077  # IAU: 9 460 730 472 580.8 km / 149 597 870.7 km


def axes(cov):
    """1-sigma semi-axes (descending) and their directions."""
    w, v = np.linalg.eigh(cov)
    return np.sqrt(w[::-1]), v[:, ::-1]


def pair_axes(d1_pc, d2_pc, sep_deg):
    """Long/short error semi-axes in au per arcsec for two stars sep_deg apart."""
    u1 = np.array([1.0, 0, 0])
    t = math.radians(sep_deg)
    u2 = np.array([math.cos(t), math.sin(t), 0])
    S = np.array([u1 * d1_pc * AU_PER_PC, u2 * d2_pc * AU_PER_PC])
    _, cov = fix(S, S / np.linalg.norm(S, axis=1, keepdims=True), ARCSEC)
    a, _ = axes(cov)
    return float(a[0]), float(a[-1])


def main():
    out = {}

    # 1. reproduce the New Horizons fix from the paper's own inputs
    names = ["Proxima Centauri", "Wolf 359"]
    S = np.array([star_position(*NH_GAIA_EPOCH[n], GAIA[n][2]) for n in names])
    D = np.array([unit(*NH_OBSERVED[n]) for n in names])
    r, cov = fix(S, D, SIGMA_NH * ARCSEC)
    a, vec = axes(cov)
    err = r - JPL_NH
    out["nh_fix"] = {
        "position_au": r.tolist(), "range_au": float(np.linalg.norm(r)),
        "jpl_au": JPL_NH.tolist(), "jpl_range_au": float(np.linalg.norm(JPL_NH)),
        "error_au": err.tolist(), "error_norm_au": float(np.linalg.norm(err)),
        "axes_au": a.tolist(), "sigma_arcsec": SIGMA_NH,
        "separation_deg": separation_deg(*(S / np.linalg.norm(S, axis=1, keepdims=True))),
        "error_sigma": float(np.sqrt(err @ np.linalg.solve(cov, err))),
    }

    # 2. Monte Carlo: re-observe the true directions with 0.18" noise
    rng = np.random.default_rng(7)
    true_dirs = observe(S, JPL_NH)
    basis = vec  # long, middle, short axes
    pts = []
    for _ in range(4000):
        rr, _ = fix(S, perturb(true_dirs, SIGMA_NH * ARCSEC, rng), SIGMA_NH * ARCSEC)
        pts.append(((rr - JPL_NH) @ basis).tolist())
    pts = np.array(pts)
    out["monte_carlo"] = {
        "n": len(pts), "std_along_axes_au": pts.std(axis=0).tolist(),
        "sample": pts[:700, :2].round(4).tolist(),
        "lorri_fix_in_basis": (err @ basis).tolist(),
    }

    # 3. star pairs from Gaia DR3 (reproduces Lauer et al. Table 4)
    keys = list(GAIA)
    pairs = []
    for i in range(len(keys)):
        for j in range(i + 1, len(keys)):
            P = np.array([star_position(*GAIA[k]) for k in (keys[i], keys[j])])
            U = P / np.linalg.norm(P, axis=1, keepdims=True)
            _, c = fix(P, U, ARCSEC)
            aa, _ = axes(c)
            pairs.append({"a": keys[i], "b": keys[j], "sep_deg": separation_deg(*U),
                          "long": float(aa[0]), "short": float(aa[-1])})
    out["pairs"] = pairs

    # 4. geometry: long axis vs separation for Proxima + a star at Wolf 359's / Barnard's distance
    seps = list(range(5, 176, 5))
    d = {n: 1000 / GAIA[n][2] for n in GAIA}  # distance in pc
    out["dop"] = {
        "sep_deg": seps,
        "with_wolf_distance": [pair_axes(d["Proxima Centauri"], d["Wolf 359"], t)[0] for t in seps],
        "with_barnard_distance": [pair_axes(d["Proxima Centauri"], d["Barnard's Star"], t)[0] for t in seps],
        "short_wolf": pair_axes(d["Proxima Centauri"], d["Wolf 359"], 90)[1],
        "dist_pc": d,
    }

    # 5. apparent shift vs baseline (maximum: baseline perpendicular to sight line)
    base = list(np.geomspace(1, 1000, 31))
    out["shift"] = {
        "baseline_au": base,
        **{n: [b / (d[n] * AU_PER_PC) / ARCSEC for b in base] for n in ("Proxima Centauri", "Barnard's Star", "Wolf 359")},
    }

    # 6. our own datasheet: what angular precision does a ±0.02 ly fix need?
    need_au = 0.02 * LY_AU
    best = min(pairs, key=lambda p: p["long"])
    out["ysnav"] = {
        "fix_au": need_au, "best_pair": best,
        "sigma_needed_arcsec": need_au / best["long"],
        "fix_at_1arcsec_au": best["long"], "fix_at_1arcsec_ly": best["long"] / LY_AU,
    }

    (HERE / "results.json").write_text(json.dumps(out, indent=1))
    f = out["nh_fix"]
    print("NH fix", np.round(f["position_au"], 2), "range", round(f["range_au"], 2),
          "| error", round(f["error_norm_au"], 3), "au", f"({f['error_sigma']:.2f} sigma)",
          "| axes", np.round(f["axes_au"], 3), "| sep", round(f["separation_deg"], 2))
    print("MC std along axes", np.round(out["monte_carlo"]["std_along_axes_au"], 3))
    for p_ in pairs:
        print(f"  {p_['a']:>16} + {p_['b']:<15} {p_['sep_deg']:6.2f}°  long {p_['long']:.2f}  short {p_['short']:.2f} au/arcsec")
    y = out["ysnav"]
    print(f"±0.02 ly = {y['fix_au']:.0f} au -> needs {y['sigma_needed_arcsec']:.0f} arcsec with {best['a']} + {best['b']}")


if __name__ == "__main__":
    main()