# ASA.py
# ============================================================
# ASA: Adaptive Singular-spectrum Analysis
#
# Components
# ----------
# 1) SSA encoder (linear):
#    - Build Hankel windows W_t (L,D) from training series F (N,D)
#    - Flatten each window into z_t in R^{L*D}
#    - Learn top-r principal directions via randomized SVD / PCA on window space
#    - Encode any window: psi = (vec(W) - mean_window) @ V_r   in R^r
#
# 2) LGP decoder (linear-kernel GP / ridge in primal):
#    - Learn global one-step map: psi_t -> x_{t+L}  (predict next sample)
#
# 3) Markovian In-Context Mechanism (nonlinear):
#    - At inference, if prefix length ell > L:
#        context pairs: (psi_i, residual_i) for i=0..K_ctx-1
#        residual_i = y_true_i - y_glob(psi_i)
#      Fit dense RBF GP on residuals in psi-space and apply during rollout:
#        y = y_glob(psi_q) + w_eff(var_q) * e_gp_mean(psi_q)
#      (Markovian: correction depends only on current psi_q.)
#
# Notes
# -----
# - This is essentially "LISA with SSA+LGP instead of NLSA+GPLM".
# - Dense IC GP is O(M^3) with M=context points used; use ctx_max_points to cap.
#
# Dependencies
# ------------
# numpy, scipy
# and LGP.py (this repo)
# ============================================================

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Literal, Optional, Tuple, Union

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

from scipy.linalg import cho_factor, cho_solve, solve_triangular

try:
    from .LGP import LGP  # type: ignore
except Exception:
    from LGP import LGP  # type: ignore


# ============================================================
# Helpers
# ============================================================

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 _flatten_windows(W_BLD: np.ndarray) -> np.ndarray:
    """
    (B,L,D) -> (B, L*D)
    """
    W = np.asarray(W_BLD, dtype=np.float64)
    if W.ndim == 2:
        W = W[None, :, :]
    B, L, D = W.shape
    return np.ascontiguousarray(W.reshape(B, L * D))


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))


# ============================================================
# Randomized SVD (Halko et al.) for SSA window PCA
# ============================================================

@dataclass
class _RandSVDParams:
    oversample: int = 8
    n_iter: int = 2
    seed: int = 0


def _randomized_top_right_singular_vectors(X: np.ndarray, r: int, *, params: _RandSVDParams) -> np.ndarray:
    """
    Return V_r (p,r) approximate top right singular vectors of X (n,p),
    using randomized SVD.

    Complexity: O(n*p*(r+o)) + O((r+o)^2*(n+p))
    """
    X = np.asarray(X, dtype=np.float64)
    n, p = X.shape
    r = int(min(max(1, r), p))
    k = int(min(p, r + int(params.oversample)))

    rng = np.random.default_rng(int(params.seed))
    Omega = rng.standard_normal(size=(p, k))  # (p,k)

    # Y = X Omega
    Y = X @ Omega  # (n,k)

    # Power iterations: (X X^T)^q X Omega improves separation
    for _ in range(int(params.n_iter)):
        Y = X @ (X.T @ Y)  # (n,k)

    # Orthonormalize Y -> Q
    Q, _ = np.linalg.qr(Y, mode="reduced")  # (n,k)

    # Small matrix B = Q^T X
    B = Q.T @ X  # (k,p)

    # SVD of B
    # B = Uhat S V^T => right vectors of X approx V
    _, _, Vt = np.linalg.svd(B, full_matrices=False)
    V = Vt.T  # (p,k)

    return np.ascontiguousarray(V[:, :r])


# ============================================================
# ASA main class
# ============================================================

SubsampleMode = Literal["uniform", "random"]


