# LISA.py
# ============================================================
# LISA: NLSA encoder + GPLM baseline decoder
#       + (optional) dense in-context Gaussian Process Regression (GPR)
#       on residuals in diffusion coordinate space.
#
# This version is compatible with the "sparse-kernel KRR" GPLM.py:
#   - GPLM trains by building a sparse kNN kernel on training latents ψ
#   - solves (K + sigma2 I) S = Y (mean GP / KRR)
#   - predicts using only pred_k neighbors per query
#
# Key advantages vs Nyström GPLM:
#   - No anchors/inducing points required
#   - Strong nonlinearity => very local sparse kernel => scalable
#
# ============================================================

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

from scipy.linalg import cho_factor, cho_solve, solve_triangular


# ------------------------------------------------------------
# Robust imports (package or flat)
# ------------------------------------------------------------
try:
    from .NLSA import NLSA  # type: ignore
except Exception:
    from NLSA import NLSA  # type: ignore

try:
    from .GPLMx import GPLM  # type: ignore
except Exception:
    from GPLMx import GPLM  # type: ignore


# ============================================================
# Small utilities
# ============================================================

def _as_2d(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=float)
    if X.ndim == 1:
        X = X[:, None]
    return X


def _sliding_windows(F_tD: np.ndarray, L: int) -> np.ndarray:
    """
    Return windows W (K,L,D) from F (N,D) with K=N-L+1.
    Handles possible (K,D,L) output from sliding_window_view.
    """
    F = _as_2d(F_tD)
    W = sliding_window_view(F, window_shape=int(L), axis=0)

    a, b = W.shape[1], W.shape[2]
    if (a, b) == (L, F.shape[1]):
        return np.ascontiguousarray(W)
    if (a, b) == (F.shape[1], L):
        return np.ascontiguousarray(np.transpose(W, (0, 2, 1)))

    raise ValueError(f"Unexpected window shape {W.shape} for L={L}, D={F.shape[1]}.")


def _pairwise_sq_dists(X: np.ndarray) -> np.ndarray:
    """
    Dense pairwise squared Euclidean distances (n,n) for X (n,r).
    """
    X = np.asarray(X, dtype=np.float64)
    x2 = np.sum(X * X, axis=1, keepdims=True)
    d2 = x2 + x2.T - 2.0 * (X @ X.T)
    np.maximum(d2, 0.0, out=d2)
    return d2


def _estimate_rbf_ell_from_d2(
    d2_mat: np.ndarray,
    q: float = 0.5,
    eps: float = 1e-12,
    min_ell: float = 1e-6,
) -> float:
    """
    If kernel is exp(-||x-y||^2/(2 ell^2)), a good heuristic is:
      ell^2 ~= quantile(d^2)/2
    """
    iu = np.triu_indices_from(d2_mat, k=1)
    vals = d2_mat[iu]
    if vals.size == 0:
        return 1.0
    qv = float(np.quantile(vals, q))
    ell = np.sqrt(max(qv, 0.0) / 2.0 + eps)
    return float(max(ell, min_ell))


# ============================================================
# LISA main class
# ============================================================

