"""
GNSS orbit fit + prediction with an ECOM-type empirical solar-radiation-pressure model.

satkit's own propagator only has a cannonball SRP. Here the equations of motion are
integrated with scipy (DOP853), vectorised over many satellites at once, using
  * Earth gravity: JGM3 spherical harmonics (from satkit's data files), Cunningham recursion
  * Sun and Moon point masses (satkit JPL DE440 ephemeris), with indirect term
  * Schwarzschild relativistic term
  * ECOM SRP in the Sun-oriented D/Y/B frame, switched off in the Earth's shadow:
        D(du) = D0 + D2c cos 2du + D2s sin 2du + D4c cos 4du + D4s sin 4du
        Y     = Y0
        B(du) = B0 + Bc cos du + Bs sin du
    du = argument of latitude of the satellite relative to the Sun (ECOM2 convention).
satkit provides time scales, ephemerides and the IAU-2010 ITRF<->GCRF rotation.

Partials for the least-squares fit are formed by finite differences, with every
perturbed copy integrated in the same vectorised run (identical step sequence).
"""
from datetime import datetime, timedelta
from functools import lru_cache
import numpy as np
from scipy.integrate import solve_ivp
from scipy.interpolate import CubicSpline

OMEGA = 7.292115146706979e-5      # Earth rotation rate [rad/s]
GM_SUN = 1.32712440018e20
GM_MOON = 4.9028001218467e12
C_LIGHT = 299792458.0
R_EARTH_SHADOW = 6378137.0
PAR_NAMES = ['D0', 'Y0', 'B0', 'Bc', 'Bs', 'D2c', 'D2s', 'D4c', 'D4s']
ACC_SCALE = 1e-9                  # ECOM parameters in nm/s^2


@lru_cache(maxsize=4)
def load_gravity(nmax, model='JGM3'):
    """Return (GM, R, C, S) un-normalised coefficient arrays up to degree/order nmax."""
    import satkit as sk
    from math import factorial
    from pathlib import Path
    f = Path(sk.utils.datadir()) / f'{model}.gfc'
    C = np.zeros((nmax + 1, nmax + 1))
    S = np.zeros((nmax + 1, nmax + 1))
    gm = r = None
    head = True
    for line in open(f):
        w = line.split()
        if head:
            if w and w[0] == 'earth_gravity_constant':
                gm = float(w[1])
            elif w and w[0] == 'radius':
                r = float(w[1])
            elif line.startswith('end_of_head'):
                head = False
            continue
        if w and w[0] == 'gfc':
            n, m = int(w[1]), int(w[2])
            if n <= nmax and m <= n:
                N = np.sqrt((2 - (m == 0)) * (2 * n + 1) * factorial(n - m) / factorial(n + m))
                C[n, m] = float(w[3].replace('D', 'E')) * N
                S[n, m] = float(w[4].replace('D', 'E')) * N
    return gm, r, C, S


def gravity_itrf(r, nmax=8, model='JGM3'):
    """Vectorised spherical-harmonic gravity acceleration. r: (M,3) ITRF [m] -> (M,3) [m/s^2]."""
    GM, R, C, S = load_gravity(nmax, model)
    x, y, z = r[:, 0], r[:, 1], r[:, 2]
    r2 = x * x + y * y + z * z
    rho = R * R / r2
    x0, y0, z0 = R * x / r2, R * y / r2, R * z / r2
    n1 = nmax + 2
    V = np.zeros((n1, n1, len(x)))
    W = np.zeros((n1, n1, len(x)))
    V[0, 0] = R / np.sqrt(r2)
    V[1, 0] = z0 * V[0, 0]
    for n in range(2, n1):
        V[n, 0] = ((2 * n - 1) * z0 * V[n - 1, 0] - (n - 1) * rho * V[n - 2, 0]) / n
    for m in range(1, n1):
        V[m, m] = (2 * m - 1) * (x0 * V[m - 1, m - 1] - y0 * W[m - 1, m - 1])
        W[m, m] = (2 * m - 1) * (x0 * W[m - 1, m - 1] + y0 * V[m - 1, m - 1])
        if m + 1 < n1:
            V[m + 1, m] = (2 * m + 1) * z0 * V[m, m]
            W[m + 1, m] = (2 * m + 1) * z0 * W[m, m]
        for n in range(m + 2, n1):
            V[n, m] = ((2 * n - 1) * z0 * V[n - 1, m] - (n + m - 1) * rho * V[n - 2, m]) / (n - m)
            W[n, m] = ((2 * n - 1) * z0 * W[n - 1, m] - (n + m - 1) * rho * W[n - 2, m]) / (n - m)
    ax = np.zeros_like(x); ay = np.zeros_like(x); az = np.zeros_like(x)
    for n in range(nmax + 1):
        c, s = C[n, 0], S[n, 0]
        ax -= c * V[n + 1, 1]
        ay -= c * W[n + 1, 1]
        az += (n + 1) * (-c * V[n + 1, 0] - s * W[n + 1, 0])
        for m in range(1, n + 1):
            c, s = C[n, m], S[n, m]
            fac = 0.5 * (n - m + 1) * (n - m + 2)
            ax += 0.5 * (-c * V[n + 1, m + 1] - s * W[n + 1, m + 1]) + fac * (c * V[n + 1, m - 1] + s * W[n + 1, m - 1])
            ay += 0.5 * (-c * W[n + 1, m + 1] + s * V[n + 1, m + 1]) + fac * (-c * W[n + 1, m - 1] + s * V[n + 1, m - 1])
            az += (n - m + 1) * (-c * V[n + 1, m] - s * W[n + 1, m])
    return GM / R ** 2 * np.stack((ax, ay, az), axis=1)


