# LGP.py
# ============================================================
# LGP: Linear-GP (linear-kernel GP / KRR) decoder
#
# Goal:
#   A "good linear GPLM" for linear encoders (e.g., SSA/ISA/ASA linear codes).
#   It behaves like a GP/KRR with a *linear kernel*:
#       k(r, r') = r^T r'
#
# Training data:
#   R_ix : (N, d)   latent features / codes
#   R_iX : (N, D)   targets in ambient/output space
#
# GP/KRR mean predictor (dual form):
#   y(r) = k(r, R)^T (K + σ^2 I)^{-1} Y
# with K = R R^T.
#
# With a linear kernel, you can do this in the *primal* cheaply:
#   V = (R^T R + σ^2 I)^{-1} R^T Y     (d x D)
#   y(r) = r^T V                       (D,)
#
# This avoids building K (N x N) entirely.
#
# Uncertainty:
#   Under the GP interpretation (and for a BLR-equivalent prior), a cheap scalar
#   function-variance proxy is:
#     var_f(r) = σ^2 * r^T (R^T R + σ^2 I)^{-1} r
# (shared across output dims if you assume independent outputs sharing features).
#
# Complexity:
#   - Form Gram in latent space:  O(N d^2)
#   - Solve dxd system:           O(d^3) via Cholesky (cheap if d is small)
#   - Predict A queries:          O(A d D)
#
# Optional solver="cg" supports very large d using matvecs:
#   A v = (R^T (R v) + σ^2 v)
#
# Dependencies:
#   numpy, scipy
# ============================================================

from __future__ import annotations

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

import numpy as np
import scipy.linalg as la
import scipy.sparse.linalg as sla


Solver = Literal["chol", "eigh", "cg"]


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


@dataclass
class _LatentPreproc:
    mean: np.ndarray
    std: np.ndarray

    def apply(self, R: np.ndarray) -> np.ndarray:
        return (R - self.mean[None, :]) / self.std[None, :]


