"""
Extend SP3 GNSS orbit/clock products beyond their last epoch.

method='ecom' (default): per-satellite dynamic fit of the last 48 h with an ECOM2 empirical
  solar-radiation-pressure model, integrated by ecom_ext.py (Earth gravity, Sun/Moon,
  relativity, conical Earth shadow; satkit supplies ephemerides, time scales and frames).
  Backtest: +24 h median 3D error ~0.3-0.7 m; eclipsing satellites form the tail.
method='cannonball': satkit's own propagator with a cannonball SRP coefficient, 12 h fit.
  Backtest: +24 h median ~2-6 m (used as per-satellite fallback for 'ecom').

Clocks are extended with a linear fit over the last 12 h. The extension is blended into the
product's last epoch with an exp(-t/2 h) decay.

Units follow tools.read_GNSS: positions in km (ITRF), clocks in microseconds,
epochs as naive GPST datetimes.
"""
from datetime import datetime, timedelta
from multiprocessing import Pool
import numpy as np

MJD0 = datetime(1858, 11, 17)
GPST_TAI = 19.0  # TAI - GPST [s], fixed


def _mjd_gpst(dts):
    return np.array([(x - MJD0) / timedelta(days=1) for x in dts])


def _sktime(mjd_gpst):
    import satkit as sk
    # satkit's from_mjd(..., timescale.GPS) treats the value as an offset from the
    # GPS epoch (v0.9.3), so go through TAI, which is exact.
    return [sk.time.from_mjd(m + GPST_TAI / 86400, sk.timescale.TAI) for m in mjd_gpst]


def _rotations(tms, to_gcrf):
    import satkit as sk
    q = sk.frametransform.qitrf2gcrf(tms) if to_gcrf else sk.frametransform.qgcrf2itrf(tms)
    q = q if isinstance(q, list) else [q]
    return np.array([x.as_rotation_matrix() for x in q])


def update_satkit_data(max_age_days=1.0, log=print):
    """Refresh satkit EOP / space-weather files if older than max_age_days (network errors are non-fatal)."""
    import satkit as sk
    from pathlib import Path
    eop = Path(sk.utils.datadir()) / 'EOP-All.csv'
    age = (datetime.now() - datetime.fromtimestamp(eop.stat().st_mtime)) / timedelta(days=1) if eop.is_file() else np.inf
    if age > max_age_days:
        try:
            sk.utils.update_datafiles()
            log(f'- satkit data files updated (EOP age was {age:.1f} d).')
        except Exception as e:
            log(f'- satkit data update failed ({e}); using existing files.')


def fit_and_propagate(rg, mjd_fit, mjd_out, gravity_order=8, niter=12, srp=True):
    """
    Fit one satellite and propagate.
      rg      : (n,3) GCRF positions [m] at mjd_fit (GPST MJD)
      mjd_out : GPST MJD epochs to output (must be >= mjd_fit[0])
    Returns (positions (m,3) GCRF [m], fit rms [m], Cr*A/m [m^2/kg]).
    """
    import satkit as sk
    from tools import lagrange
    settings = sk.propsettings()
    settings.gravity_order = gravity_order
    settings.use_spaceweather = False  # no drag at GNSS altitude
    tms_fit = _sktime(mjd_fit)
    tms_out = _sktime(mjd_out)
    t0, tend = tms_fit[0], max(tms_fit[-1], tms_out[-1], key=lambda t: t.as_mjd())
    tsec = (mjd_fit - mjd_fit[0]) * 86400
    n0 = min(9, len(tsec))
    _, v0 = lagrange(tsec[:n0], rg[:n0], tsec[:1], n0)
    x, cr = np.hstack((rg[0], v0[0])), 0.0
    srp_on = False
    best = (np.inf, x, cr)

    def run(x, cr):
        kw = dict(output_phi=True, propsettings=settings)
        if srp_on or cr != 0.0:
            kw['satproperties'] = sk.satproperties_static(craoverm=cr)
        return sk.propagate(x, t0, end=tend, **kw)

    dy = np.zeros(1)
    for it in range(niter):
        res = run(x, cr)
        H, dy = [], []
        for tm, r in zip(tms_fit, rg):
            st, phi = res.interp(tm, output_phi=True)
            H.append(phi[:3, :])
            dy.append(r - st[:3])
        H, dy = np.vstack(H), np.hstack(dy)
        rms = np.sqrt(np.mean(dy ** 2))
        if srp_on:
            if rms < best[0]:
                best = (rms, x.copy(), cr)
            elif rms > 0.98 * best[0]:
                break
            dc = 2e-3
            res2 = run(x, cr + dc)
            hc = np.hstack([(res2.interp(tm)[:3] - res.interp(tm)[:3]) / dc for tm in tms_fit])
            H = np.hstack((H, hc[:, None]))
        sol = np.linalg.lstsq(H, dy, rcond=None)[0]
        x = x + sol[:6]
        if srp_on:
            cr += sol[6]
        if np.max(np.abs(sol[:3])) < 1e-2 and not srp_on:
            if not srp:
                best = (rms, x.copy(), cr)
                break
            srp_on = True
    rms, x, cr = best if np.isfinite(best[0]) else (np.sqrt(np.mean(dy ** 2)), x, cr)
    srp_on = cr != 0.0
    res = run(x, cr)
    return np.array([res.interp(tm)[:3] for tm in tms_out]), rms, cr