R_SUN = 696000e3


def shadow_conical(r, s):
    """Fraction of the solar disc visible (0 umbra .. 1 full sun), Earth occultation, conical model.
    r: (M,3) satellite GCRF [m]; s: (3,) Sun GCRF [m]. Montenbruck & Gill (2000) eq. 3.85-3.87."""
    d = s[None, :] - r
    dn = np.linalg.norm(d, axis=1)
    rn = np.linalg.norm(r, axis=1)
    a = np.arcsin(np.clip(R_SUN / dn, -1, 1))            # apparent Sun radius
    b = np.arcsin(np.clip(R_EARTH_SHADOW / rn, -1, 1))   # apparent Earth radius
    c = np.arccos(np.clip(np.sum(-r * d, axis=1) / (rn * dn), -1, 1))  # Sun-Earth angular separation
    nu = np.ones(len(r))
    umb = c <= b - a
    nu[umb] = 0.0
    pen = (np.abs(a - b) < c) & (c < a + b)
    if pen.any():
        aa, bb, cc = a[pen], b[pen], c[pen]
        x = (cc ** 2 + aa ** 2 - bb ** 2) / (2 * cc)
        y = np.sqrt(np.maximum(aa ** 2 - x ** 2, 0))
        A = aa ** 2 * np.arccos(np.clip(x / aa, -1, 1)) + bb ** 2 * np.arccos(np.clip((cc - x) / bb, -1, 1)) - cc * y
        nu[pen] = 1 - A / (np.pi * aa ** 2)
    return nu