class LGP:
    """
    LGP = Linear-kernel GP mean (KRR) decoder.

    API is intentionally similar to GPLM:
      - __init__(...) fits immediately
      - __call__(R_ax, batch_size=None) predicts mean
      - predict(R_ax, return_var=True) returns mean + scalar var proxy

    Notes:
      - If center_X=True, we center outputs during training and add mean back at inference.
      - If whiten_latent=True, we standardize latents dimension-wise (recommended if feature scales vary).
      - sigma2 is the ridge/noise term (λ). Larger sigma2 -> smoother / more stable rollout.
    """

    def __init__(
        self,
        R_ix: np.ndarray,
        R_iX: np.ndarray,
        *,
        sigma2: float = 1e-5,
        jitter: float = 1e-10,
        center_X: bool = True,
        whiten_latent: bool = False,
        dtype: Any = np.float32,
        solver: Solver = "chol",
        cg_maxiter: int = 500,
        cg_rtol: float = 1e-6,
        cg_atol: float = 0.0,
        # numeric safety
        eig_clip: float = 1e-12,
    ):
        self.sigma2 = float(sigma2)
        self.jitter = float(jitter)
        self.center_X = bool(center_X)
        self.whiten_latent = bool(whiten_latent)
        self.dtype = dtype
        self.solver: Solver = str(solver)  # type: ignore

        self.cg_maxiter = int(cg_maxiter)
        self.cg_rtol = float(cg_rtol)
        self.cg_atol = float(cg_atol)
        self.eig_clip = float(eig_clip)

        # cast/validate
        R_ix = np.ascontiguousarray(np.asarray(R_ix).astype(self.dtype, copy=False))
        R_iX = np.ascontiguousarray(np.asarray(R_iX).astype(self.dtype, copy=False))
        if R_ix.ndim != 2 or R_iX.ndim != 2 or R_ix.shape[0] != R_iX.shape[0]:
            raise ValueError("R_ix must be (N,d) and R_iX must be (N,D) with same N.")

        self.R_ix = R_ix
        self.R_iX = R_iX
        self.N, self.d_lat = R_ix.shape
        _, self.D = R_iX.shape

        # output centering
        Y = R_iX.astype(np.float64)
        if self.center_X:
            self.mean_X = Y.mean(axis=0)
            Yc = Y - self.mean_X[None, :]
        else:
            self.mean_X = np.zeros((self.D,), dtype=np.float64)
            Yc = Y

        # latent whitening
        R = R_ix.astype(np.float64)
        if self.whiten_latent:
            mu = R.mean(axis=0)
            sd = np.maximum(R.std(axis=0), 1e-12)
            self._pp = _LatentPreproc(mu, sd)
            Rw = self._pp.apply(R)
        else:
            self._pp = _LatentPreproc(np.zeros((self.d_lat,), dtype=np.float64),
                                      np.ones((self.d_lat,), dtype=np.float64))
            Rw = R

        self.R_ix_w = np.ascontiguousarray(Rw, dtype=np.float64)  # (N,d)

        # Fit: V = (R^T R + (sigma2 + jitter) I)^-1 R^T Yc
        lam = float(self.sigma2)
        jit = float(self.jitter)

        # We'll store A = R^T R + lam I (+ jitter)
        # and a factorization/solver representation.
        self._A = None
        self._chol = None
        self._eig = None  # (w, V)
        self.V_xX = None  # (d,D)

        if self.solver in ("chol", "eigh"):
            self._fit_closed_form(Rw=self.R_ix_w, Yc=Yc, lam=lam, jit=jit)
        elif self.solver == "cg":
            self._fit_cg(Rw=self.R_ix_w, Yc=Yc, lam=lam, jit=jit)
        else:
            raise ValueError("solver must be one of {'chol','eigh','cg'}")

    # ------------------------------------------------------------
    # Training implementations
    # ------------------------------------------------------------
    def _fit_closed_form(self, *, Rw: np.ndarray, Yc: np.ndarray, lam: float, jit: float) -> None:
        # A = R^T R + lam I
        G = Rw.T @ Rw  # (d,d)
        d = G.shape[0]
        A = G + (lam + jit) * np.eye(d, dtype=np.float64)
        B = Rw.T @ Yc  # (d,D)

        if self.solver == "eigh":
            w, V = np.linalg.eigh(A)
            w = np.maximum(w, float(self.eig_clip))
            self._eig = (w, V)
            # V_xX = A^{-1} B = V diag(1/w) V^T B
            self.V_xX = (V @ ((V.T @ B) / w[:, None])).astype(np.float64)
            self._A = A
            return

        # chol (default)
        cF = la.cho_factor(A, lower=True, check_finite=False)
        self._chol = cF
        self.V_xX = la.cho_solve(cF, B, check_finite=False).astype(np.float64)
        self._A = A

    def _fit_cg(self, *, Rw: np.ndarray, Yc: np.ndarray, lam: float, jit: float) -> None:
        """
        Solve (R^T R + lam I) V = R^T Y with CG in d-space using matvecs.
        Good when d is too large to factorize directly.
        """
        N, d = Rw.shape
        lam_eff = lam + jit

        def matvec(v: np.ndarray) -> np.ndarray:
            # A v = R^T (R v) + lam v
            v = np.asarray(v, dtype=np.float64)
            return (Rw.T @ (Rw @ v)) + lam_eff * v

        Aop = sla.LinearOperator((d, d), matvec=matvec, dtype=np.float64)

        B = Rw.T @ Yc  # (d,D)
        V = np.zeros((d, self.D), dtype=np.float64)

        # Solve each output dim independently (D is often small; if D is large, consider block-CG)
        for j in range(self.D):
            x0 = np.zeros((d,), dtype=np.float64)
            sol, info = sla.cg(
                Aop,
                B[:, j],
                x0=x0,
                maxiter=self.cg_maxiter,
                rtol=self.cg_rtol,
                atol=self.cg_atol,
            )
            if info != 0:
                raise RuntimeError(f"LGP(CD) CG failed for output dim {j} with info={info}. "
                                   f"Try larger cg_maxiter, looser rtol, or larger sigma2.")
            V[:, j] = sol

        self.V_xX = V
        self._A = None
        self._chol = None
        self._eig = None

    # ------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------
    def __call__(self, R_ax: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
        """
        Mean prediction only.
        R_ax: (A,d) or (d,)
        returns: (A,D) or (D,)
        """
        R_ax = np.asarray(R_ax)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]

        Ra = np.ascontiguousarray(R_ax.astype(np.float64, copy=False))
        if Ra.shape[1] != self.d_lat:
            raise ValueError(f"Expected latent dim d={self.d_lat}, got {Ra.shape[1]}.")

        if batch_size is None:
            Y = self._decode_mean(Ra)
        else:
            bs = int(batch_size)
            out = []
            for s in range(0, Ra.shape[0], bs):
                out.append(self._decode_mean(Ra[s:s+bs]))
            Y = np.vstack(out)

        return Y[0] if single else Y

    def predict(
        self,
        R_ax: Union[np.ndarray, list],
        *,
        return_var: bool = False,
        batch_size: Optional[int] = None,
        include_obs_noise: bool = False,
    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
        """
        Predict mean and optional scalar variance proxy per query.

        var_f(a) = sigma2 * r_a^T (R^T R + sigma2 I)^{-1} r_a

        If include_obs_noise=True, returns var_y = var_f + sigma2.
        """
        R_ax = np.asarray(R_ax)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]

        Ra = np.ascontiguousarray(R_ax.astype(np.float64, copy=False))
        if Ra.shape[1] != self.d_lat:
            raise ValueError(f"Expected latent dim d={self.d_lat}, got {Ra.shape[1]}.")

        if not return_var:
            mean = self.__call__(Ra, batch_size=batch_size)
            if single and mean.ndim == 1:
                mean = mean[None, :]
            return mean[0] if single else mean

        if batch_size is None:
            mean, var = self._decode_mean_var(Ra, include_obs_noise=include_obs_noise)
        else:
            bs = int(batch_size)
            ms, vs = [], []
            for s in range(0, Ra.shape[0], bs):
                m, v = self._decode_mean_var(Ra[s:s+bs], include_obs_noise=include_obs_noise)
                ms.append(m)
                vs.append(v)
            mean = np.vstack(ms)
            var = np.concatenate(vs, axis=0)

        if single:
            return mean[0], var[0]
        return mean, var

    # ------------------------------------------------------------
    # core decode routines
    # ------------------------------------------------------------
    def _decode_mean(self, Ra: np.ndarray) -> np.ndarray:
        assert self.V_xX is not None

        # apply latent whitening
        Raw = self._pp.apply(Ra)

        Yc = Raw @ self.V_xX  # (A,D)
        Y = Yc + self.mean_X[None, :]
        return Y

    def _solve_Ainv_vecs(self, X_dA: np.ndarray) -> np.ndarray:
        """
        Solve (R^T R + sigma2 I)^{-1} X for X shape (d,A).
        Only available for chol/eigh solvers. For CG solver we can do per-column CG.
        """
        X = np.asarray(X_dA, dtype=np.float64)

        if self.solver == "chol":
            assert self._chol is not None
            return la.cho_solve(self._chol, X, check_finite=False)

        if self.solver == "eigh":
            assert self._eig is not None
            w, V = self._eig
            return V @ ((V.T @ X) / w[:, None])

        # cg fallback
        # (Ainv X) column-by-column with CG matvecs
        Rw = self.R_ix_w
        lam_eff = self.sigma2 + self.jitter
        d = self.d_lat

        def matvec(v: np.ndarray) -> np.ndarray:
            return (Rw.T @ (Rw @ v)) + lam_eff * v

        Aop = sla.LinearOperator((d, d), matvec=matvec, dtype=np.float64)

        out = np.empty_like(X)
        for j in range(X.shape[1]):
            sol, info = sla.cg(
                Aop,
                X[:, j],
                x0=np.zeros((d,), dtype=np.float64),
                maxiter=self.cg_maxiter,
                rtol=self.cg_rtol,
                atol=self.cg_atol,
            )
            if info != 0:
                raise RuntimeError(f"LGP variance CG solve failed with info={info}.")
            out[:, j] = sol
        return out

    def _decode_mean_var(self, Ra: np.ndarray, *, include_obs_noise: bool) -> Tuple[np.ndarray, np.ndarray]:
        """
        Mean + scalar variance proxy:
          var_f[a] = sigma2 * r_a^T A^{-1} r_a
        """
        assert self.V_xX is not None

        Raw = self._pp.apply(Ra)  # (A,d)

        # mean
        Yc = Raw @ self.V_xX
        mean = Yc + self.mean_X[None, :]

        # var_f: sigma2 * diag(Raw A^{-1} Raw^T)
        # Compute A^{-1} Raw^T (d,A), then diag = sum(Raw^T * AinvRawT, axis=0)
        Ainv_Rt = self._solve_Ainv_vecs(Raw.T)  # (d,A)
        quad = np.sum(Raw.T * Ainv_Rt, axis=0)  # (A,)
        var_f = self.sigma2 * quad

        if include_obs_noise:
            var_f = var_f + self.sigma2

        return mean, var_f

    # ------------------------------------------------------------
    # Diagnostics
    # ------------------------------------------------------------
    def kernel_mass(self, R_ax: Union[np.ndarray, list]) -> np.ndarray:
        """
        For linear kernel, a simple "mass" proxy is ||r||^2.
        """
        R_ax = np.asarray(R_ax, dtype=np.float64)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]
        Raw = self._pp.apply(R_ax)
        mass = np.sum(Raw * Raw, axis=1)
        return mass[0] if single else mass


__all__ = ["LGP"]