def _worker(args):
    rg, mjd_fit, mjd_out, gorder = args
    try:
        p, rms, cr = fit_and_propagate(rg, mjd_fit, mjd_out, gorder)
        if not np.isfinite(rms) or rms > 2.0:  # SRP fit unreliable -> plain dynamic fit
            p2, rms2, cr2 = fit_and_propagate(rg, mjd_fit, mjd_out, gorder, srp=False)
            if rms2 < rms:
                p, rms, cr = p2, rms2, cr2
        return p, rms, cr
    except Exception:
        return None, np.nan, np.nan


def _ecom_worker(args):
    rg, mjd_fit, mjd_out, est, nmax, mjd_dense = args
    from ecom_ext import fit_ecom
    try:
        r = fit_ecom(rg, mjd_fit, mjd_out, est=est, nmax=nmax, mjd_dense=mjd_dense)
        return r if mjd_dense is not None else r + (None, None)
    except Exception:
        n = len(rg)
        return np.full((n, len(mjd_out), 3), np.nan), np.full(n, np.nan), np.full((n, 9), np.nan), None, None


class SmoothRotation:
    """Smooth ITRF->GCRF rotation R(t) = Q(t) Rz(OMEGA t) W(t) on [t_min, t_max] s from mjd_ref (GPST MJD).
    Q (precession-nutation, UT1 drift) and W (polar motion) are Chebyshev fits to satkit's IAU-2010
    matrices at 30-min anchors. satkit's per-epoch matrices jitter by ~2e-9 rad (cm at GNSS altitude,
    floating-point time precision); the fit removes the jitter, slow differences are < 1e-9 rad."""

    def __init__(self, mjd_ref, t_min, t_max, anchor=1800.0, deg_q=8, deg_w=4):
        import satkit as sk
        from ecom_ext import OMEGA
        self.omega, self.mjd_ref = OMEGA, mjd_ref
        t = np.arange(t_min - anchor, t_max + 2 * anchor, anchor)
        tm = _sktime(mjd_ref + t / 86400)
        R = _rotations(tm, True)
        W = np.array([q.as_rotation_matrix() for q in sk.frametransform.qitrf2tirs(tm)])
        Q = np.einsum('nij,nkj,nlk->nil', R, W, self._rz(t))          # R W^T Z^T
        self.a, self.b = t[0], t[-1]
        tn = self._tn(t)
        self.cq = np.polynomial.chebyshev.chebfit(tn, Q.reshape(len(t), 9), deg_q)
        self.cw = np.polynomial.chebyshev.chebfit(tn, W.reshape(len(t), 9), deg_w)

    def _tn(self, t):
        return (2 * np.asarray(t, float) - self.a - self.b) / (self.b - self.a)

    def _rz(self, t):
        th = self.omega * np.asarray(t, float)
        c, s = np.cos(th), np.sin(th)
        Z = np.zeros((len(th), 3, 3))
        Z[:, 0, 0], Z[:, 0, 1], Z[:, 1, 0], Z[:, 1, 1], Z[:, 2, 2] = c, -s, s, c, 1.0
        return Z

    def itrf2gcrf(self, t):
        t = np.atleast_1d(np.asarray(t, float))
        tn = self._tn(t)
        Q = np.polynomial.chebyshev.chebval(tn, self.cq).T.reshape(-1, 3, 3)
        W = np.polynomial.chebyshev.chebval(tn, self.cw).T.reshape(-1, 3, 3)
        return Q @ self._rz(t) @ W

    def gcrf2itrf_state(self, t, rg, vg):
        """rg, vg: (n,3) GCRF [m, m/s] at t (s) -> ITRF position and Earth-fixed velocity."""
        RT = np.transpose(self.itrf2gcrf(t), (0, 2, 1))
        dRT = (np.transpose(self.itrf2gcrf(t + 1.0), (0, 2, 1)) - np.transpose(self.itrf2gcrf(t - 1.0), (0, 2, 1))) / 2
        return np.einsum('nij,nj->ni', RT, rg), np.einsum('nij,nj->ni', RT, vg) + np.einsum('nij,nj->ni', dRT, rg)