class Environment:
    """Time-dependent quantities on [mjd0 - pad, mjd0 + span + pad] (GPST MJD), t in seconds from mjd0."""

    def __init__(self, mjd0, span_sec, nmax=8, rot_step=1800.0, eph_step=600.0, t_min=0.0):
        import satkit as sk
        from satkit_ext import _sktime
        self.mjd0, self.nmax = mjd0, nmax
        self.GM = load_gravity(nmax)[0]
        t_min = min(0.0, t_min)
        tr = np.arange(t_min - rot_step, span_sec + 2 * rot_step, rot_step)
        q = sk.frametransform.qitrf2gcrf(_sktime(mjd0 + tr / 86400))
        self.rot_t = tr
        self.rot_R = np.array([x.as_rotation_matrix() for x in q])
        self.rot_step = rot_step
        te = np.arange(t_min - 2 * eph_step, span_sec + 3 * eph_step, eph_step)
        tms = _sktime(mjd0 + te / 86400)
        self.sun = CubicSpline(te, np.array(sk.jplephem.geocentric_pos(sk.solarsystem.Sun, tms)))
        self.moon = CubicSpline(te, np.array(sk.jplephem.geocentric_pos(sk.solarsystem.Moon, tms)))

    def R_itrf2gcrf(self, t):
        k = int(np.clip(np.rint((t - self.rot_t[0]) / self.rot_step), 0, len(self.rot_t) - 1))
        th = OMEGA * (t - self.rot_t[k])
        c, s = np.cos(th), np.sin(th)
        return self.rot_R[k] @ np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]])

    def accel(self, t, r, v, p):
        """r, v: (M,3) GCRF; p: (M,9) ECOM params [nm/s^2]. Returns (M,3)."""
        R = self.R_itrf2gcrf(t)
        a = gravity_itrf(r @ R, self.nmax) @ R.T          # r_itrf = R^T r ; a_gcrf = R a_itrf
        s, mn = self.sun(t), self.moon(t)
        for body, gm in ((s, GM_SUN), (mn, GM_MOON)):
            d = body[None, :] - r
            a += gm * (d / np.linalg.norm(d, axis=1)[:, None] ** 3 - body / np.linalg.norm(body) ** 3)
        rn = np.linalg.norm(r, axis=1)[:, None]
        v2 = np.sum(v * v, axis=1)[:, None]
        rv = np.sum(r * v, axis=1)[:, None]
        a += self.GM / (C_LIGHT ** 2 * rn ** 3) * ((4 * self.GM / rn - v2) * r + 4 * rv * v)
        # ECOM SRP
        eD = s[None, :] - r
        eD /= np.linalg.norm(eD, axis=1)[:, None]
        er = r / rn
        eY = np.cross(eD, er)
        eY /= np.linalg.norm(eY, axis=1)[:, None]
        eB = np.cross(eD, eY)
        h = np.cross(r, v)
        h /= np.linalg.norm(h, axis=1)[:, None]
        sh = s / np.linalg.norm(s)
        sp = sh[None, :] - np.sum(sh[None, :] * h, axis=1)[:, None] * h  # Sun direction projected in orbit plane
        sp /= np.linalg.norm(sp, axis=1)[:, None]
        cu = np.sum(er * sp, axis=1)
        su = np.sum(np.cross(sp, er) * h, axis=1)
        du = np.arctan2(su, cu)
        D = p[:, 0] + p[:, 5] * np.cos(2 * du) + p[:, 6] * np.sin(2 * du) + p[:, 7] * np.cos(4 * du) + p[:, 8] * np.sin(4 * du)
        Y = p[:, 1]
        B = p[:, 2] + p[:, 3] * np.cos(du) + p[:, 4] * np.sin(du)
        lit = shadow_conical(r, s)
        a += (ACC_SCALE * lit)[:, None] * (D[:, None] * eD + Y[:, None] * eY + B[:, None] * eB)
        return a


STEP = 600.0  # fixed integration step [s] (150-600 s agree to < 1 mm over 72 h)


def integrate(env, x0, p, t_eval, step=STEP, with_vel=False):
    """x0: (M,6) GCRF state at t=0; p: (M,9). Returns positions (len(t_eval), M, 3)
    [and velocities if with_vel]. t_eval sorted; negative times are integrated backwards from t=0.
    DOP853 with a fixed step (huge tolerances + max_step), so that every satellite / perturbed copy
    in the batch sees the same step sequence and the error does not depend on the batch size.
    step=None -> adaptive with tight tolerances (reference)."""
    M = len(x0)
    t_eval = np.asarray(t_eval, float)

    def f(t, y):
        y = y.reshape(M, 6)
        return np.hstack((y[:, 3:], env.accel(t, y[:, :3], y[:, 3:], p))).ravel()

    kw = dict(rtol=1e-13, atol=1e-7) if step is None else dict(rtol=1e3, atol=1e12, first_step=step, max_step=step)
    parts = []
    for te in (t_eval[t_eval < 0][::-1], t_eval[t_eval >= 0]):
        if len(te) == 0:
            continue
        if te[-1] == 0.0:
            y = x0.reshape(-1, 1)
        else:
            sol = solve_ivp(f, (0.0, float(te[-1])), x0.ravel(), method='DOP853', t_eval=te, **kw)
            if not sol.success:
                raise RuntimeError(sol.message)
            y = sol.y
        y = y.reshape(M, 6, -1)
        parts.append(y[:, :, ::-1] if te[0] < 0 or (len(te) > 1 and te[1] < te[0]) else y)
    y = np.concatenate(parts, axis=2)
    pos = y[:, :3, :].transpose(2, 0, 1)
    return (pos, y[:, 3:, :].transpose(2, 0, 1)) if with_vel else pos