class LISA:
    """
    LISA: NLSA encoder + GPLM baseline + optional dense GPR residual correction.

    - Baseline:
        window W (L,D) -> ψ(W) in R^r (via NLSA Nyström OOS)
        y_glob = GPLM(ψ) in R^D

    - In-context (optional):
        Given prefix length ℓ > L, fit dense GP on residuals e(ψ)
        using context windows inside prefix, then apply during rollout.

    Note:
      The IC GP here is dense and scales O(K_ctx^2) memory / O(K_ctx^3) time.
      A default cap gp_max_ctx prevents accidental blowups with huge prefixes.
    """

    def __init__(
        self,
        F_tX: np.ndarray,
        *,
        L: int,
        rank: int,
        # ------------------ NLSA encoder hyperparams ------------------
        beta: float | None = None,
        alpha: float = 1.0,
        center_outputs: bool = True,
        drop_first: bool = True,
        seed: int = 0,
        nlsa_kwargs: Optional[dict] = None,
        # ------------------ GPLM baseline decoder hyperparams ----------
        gplm_kwargs: Optional[dict] = None,
        # ------------------ IC / GPR controls -------------------------
        ctx_min_windows: Optional[int] = None,
        ctx_k0: float = 10.0,                 # base mixing K_ctx/(K_ctx+ctx_k0)
        gp_noise2: float = 1e-3,              # σ_n^2 in (K + σ_n^2 I)
        gp_kernel: str = "rbf",               # "rbf" or "linear"
        gp_rbf_ell: float | None = None,      # fixed, or auto from context
        gp_rbf_q: float = 0.5,
        # ---- NEW safety: cap dense GP context used in __call__ -------
        gp_max_ctx: Optional[int] = 4096,     # None disables cap
        gp_ctx_mode: str = "recent",          # "recent" or "uniform"
        # ------------------ stability / trust gating -----------------
        use_var_gate: bool = True,
        gate_tau2: float = 1.0,
        gate_mode: str = "rational",          # "rational" or "exp"
    ):
        if nlsa_kwargs is None:
            nlsa_kwargs = {}
        if gplm_kwargs is None:
            gplm_kwargs = {}

        self._rng = np.random.default_rng(int(seed))

        # ---- Encoder ----
        self.base = NLSA(
            F_tX,
            L=int(L),
            rank=int(rank),
            beta=beta,
            alpha=float(alpha),
            center=bool(center_outputs),
            drop_first=bool(drop_first),
            seed=int(seed),
            **nlsa_kwargs,
        )

        base = self.base
        if base.psi_ is None or base.psi_.shape[1] == 0:
            raise ValueError("NLSA produced zero latent dims. Increase rank or set drop_first=False.")

        self.L = int(base.L)
        self.D = int(base.D_)
        self.r = int(base.psi_.shape[1])

        # mean handling (consistent w/ NLSA centering)
        self.center_outputs = bool(base.center)
        self.mu_X = base.mu_.reshape(-1).astype(np.float64) if self.center_outputs else np.zeros((self.D,), dtype=np.float64)

        # ---- Train baseline GPLM: ψ -> next sample (centered) ----
        K = int(base.K_)
        N_pairs = K - 1
        if N_pairs < 4:
            raise ValueError("Training series too short for LISA with this L.")

        X_train = np.asarray(base.psi_[:N_pairs, :], dtype=np.float64)  # (K-1,r)
        Y_train_c = np.asarray(
            base.R_[self.L : self.L + N_pairs, :],
            dtype=np.float64
        )  # (K-1,D) centered if base.center

        # ---- Defaults tuned for sparse-kernel GPLM ----
        gkw = dict(gplm_kwargs)
        gkw.setdefault("seed", int(seed))

        # IMPORTANT: LISA already provides centered Y, so disable GPLM centering
        gkw.setdefault("center_X", False)

        # Regularization for sparse KRR (lambda in K + lambda I)
        # (You may tune this up if solves are noisy/unstable.)
        gkw.setdefault("sigma2", 1e-4)

        # Numerical diagonal stabilizer
        gkw.setdefault("jitter", 1e-8)

        # Sparse kernel build on ψ_train
        gkw.setdefault("k_graph", 64)          # training graph sparsity
        gkw.setdefault("mutual_knn", True)
        gkw.setdefault("include_self", True)
        gkw.setdefault("normalize_rows", False)

        # eps estimation for RBF weights (median kNN distance)
        gkw.setdefault("k_eps", 256)
        gkw.setdefault("eps_use_kth", True)
        gkw.setdefault("eps_mul", 1.0)

        # Solver controls for (K + lambda I) S = Y
        gkw.setdefault("solve_method", "auto")
        gkw.setdefault("solve_tol", 1e-6)
        gkw.setdefault("solve_maxiter", 800)
        gkw.setdefault("solve_verbose", False)

        # Inference neighbor truncation
        gkw.setdefault("pred_k", 128)

        # Latent preconditioning (optional)
        # On ψ it sometimes helps, sometimes hurts. Leave default False unless needed.
        gkw.setdefault("whiten_latent", False)

        # dtype for ANN storage
        gkw.setdefault("dtype", np.float32)

        self.gplm = GPLM(X_train, Y_train_c, **gkw)

        # ---- IC controls ----
        self.ctx_k0 = float(ctx_k0)
        self.ctx_min_windows = int(ctx_min_windows) if ctx_min_windows is not None else max(1, self.r + 1)

        # ---- GP residual controls ----
        self.gp_noise2 = float(gp_noise2)
        self.gp_kernel = str(gp_kernel).lower().strip()
        if self.gp_kernel not in ("rbf", "linear"):
            raise ValueError("gp_kernel must be 'rbf' or 'linear'")

        self.gp_rbf_ell = None if gp_rbf_ell is None else float(gp_rbf_ell)
        self.gp_rbf_q = float(gp_rbf_q)
        self.last_gp_rbf_ell_: Optional[float] = None

        # Dense GP context cap
        self.gp_max_ctx = None if gp_max_ctx is None else int(gp_max_ctx)
        self.gp_ctx_mode = str(gp_ctx_mode).lower().strip()
        if self.gp_ctx_mode not in ("recent", "uniform"):
            raise ValueError("gp_ctx_mode must be 'recent' or 'uniform'")

        # ---- variance gating ----
        self.use_var_gate = bool(use_var_gate)
        self.gate_tau2 = float(gate_tau2)
        self.gate_mode = str(gate_mode).lower().strip()
        if self.gate_mode not in ("rational", "exp"):
            raise ValueError("gate_mode must be 'rational' or 'exp'")

    # ============================================================
    # Baseline utilities
    # ============================================================

    def _encode_batch(self, W_BLD: np.ndarray) -> np.ndarray:
        """
        Encode windows (B,L,D) -> (B,r) using base.encode_window in a loop.
        """
        W = np.asarray(W_BLD, dtype=np.float64)
        if W.ndim == 2:
            W = W[None, :, :]
        B = W.shape[0]
        out = np.zeros((B, self.r), dtype=np.float64)
        for i in range(B):
            out[i] = self.base.encode_window(W[i])
        return out

    def _baseline_centered_batch(self, Psi_Br: np.ndarray) -> np.ndarray:
        """
        GPLM baseline in centered output space: (B,r) -> (B,D)
        """
        Psi = np.asarray(Psi_Br, dtype=np.float64)
        if Psi.ndim == 1:
            Psi = Psi[None, :]
        Yc = self.gplm(Psi)  # (B,D) (already centered because center_X=False + mean_X=0)
        return np.asarray(Yc, dtype=np.float64)

    def _baseline_centered_one(self, psi_r: np.ndarray) -> np.ndarray:
        return self._baseline_centered_batch(np.asarray(psi_r, dtype=np.float64))[0]

    # ============================================================
    # GP kernel utilities (residual GP)
    # ============================================================

    def _gp_kernel_matrix(self, Psi_ctx: np.ndarray) -> Tuple[np.ndarray, Optional[float]]:
        """
        Build dense kernel matrix K(Ψ,Ψ) on context points.
        Returns (K_mat, ell_used).
        """
        Psi_ctx = np.asarray(Psi_ctx, dtype=np.float64)

        if self.gp_kernel == "linear":
            return Psi_ctx @ Psi_ctx.T, None

        d2 = _pairwise_sq_dists(Psi_ctx)
        if self.gp_rbf_ell is None:
            ell_used = _estimate_rbf_ell_from_d2(d2, q=self.gp_rbf_q)
        else:
            ell_used = float(self.gp_rbf_ell)

        self.last_gp_rbf_ell_ = ell_used
        K = np.exp(-0.5 * d2 / (ell_used**2 + 1e-12))
        return K, ell_used

    def _gp_kernel_eval(self, Psi_ctx: np.ndarray, psi_q: np.ndarray, ell_used: Optional[float]) -> np.ndarray:
        """
        k(Ψ,ψ_q) for query ψ_q. Returns (K_ctx,).
        """
        Psi_ctx = np.asarray(Psi_ctx, dtype=np.float64)
        psi_q = np.asarray(psi_q, dtype=np.float64)

        if self.gp_kernel == "linear":
            return Psi_ctx @ psi_q

        assert ell_used is not None
        diff = Psi_ctx - psi_q[None, :]
        d2 = np.einsum("kr,kr->k", diff, diff, optimize=True)
        return np.exp(-0.5 * d2 / (ell_used**2 + 1e-12))

    def _k_qq(self, psi_q: np.ndarray) -> float:
        """
        Kernel self-similarity k(ψ,ψ).
        """
        psi_q = np.asarray(psi_q, dtype=np.float64)
        if self.gp_kernel == "linear":
            return float(np.dot(psi_q, psi_q))
        return 1.0  # RBF

    def _gate_from_var(self, var_f: float) -> float:
        """
        Convert predictive function variance -> trust weight in [0,1].
        """
        if not self.use_var_gate:
            return 1.0

        v = max(float(var_f), 0.0)
        tau2 = max(float(self.gate_tau2), 1e-18)

        if self.gate_mode == "exp":
            return float(np.exp(-v / tau2))
        return float(tau2 / (tau2 + v))

    # ============================================================
    # Public API
    # ============================================================

    def predict_one_step(self, W_LD: np.ndarray) -> np.ndarray:
        """
        Predict one step from a single window (L,D) using baseline only.
        Returns (D,) in original (uncentered) scale.
        """
        W = np.asarray(W_LD, dtype=np.float64)
        if W.ndim == 1:
            W = W[:, None]
        if W.shape != (self.L, self.D):
            raise ValueError(f"Expected window shape {(self.L, self.D)}, got {W.shape}")

        psi = self.base.encode_window(W)
        y_c = self._baseline_centered_one(psi)
        return y_c + self.mu_X if self.center_outputs else y_c

    def __call__(
        self,
        prefix: np.ndarray,
        steps: int = 1,
        *,
        return_var: bool = False,
        sample: bool = False,
        rng: Optional[np.random.Generator] = None,
        include_obs_noise: bool = True,
    ):
        """
        Autoregressive forecast from prefix (ℓ,D), ℓ >= L.

        If ℓ == L: baseline AR only (NLSA + sparse GPLM).
        If ℓ >  L: uses dense in-context GPR on residuals (optional).
        """
        prefix = _as_2d(prefix)
        ell, D = prefix.shape
        if D != self.D:
            raise ValueError(f"LISA trained with D={self.D}, got prefix D={D}.")
        if ell < self.L:
            raise ValueError(f"Need prefix length ell >= L={self.L}.")

        H = int(steps)
        if H <= 0:
            out = np.zeros((0, self.D), dtype=np.float64)
            return (out, np.zeros((0,), dtype=np.float64)) if return_var else out

        if rng is None:
            rng = self._rng

        # seed window for AR rollout
        cur = prefix[-self.L:, :].copy()

        # ---------------------------
        # If IC nullset -> baseline AR
        # ---------------------------
        K_ctx_full = ell - self.L
        if K_ctx_full <= 0 or self.r == 0:
            preds = self._rollout_baseline(cur, H)
            if return_var:
                return preds[0] if H == 1 else preds, np.zeros((H,), dtype=np.float64)
            return preds[0] if H == 1 else preds

        if K_ctx_full < self.ctx_min_windows:
            preds = self._rollout_baseline(cur, H)
            if return_var:
                return preds[0] if H == 1 else preds, np.zeros((H,), dtype=np.float64)
            return preds[0] if H == 1 else preds

        # ---------------------------
        # Dense GP context cap (safety)
        # ---------------------------
        if self.gp_max_ctx is None:
            K_ctx = int(K_ctx_full)
            start = 0
            ctx_idx = None
        else:
            K_ctx = int(min(K_ctx_full, self.gp_max_ctx))
            if self.gp_ctx_mode == "recent":
                start = int(K_ctx_full - K_ctx)
                ctx_idx = None
            else:
                ctx_idx = np.linspace(0, K_ctx_full - 1, num=K_ctx, dtype=np.int64)
                start = 0

        # ---------------------------
        # Build context windows + targets
        # ---------------------------
        W_all = _sliding_windows(prefix, self.L)  # (ell-L+1,L,D)

        if ctx_idx is None:
            W_ctx = np.ascontiguousarray(W_all[start : start + K_ctx, :, :])
            Y_ctx = np.asarray(prefix[self.L + start : self.L + start + K_ctx, :], dtype=np.float64)
        else:
            W_ctx = np.ascontiguousarray(W_all[ctx_idx, :, :])
            Y_ctx = np.asarray(prefix[self.L + ctx_idx, :], dtype=np.float64)

        # center targets consistently with baseline training
        Y_ctx_c = (Y_ctx - self.mu_X[None, :]) if self.center_outputs else Y_ctx

        # encode ψ for context
        Psi_ctx = self._encode_batch(W_ctx)  # (K_ctx,r)

        # baseline predictions on context
        Y_glob_ctx_c = self._baseline_centered_batch(Psi_ctx)  # (K_ctx,D)

        # residual table
        E_ctx = Y_ctx_c - Y_glob_ctx_c  # (K_ctx,D)

        # ---------------------------
        # Fit dense GP on residuals
        # ---------------------------
        K_mat, ell_used = self._gp_kernel_matrix(Psi_ctx)  # (K_ctx,K_ctx)
        K_reg = K_mat + self.gp_noise2 * np.eye(K_ctx, dtype=np.float64)

        cf = cho_factor(K_reg, lower=True, check_finite=False)
        alpha = cho_solve(cf, E_ctx, check_finite=False)  # (K_ctx,D)
        Lfac, lower = cf

        w_ctx_base = float(K_ctx) / float(K_ctx + self.ctx_k0) if self.ctx_k0 > 0 else 1.0

        preds = np.zeros((H, self.D), dtype=np.float64)
        vars_out = np.zeros((H,), dtype=np.float64) if return_var else None

        # ---------------------------
        # AR rollout with GP residual
        # ---------------------------
        for h in range(H):
            psi_q = self.base.encode_window(cur)

            # baseline
            y_glob_c = self._baseline_centered_one(psi_q)

            # GP residual mean
            k_eval = self._gp_kernel_eval(Psi_ctx, psi_q, ell_used)
            e_mean = k_eval @ alpha

            # GP residual variance (function variance)
            u = solve_triangular(Lfac, k_eval, lower=lower, check_finite=False)
            quad = float(np.dot(u, u))
            var_f = max(0.0, self._k_qq(psi_q) - quad)

            if return_var:
                vars_out[h] = float(var_f)

            w_gate = self._gate_from_var(var_f)
            w_eff = w_ctx_base * w_gate

            e_use = e_mean
            if sample:
                var_y = var_f + (self.gp_noise2 if include_obs_noise else 0.0)
                var_y = max(0.0, float(var_y))
                if var_y > 0:
                    e_use = e_mean + np.sqrt(var_y) * rng.standard_normal(size=(self.D,))

            # combine
            y_c = y_glob_c + w_eff * e_use
            y = (y_c + self.mu_X) if self.center_outputs else y_c

            preds[h] = y

            # update rolling window
            if self.L > 1:
                cur[:-1] = cur[1:]
            cur[-1] = y

        if return_var:
            if H == 1:
                return preds[0], vars_out
            return preds, vars_out
        return preds[0] if H == 1 else preds

    def _rollout_baseline(self, seed_LD: np.ndarray, H: int) -> np.ndarray:
        """
        Baseline AR rollout only: NLSA encode + sparse GPLM decode.
        """
        cur = np.asarray(seed_LD, dtype=np.float64).copy()
        out = np.zeros((H, self.D), dtype=np.float64)

        for h in range(int(H)):
            psi = self.base.encode_window(cur)
            y_c = self._baseline_centered_one(psi)
            y = (y_c + self.mu_X) if self.center_outputs else y_c
            out[h] = y

            if self.L > 1:
                cur[:-1] = cur[1:]
            cur[-1] = y

        return out


__all__ = ["LISA"]