class ASA:
    """
    ASA: SSA encoder + LGP global head + nonlinear Markovian IC GP on residuals.

    Baseline:
      window W (L,D) -> psi(W) in R^r (SSA linear projection)
      y_glob = LGP(psi) in R^D

    In-context (Markovian nonlinear residual GP):
      For prefix length ell >= L:
        K_ctx = ell - L context pairs
        psi_i = encode(window_i)
        e_i = y_true_i - y_glob(psi_i)
      Fit RBF GP on (psi_i -> e_i), apply during rollout.
    """

    def __init__(
        self,
        F_tX: np.ndarray,
        *,
        L: int,
        rank: int,
        # ---------------- SSA / PCA encoder ----------------
        center_windows: bool = True,
        whiten_windows: bool = False,         # standardize window coords before PCA (optional)
        randomized_svd: bool = True,
        rsvd_oversample: int = 8,
        rsvd_n_iter: int = 2,
        seed: int = 0,
        # ---------------- LGP decoder ----------------------
        lgp_kwargs: Optional[dict] = None,
        # ---------------- IC GP controls -------------------
        ctx_min_windows: Optional[int] = None,
        ctx_max_points: Optional[int] = 2000,  # cap dense GP points (set None to use all)
        ctx_subsample: SubsampleMode = "uniform",
        ctx_k0: float = 10.0,                 # base mixing K/(K+ctx_k0)
        gp_noise2: float = 1e-3,              # σ_n^2 in (K + σ_n^2 I)
        gp_rbf_ell: Optional[float] = None,   # fixed, or auto from context
        gp_rbf_q: float = 0.5,
        # ---------------- trust gating ---------------------
        use_var_gate: bool = True,
        gate_tau2: float = 1.0,
        gate_mode: Literal["rational", "exp"] = "rational",
    ):
        if lgp_kwargs is None:
            lgp_kwargs = {}

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

        F = _as_2d(F_tX).astype(np.float64, copy=False)
        self.N, self.D = F.shape

        self.L = int(L)
        self.r = int(rank)
        if self.L < 2:
            raise ValueError("ASA requires L>=2 for meaningful SSA windows.")
        if self.N <= self.L + 2:
            raise ValueError("Training series too short for this L.")

        # ---- Build training windows ----
        W_all = _sliding_windows(F, self.L)  # (K,L,D)
        K = W_all.shape[0]
        N_pairs = K - 1
        if N_pairs < 4:
            raise ValueError("Not enough training pairs for ASA (increase N or reduce L).")

        Z_all = _flatten_windows(W_all)       # (K, P), P=L*D
        P = Z_all.shape[1]

        # ---- Window preprocessing for encoder ----
        self.center_windows = bool(center_windows)
        self.whiten_windows = bool(whiten_windows)

        if self.center_windows:
            self.win_mean = Z_all.mean(axis=0)
        else:
            self.win_mean = np.zeros((P,), dtype=np.float64)

        Zc = Z_all - self.win_mean[None, :]

        if self.whiten_windows:
            self.win_std = np.maximum(Zc.std(axis=0), 1e-12)
        else:
            self.win_std = np.ones((P,), dtype=np.float64)

        Zp = Zc / self.win_std[None, :]

        # ---- SSA basis via PCA (top right singular vectors) ----
        r_eff = int(min(max(1, self.r), P))
        self.r = r_eff

        if randomized_svd:
            V_r = _randomized_top_right_singular_vectors(
                Zp, r_eff,
                params=_RandSVDParams(
                    oversample=int(rsvd_oversample),
                    n_iter=int(rsvd_n_iter),
                    seed=int(seed),
                ),
            )
        else:
            # Exact SVD (can be heavy for large K,P)
            _, _, Vt = np.linalg.svd(Zp, full_matrices=False)
            V_r = Vt.T[:, :r_eff]

        self.V_r = np.ascontiguousarray(V_r, dtype=np.float64)  # (P,r)

        # ---- Encode training windows (psi) ----
        Psi_all = Zp @ self.V_r                  # (K,r)
        Psi_train = np.ascontiguousarray(Psi_all[:N_pairs, :])  # (K-1,r)

        # ---- Targets: next sample after each window ----
        Y_train = np.ascontiguousarray(F[self.L:self.L + N_pairs, :])  # (K-1,D)

        # ---- Fit global decoder (LGP) ----
        lkw = dict(lgp_kwargs)
        lkw.setdefault("sigma2", 1e-5)
        lkw.setdefault("jitter", 1e-10)
        lkw.setdefault("center_X", True)
        lkw.setdefault("whiten_latent", True)   # often helps linear features
        lkw.setdefault("solver", "chol")

        self.lgp = LGP(Psi_train, Y_train, **lkw)

        # ---- 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(8, self.r + 1)

        self.ctx_max_points = None if ctx_max_points is None else int(ctx_max_points)
        self.ctx_subsample = str(ctx_subsample).lower().strip()
        if self.ctx_subsample not in ("uniform", "random"):
            raise ValueError("ctx_subsample must be 'uniform' or 'random'")

        # ---- GP residual controls ----
        self.gp_noise2 = float(gp_noise2)
        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

        # ---- 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'")

    # ============================================================
    # SSA encoder
    # ============================================================

    def encode_window(self, W_LD: np.ndarray) -> np.ndarray:
        """
        Encode one window (L,D) -> (r,)
        """
        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}")

        z = W.reshape(-1).astype(np.float64, copy=False)  # (P,)
        zc = z - self.win_mean
        zp = zc / self.win_std
        return zp @ self.V_r

    def _encode_batch(self, W_BLD: np.ndarray) -> np.ndarray:
        """
        Vectorized encode windows (B,L,D) -> (B,r)
        """
        Z = _flatten_windows(W_BLD)  # (B,P)
        Zc = Z - self.win_mean[None, :]
        Zp = Zc / self.win_std[None, :]
        return np.ascontiguousarray(Zp @ self.V_r)

    # ============================================================
    # GP kernel utilities (RBF on psi-space)
    # ============================================================

    def _gp_kernel_matrix(self, Psi_ctx: np.ndarray) -> Tuple[np.ndarray, float]:
        Psi_ctx = np.asarray(Psi_ctx, dtype=np.float64)
        d2 = _pairwise_sq_dists(Psi_ctx)
        if self.gp_rbf_ell is None:
            ell = _estimate_rbf_ell_from_d2(d2, q=self.gp_rbf_q)
        else:
            ell = float(self.gp_rbf_ell)
        self.last_gp_rbf_ell_ = float(ell)
        K = np.exp(-0.5 * d2 / (ell**2 + 1e-12))
        return K, float(ell)

    def _gp_kernel_eval(self, Psi_ctx: np.ndarray, psi_q: np.ndarray, ell: float) -> np.ndarray:
        Psi_ctx = np.asarray(Psi_ctx, dtype=np.float64)
        psi_q = np.asarray(psi_q, dtype=np.float64)
        diff = Psi_ctx - psi_q[None, :]
        d2 = np.einsum("kr,kr->k", diff, diff, optimize=True)
        return np.exp(-0.5 * d2 / (ell**2 + 1e-12))

    def _gate_from_var(self, var_f: float) -> float:
        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:
        """
        Baseline one-step prediction from a single window (L,D).
        """
        psi = self.encode_window(W_LD)
        return self.lgp(psi)

    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 (ell,D), ell >= L.

        If ell == L: baseline AR only (SSA+LGP).
        If ell  > L: Markovian IC GP on residuals in psi-space.
        """
        prefix = _as_2d(prefix).astype(np.float64, copy=False)
        ell, D = prefix.shape
        if D != self.D:
            raise ValueError(f"ASA 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 for rollout
        cur = prefix[-self.L:, :].copy()

        # ---------------------------
        # Baseline only if no context
        # ---------------------------
        K_ctx = ell - self.L
        if K_ctx <= 0 or K_ctx < 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

        # ---------------------------
        # Build context (windows + targets)
        # ---------------------------
        W_all = _sliding_windows(prefix, self.L)            # (ell-L+1, L, D)
        W_ctx = np.ascontiguousarray(W_all[:K_ctx, :, :])   # (K_ctx, L, D)
        Y_ctx = np.ascontiguousarray(prefix[self.L:self.L + K_ctx, :])  # (K_ctx, D)

        # encode all context windows
        Psi_ctx_full = self._encode_batch(W_ctx)            # (K_ctx, r)

        # optionally subsample context to keep dense GP feasible
        if self.ctx_max_points is not None and K_ctx > self.ctx_max_points:
            M = int(self.ctx_max_points)
            if self.ctx_subsample == "uniform":
                idx = np.linspace(0, K_ctx - 1, M).round().astype(np.int64)
            else:
                idx = self._rng.choice(K_ctx, size=M, replace=False).astype(np.int64)
            Psi_ctx = Psi_ctx_full[idx]
            Y_ctx_use = Y_ctx[idx]
        else:
            Psi_ctx = Psi_ctx_full
            Y_ctx_use = Y_ctx

        M = Psi_ctx.shape[0]

        # baseline predictions on context
        Y_glob_ctx = self.lgp(Psi_ctx)                      # (M, D)

        # residuals to learn in-context
        E_ctx = Y_ctx_use - Y_glob_ctx                      # (M, D)

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

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

        # base mixing strength from amount of context
        w_ctx_base = float(M) / float(M + 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

        # ---------------------------
        # Rollout with Markovian IC correction
        # ---------------------------
        for h in range(H):
            psi_q = self.encode_window(cur)                 # (r,)

            # baseline
            y_glob = self.lgp(psi_q)                        # (D,)

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

            # function variance: var_f = 1 - k^T (K+σ^2I)^{-1} k
            u = solve_triangular(Lfac, k_eval, lower=lower, check_finite=False)
            quad = float(np.dot(u, u))
            var_f = max(0.0, 1.0 - quad)

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

            # trust gate
            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,))

            y = y_glob + w_eff * e_use
            preds[h] = y

            # update 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

    # ============================================================
    # Baseline rollout
    # ============================================================

    def _rollout_baseline(self, seed_LD: np.ndarray, H: int) -> np.ndarray:
        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.encode_window(cur)
            y = self.lgp(psi)
            out[h] = y
            if self.L > 1:
                cur[:-1] = cur[1:]
            cur[-1] = y
        return out


__all__ = ["ASA"]