def fit_ecom(rg, mjd_fit, mjd_out, est=('D0', 'Y0', 'B0', 'Bc', 'Bs', 'D2c', 'D2s'), nmax=8,
             niter=8, sigma_par=None, mjd_dense=None, dense_step=300.0):
    """
    Batched fit for N satellites.
      rg      : (N, nf, 3) GCRF positions [m] at mjd_fit (GPST MJD, common to all satellites)
      mjd_out : output epochs (>= mjd_fit[0])
      est     : ECOM parameters to estimate (others fixed at 0)
      sigma_par: optional a-priori sigma [nm/s^2] per estimated parameter (dict) for regularisation
      mjd_dense: optional extra epochs (may precede mjd_fit[0]: integrated backwards) at which the
                 fitted orbit is also returned with velocities (integration step dense_step)
    Returns pos_out (N, nout, 3) GCRF [m], fit rms (N,) [m], params (N, 9) [nm/s^2]
    [, dense_pos (N, nd, 3) [m], dense_vel (N, nd, 3) [m/s]  if mjd_dense is given].
    """
    from tools import lagrange
    N, nf, _ = rg.shape
    t_fit = (mjd_fit - mjd_fit[0]) * 86400
    t_out = (np.asarray(mjd_out) - mjd_fit[0]) * 86400
    t_dense = None if mjd_dense is None else (np.asarray(mjd_dense) - mjd_fit[0]) * 86400
    env = Environment(mjd_fit[0], max(t_fit[-1], t_out[-1], -np.inf if t_dense is None else t_dense[-1]), nmax,
                      t_min=0.0 if t_dense is None else t_dense[0])
    pidx = [PAR_NAMES.index(x) for x in est]
    nu = 6 + len(pidx)
    n0 = min(9, nf)
    x = np.zeros((N, 6))
    for i in range(N):
        _, v0 = lagrange(t_fit[:n0], rg[i, :n0], t_fit[:1], n0)
        x[i] = np.hstack((rg[i, 0], v0[0]))
    p = np.zeros((N, 9))
    dpert = np.array([1.0] * 3 + [1e-4] * 3 + [0.1] * len(pidx))
    W0 = np.zeros(nu)
    if sigma_par:
        for k, name in enumerate(est):
            if name in sigma_par:
                W0[6 + k] = 1.0 / sigma_par[name] ** 2
    rms = np.full(N, np.inf)
    done = np.zeros(N, bool)
    for it in range(niter):
        act = np.where(~done)[0]
        if len(act) == 0:
            break
        na = len(act)
        X = np.repeat(x[act], nu + 1, axis=0).reshape(na, nu + 1, 6)
        P = np.repeat(p[act], nu + 1, axis=0).reshape(na, nu + 1, 9)
        for k in range(nu):
            if k < 6:
                X[:, k + 1, k] += dpert[k]
            else:
                P[:, k + 1, pidx[k - 6]] += dpert[k]
        # first iteration: the initial velocity is rough, skip partial-free shortcuts
        pos = integrate(env, X.reshape(-1, 6), P.reshape(-1, 9), t_fit).reshape(nf, na, nu + 1, 3)
        for j, i in enumerate(act):
            nom = pos[:, j, 0, :]
            H = ((pos[:, j, 1:, :] - nom[:, None, :]) / dpert[None, :, None]).transpose(0, 2, 1).reshape(-1, nu)
            dy = (rg[i] - nom).ravel()
            new_rms = np.sqrt(np.mean(dy ** 2))
            # a-priori constraint pulls the parameter towards 0
            pv = np.hstack((np.zeros(6), p[i, pidx]))
            Nm = H.T @ H + np.diag(W0)
            b = H.T @ dy - W0 * pv
            sol = np.linalg.solve(Nm, b)
            if new_rms > 0.999 * rms[i] and it > 2:
                done[i] = True
                continue
            rms[i] = new_rms
            x[i] += sol[:6]
            p[i, pidx] += sol[6:]
            if np.max(np.abs(sol[:3])) < 1e-3:
                done[i] = True
    # final run: fitted orbit at fit and output epochs
    t_all = np.union1d(t_fit, t_out)
    pos = integrate(env, x, p, t_all)
    fit_pos = pos[np.searchsorted(t_all, t_fit)]
    rms = np.sqrt(np.mean((fit_pos.transpose(1, 0, 2) - rg) ** 2, axis=(1, 2)))
    pos_out = pos[np.searchsorted(t_all, t_out)].transpose(1, 0, 2)
    if t_dense is None:
        return pos_out, rms, p
    dp, dv = integrate(env, x, p, t_dense, step=dense_step, with_vel=True)
    return pos_out, rms, p, dp.transpose(1, 0, 2), dv.transpose(1, 0, 2)