ECOM2 = ('D0', 'Y0', 'B0', 'Bc', 'Bs', 'D2c', 'D2s')
MAX_NPROC = 16  # hard cap on worker processes


def _fit_epochs(Gdt0, fit_hours, fit_step_min, step):
    fidx = np.where(Gdt0 >= Gdt0[-1] - timedelta(hours=fit_hours))[0]
    fidx_orb = fidx[::max(1, int(timedelta(minutes=fit_step_min) / step))]
    if fidx_orb[-1] != fidx[-1]:
        fidx_orb = np.append(fidx_orb, fidx[-1])
    return fidx, fidx_orb


def extend_gnss(Gdata0, Gdt0, Gsats, ext_hours=24, fit_hours=None, fit_step_min=15,
                gravity_order=8, max_fit_rms=5.0, blend_hours=2.0, nproc=16, log=print,
                method='ecom', ecom_params=ECOM2, ecom_max_rms=1.0, clock_fit_hours=12, dense_sec=30):
    """
    Append `ext_hours` of predicted epochs to an SP3 product.
      Gdata0 : (nt, nsat, >=4) [x, y, z (km, ITRF), clk (us), ...], Gdt0 sorted GPST datetimes
      method : 'ecom'       - own integrator with ECOM SRP (ecom_ext.py), default fit window 48 h;
                              satellites whose ECOM fit rms > ecom_max_rms fall back to 'cannonball'
               'cannonball' - satkit propagator, cannonball SRP, default fit window 12 h
    Returns Gdata (nt+next, nsat, same), Gdt, ext_info dict (fit rms [m], model used, SRP parameters).
    Satellites whose fit fails or exceeds max_fit_rms get NaN in the extension.
    The offset between the fitted orbit and the product at its last epoch is added to
    the extension with an exp(-t/blend_hours) decay, so the junction is continuous.
    Clocks: linear fit over the last clock_fit_hours.
    dense_sec (ECOM satellites only, None to disable): ext_info['dense'] holds positions [km] and
    Earth-fixed velocities [dm/s] on a dense_sec grid over the whole (extended) span, computed as
    fitted orbit + cubic spline of (product - fitted orbit) at the product epochs. This replaces
    high-order Lagrange interpolation, which oscillates near the ends of the span.
    """
    nproc = max(1, min(int(nproc), MAX_NPROC))
    Gdt0 = np.asarray(Gdt0)
    step = Gdt0[-1] - Gdt0[-2]
    n_ext = int(timedelta(hours=ext_hours) / step)
    Gdt_ext = np.array([Gdt0[-1] + step * (k + 1) for k in range(n_ext)])
    mjd_out = _mjd_gpst(np.concatenate((Gdt0[-1:], Gdt_ext)))  # last product epoch + extension
    R_out = _rotations(_sktime(mjd_out), False)  # GCRF->ITRF
    decay = np.exp(-(mjd_out[1:] - mjd_out[0]) * 24 / blend_hours)[:, None] if blend_hours > 0 else 0.0
    nsat = len(Gsats)
    pos = np.full((nsat, len(mjd_out), 3), np.nan)  # GCRF [m]
    fit_rms = np.full(nsat, np.nan)
    model = np.full(nsat, '', dtype='<U10')
    srp = np.full((nsat, 9), np.nan)

    def gcrf_tail(fh):
        _, fo = _fit_epochs(Gdt0, fh, fit_step_min, step)
        mjd = _mjd_gpst(Gdt0[fo])
        R = _rotations(_sktime(mjd), True)  # ITRF->GCRF, per epoch
        rg = np.einsum('nij,snj->sni', R, Gdata0[fo][:, :, :3].transpose(1, 0, 2) * 1e3)
        return rg, mjd, np.all(np.isfinite(rg), axis=(1, 2))

    todo = np.ones(nsat, bool)
    use_dense = method == 'ecom' and dense_sec
    if use_dense:
        dstep = timedelta(seconds=dense_sec)
        Gdt_d = np.arange(Gdt0[0], Gdt_ext[-1] + dstep, dstep).astype(datetime)
        mjd_d = _mjd_gpst(Gdt_d)
        dpos = np.full((nsat, len(Gdt_d), 3), np.nan)   # GCRF m
        dvel = np.full((nsat, len(Gdt_d), 3), np.nan)   # GCRF m/s
    if method == 'ecom':
        rg, mjd_fit, ok = gcrf_tail(fit_hours or 48)
        js = np.where(ok)[0]
        chunks = [c for c in np.array_split(js, max(1, min(nproc, len(js)))) if len(c)]
        jobs = [(rg[c], mjd_fit, mjd_out, tuple(ecom_params), gravity_order, mjd_d if use_dense else None)
                for c in chunks]
        if nproc > 1:
            with Pool(len(jobs)) as pool:
                results = pool.map(_ecom_worker, jobs)
        else:
            results = [_ecom_worker(x) for x in jobs]
        for c, (p, rms, par, dp, dv) in zip(chunks, results):
            good = np.isfinite(rms) & (rms <= ecom_max_rms)
            pos[c[good]], fit_rms[c[good]], srp[c[good]] = p[good], rms[good], par[good]
            if use_dense and dp is not None:
                dpos[c[good]], dvel[c[good]] = dp[good], dv[good]
            model[c[good]] = 'ecom'
            todo[c[good]] = False
        if todo.any():
            log(f'  ECOM fit not accepted for {todo.sum()} sats, using cannonball: {" ".join(Gsats[todo])}')
    if todo.any():
        rg, mjd_fit, ok = gcrf_tail(fit_hours if (fit_hours and method != 'ecom') else 12)
        js = np.where(todo & ok)[0]
        jobs = [(rg[j], mjd_fit, mjd_out, gravity_order) for j in js]
        if nproc > 1 and len(jobs) > 1:
            with Pool(min(nproc, len(jobs))) as pool:
                results = pool.map(_worker, jobs)
        else:
            results = [_worker(x) for x in jobs]
        for j, (p, rms, cr) in zip(js, results):
            if p is not None:
                pos[j], fit_rms[j], srp[j, 0] = p, rms, cr
                model[j] = 'cannonball'

    ext = np.full((n_ext, Gdata0.shape[1], Gdata0.shape[2]), np.nan)
    dense = None
    if use_dense:
        from scipy.interpolate import CubicSpline
        td = (mjd_d - mjd_d[0]) * 86400
        srot = SmoothRotation(mjd_d[0], td[0], td[-1])
        k_node = np.rint(np.array([(x - Gdt_d[0]) / dstep for x in np.concatenate((Gdt0, Gdt_ext))])).astype(int)
        t_node = td[k_node]
        n0 = len(Gdt0)
        dense = dict(t=Gdt_d, pos=np.full((len(Gdt_d), nsat, 3), np.nan), vel=np.full((len(Gdt_d), nsat, 3), np.nan),
                     ok=np.zeros(nsat, bool))
    for j in range(nsat):
        if not fit_rms[j] <= max_fit_rms or not np.all(np.isfinite(pos[j])):
            continue
        if use_dense and model[j] == 'ecom' and np.all(np.isfinite(dpos[j])) \
                and np.all(np.isfinite(Gdata0[:, j, :3])):
            ri, vi = srot.gcrf2itrf_state(td, dpos[j], dvel[j])      # smooth model, ITRF [m, m/s]
            ri /= 1e3
            ext[:, j, :3] = ri[k_node[n0:]] + (Gdata0[-1, j, :3] - ri[k_node[n0 - 1]]) * decay
            res = np.concatenate((Gdata0[:, j, :3], ext[:, j, :3])) - ri[k_node]   # km, smooth & small
            cs = CubicSpline(t_node, res, axis=0)
            dense['pos'][:, j] = ri + cs(td)
            dense['vel'][:, j] = vi * 10 + cs(td, 1) * 1e4                   # dm/s
            dense['ok'][j] = True
            continue
        p = np.einsum('nij,nj->ni', R_out, pos[j]) / 1e3
        ext[:, j, :3] = p[1:] + (Gdata0[-1, j, :3] - p[0]) * decay
    bad = [Gsats[j] for j in range(nsat) if not np.isfinite(ext[0, j, 0])]
    if bad:
        log(f'  orbit extension failed for {len(bad)} sats: {" ".join(bad)}')

    # clock: linear fit over the clock window
    fidx, _ = _fit_epochs(Gdt0, clock_fit_hours, fit_step_min, step)
    tt = (_mjd_gpst(Gdt0[fidx]) - mjd_out[0]) * 86400
    te = (mjd_out[1:] - mjd_out[0]) * 86400
    for j in range(nsat):
        c = Gdata0[fidx, j, 3]
        m = np.isfinite(c)
        if m.sum() >= 2 and np.isfinite(ext[0, j, 0]):
            pc = np.polyfit(tt[m], c[m], 1)
            off = (Gdata0[-1, j, 3] - pc[1]) if np.isfinite(Gdata0[-1, j, 3]) else 0.0
            ext[:, j, 3] = np.polyval(pc, te) + off * np.ravel(decay)

    info = dict(fit_rms=fit_rms, model=model, srp=srp, ext_start=Gdt_ext[0], method=method, dense=dense)
    return np.concatenate((Gdata0, ext), axis=0), np.concatenate((Gdt0, Gdt_ext)), info