# GPLM.py
# ============================================================
# GPLM (drop-in replacement):
#   Sparse-kernel KRR / mean-GP decoder on latent coordinates.
#
# This version:
#   - Builds a sparse kernel matrix K using kNN (ANN) on training latents.
#   - Solves (K + sigma2 * I) S = Y  for S (N,D) via iterative sparse solvers.
#   - Predicts for new points using only k_pred nearest training points:
#       y(x) ≈ sum_{j in kNN(x)} k(x, x_j) * S_j
#
# Key properties:
#   - No Nyström anchors / no inducing points.
#   - "Nonlinearity" corresponds to strong locality (small eps, small k_graph):
#       sparse + high-rank-ish operator, but scalable because it's sparse.
#
# API compatibility:
#   - class GPLM
#   - __init__(...), fit(...)
#   - __call__(R_ax), predict(..., return_var=True)
#   - kernel_mass(...)
#   - flow(...) present but NotImplemented (optional advanced geometry)
#
# Dependencies:
#   numpy, scipy
#   ann.py + utils.py (same repo assumptions as your previous GPLM)
# ============================================================

from __future__ import annotations

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

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla

# --- same repo dependencies as before ---
from ann import ANNBackend, make_ann
from utils import median_eps_from_knn_d2

InducingMode = Literal["random_subset", "fps", "kmeans_medoids", "given"]  # kept for API compat
SolveMethod = Literal["cg", "minres", "auto"]


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


def _row_norm2(X: np.ndarray) -> np.ndarray:
    return np.sum(X * X, axis=1)


def _rbf_weights_from_d2(D2: np.ndarray, beta: float, eps: float) -> np.ndarray:
    # weight = exp(-beta * d^2 / eps)
    return np.exp(-float(beta) * (D2.astype(np.float64) / float(eps)))


def _symmetrize_coo(i: np.ndarray, j: np.ndarray, v: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Add transpose entries and return concatenated arrays."""
    i2 = np.concatenate([i, j], axis=0)
    j2 = np.concatenate([j, i], axis=0)
    v2 = np.concatenate([v, v], axis=0)
    return i2, j2, v2


def _unique_coo_sum(N: int, i: np.ndarray, j: np.ndarray, v: np.ndarray) -> sp.csr_matrix:
    """Build CSR matrix with duplicates summed."""
    K = sp.coo_matrix((v, (i, j)), shape=(N, N), dtype=np.float64).tocsr()
    K.sum_duplicates()
    return K


def _solve_multi_rhs(
    A: sp.csr_matrix,
    Y: np.ndarray,
    *,
    method: SolveMethod = "auto",
    tol: float = 1e-6,
    maxiter: int = 500,
    verbose: bool = False,
) -> np.ndarray:
    """
    Solve A X = Y for X with multiple RHS columns using iterative solvers.
    A is expected sparse and (typically) symmetric.
    """
    Y = np.asarray(Y, dtype=np.float64, order="C")
    N, D = Y.shape
    X = np.zeros((N, D), dtype=np.float64)

    # Choose solver
    if method == "auto":
        # CG is fastest if SPD; minres is safer if indefinite
        method_use: SolveMethod = "cg"
    else:
        method_use = method

    # Wrapper per RHS
    for d in range(D):
        b = Y[:, d]
        x0 = None  # could warm-start if you want

        if method_use == "cg":
            x, info = spla.cg(A, b, x0=x0, rtol=tol, maxiter=maxiter)
            if info != 0:
                # fallback to MINRES
                x, info2 = spla.minres(A, b, x0=x0, rtol=tol, maxiter=maxiter)
                if verbose:
                    print(f"[GPLM sparse] cg failed info={info}, minres info={info2} on dim {d}")
            X[:, d] = x
        elif method_use == "minres":
            x, info = spla.minres(A, b, x0=x0, rtol=tol, maxiter=maxiter)
            if verbose and info != 0:
                print(f"[GPLM sparse] minres info={info} on dim {d}")
            X[:, d] = x
        else:
            raise ValueError(f"Unknown solve method: {method_use}")

    return X


@dataclass
class _SparseKernelConfig:
    k_graph: int = 64          # neighbors per training point for building sparse K
    mutual: bool = True        # mutual-kNN symmetrization (recommended)
    include_self: bool = True  # ensure diagonal has 1.0
    normalize_rows: bool = False  # optional row-normalization for stability


class GPLM:
    """
    GPLM (Sparse-kernel KRR decoder)

    Training:
      Inputs:  R_ix (N,d) latents, R_iX (N,D) outputs
      Build sparse kernel K via kNN on R_ix:
        K_ij = exp(-beta ||R_i-R_j||^2 / eps)  for neighbors only

      Solve for weights S (N,D):
        (K + sigma2 * I + jitter*I) S = Y_centered

    Inference:
      For query R_ax:
        Find k_pred nearest training points j_aK
        Compute weights w_aK = exp(-beta d2/eps)
        Predict:
          Yc = sum_k w[a,k] * S[j_aK[a,k], :]
        Return Y = Yc + mean_X

    Variance proxy:
      Not full GP variance; return a support-based scalar:
        mass = sum_k w[a,k]
        var ≈ sigma2 / (mass + 1e-12)
    """

    def __init__(
        self,
        R_ix: np.ndarray,
        R_iX: np.ndarray,
        *,
        # Kernel params
        beta: float = 1.0,
        eps: Optional[float] = None,
        k_eps: int = 256,
        eps_use_kth: bool = True,
        eps_mul: float = 1.0,
        # Regularization (acts like ridge lambda)
        sigma2: float = 1e-5,
        jitter: float = 1e-8,
        # "Inducing" params kept for API compat (ignored)
        m: int = 1024,
        inducing: InducingMode = "kmeans_medoids",
        Z_mx: Optional[np.ndarray] = None,
        seed: int = 0,
        # Preprocess
        center_X: bool = True,
        whiten_latent: bool = False,
        dtype: Any = np.float32,
        # Sparse kernel build
        k_graph: int = 64,
        mutual_knn: bool = True,
        include_self: bool = True,
        normalize_rows: bool = False,
        # Solve
        solve_method: SolveMethod = "auto",
        solve_tol: float = 1e-6,
        solve_maxiter: int = 800,
        solve_verbose: bool = False,
        # Inference neighbor truncation
        pred_k: Optional[int] = 128,
        ann_backend: ANNBackend = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        # accept unicode kwargs (β, ε, κ_eps, σ2, pred_κ, ...)
        **kwargs: Any,
    ):
        # ---- map unicode kwargs -> ascii ----
        if "β" in kwargs:
            beta = kwargs.pop("β")
        if "ε" in kwargs:
            eps = kwargs.pop("ε")
        if "κ_eps" in kwargs:
            k_eps = kwargs.pop("κ_eps")
        if "ε_use_kth" in kwargs:
            eps_use_kth = kwargs.pop("ε_use_kth")
        if "ε_mul" in kwargs:
            eps_mul = kwargs.pop("ε_mul")
        if "σ2" in kwargs:
            sigma2 = kwargs.pop("σ2")
        if "pred_κ" in kwargs:
            pred_k = kwargs.pop("pred_κ")

        # ignore anchor arguments quietly (compat)
        _ = (m, inducing, Z_mx)

        if kwargs:
            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")

        self.beta = float(beta)
        self.β = self.beta

        self.sigma2 = float(sigma2)   # ridge lambda in (K + sigma2 I)
        self.σ2 = self.sigma2

        self.jitter = float(jitter)
        self.seed = int(seed)
        self.dtype = dtype

        # ---- validate / cast ----
        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

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

        # ---- latent whitening (optional) ----
        self.whiten_latent = bool(whiten_latent)
        Ztrain = R_ix.astype(np.float64)
        if self.whiten_latent:
            self.lat_mean_x = Ztrain.mean(axis=0)
            self.lat_std_x = np.maximum(Ztrain.std(axis=0), 1e-12)
            Ztrain_w = (Ztrain - self.lat_mean_x) / self.lat_std_x
        else:
            self.lat_mean_x = np.zeros((self.d_lat,), dtype=np.float64)
            self.lat_std_x = np.ones((self.d_lat,), dtype=np.float64)
            Ztrain_w = Ztrain

        self.R_ix_w = np.ascontiguousarray(Ztrain_w.astype(np.float64, copy=False))  # (N,d) float64

        # ---- ANN on training latents ----
        self.ann_train, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
        self.ann_train.build(self.R_ix_w.astype(self.dtype, copy=False))

        # ---- eps via kNN distances ----
        if eps is None:
            k_eps_eff = int(min(max(8, int(k_eps)), self.N - 1))
            # ask for k_eps+1 to try include self
            j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), k_eps_eff + 1)

            i = np.arange(self.N)[:, None]
            is_self = (j_iK1 == i)

            if np.any(is_self):
                D2_iK = np.empty((self.N, k_eps_eff), dtype=np.float64)
                for ii in range(self.N):
                    keep = (j_iK1[ii] != ii)
                    D2_iK[ii] = D2_iK1[ii][keep][:k_eps_eff]
            else:
                D2_iK = D2_iK1[:, :k_eps_eff].astype(np.float64, copy=False)

            eps_hat = median_eps_from_knn_d2(D2_iK, use_kth=bool(eps_use_kth))
        else:
            eps_hat = float(eps)

        eps_hat *= float(eps_mul)
        if eps_hat <= 0:
            raise ValueError("eps must be > 0.")
        self.eps = float(eps_hat)
        self.ε = self.eps

        # ---- build sparse kernel K ----
        cfg = _SparseKernelConfig(
            k_graph=int(min(max(4, int(k_graph)), self.N - 1)),
            mutual=bool(mutual_knn),
            include_self=bool(include_self),
            normalize_rows=bool(normalize_rows),
        )
        self._cfg = cfg

        # query kNN on training set for graph edges
        j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), cfg.k_graph + 1)

        # drop self if present
        rows = []
        cols = []
        vals = []

        for i in range(self.N):
            nbrs = j_iK1[i]
            d2 = D2_iK1[i].astype(np.float64, copy=False)
            # remove self
            mask = (nbrs != i)
            nbrs = nbrs[mask][: cfg.k_graph]
            d2 = d2[mask][: cfg.k_graph]

            w = _rbf_weights_from_d2(d2, beta=self.beta, eps=self.eps)

            rows.append(np.full(nbrs.shape[0], i, dtype=np.int64))
            cols.append(nbrs.astype(np.int64, copy=False))
            vals.append(w.astype(np.float64, copy=False))

        i_idx = np.concatenate(rows, axis=0)
        j_idx = np.concatenate(cols, axis=0)
        v_idx = np.concatenate(vals, axis=0)

        # symmetric adjacency
        if cfg.mutual:
            i_idx, j_idx, v_idx = _symmetrize_coo(i_idx, j_idx, v_idx)

        # build sparse K (CSR)
        K = _unique_coo_sum(self.N, i_idx, j_idx, v_idx)

        # set diagonal to 1 (kernel self-sim), improves conditioning
        if cfg.include_self:
            K = K.tolil(copy=False)
            diag = K.diagonal()
            # if diagonal already has values from sym edges, top it up to 1
            diag_new = np.maximum(np.asarray(diag).reshape(-1), 1.0)
            for ii in range(self.N):
                K[ii, ii] = float(diag_new[ii])
            K = K.tocsr(copy=False)

        # optional row normalization (turns kernel into a diffusion-like operator)
        if cfg.normalize_rows:
            rs = np.asarray(K.sum(axis=1)).reshape(-1)
            rs = np.maximum(rs, 1e-12)
            inv = 1.0 / rs
            K = sp.diags(inv, format="csr") @ K

        K.sum_duplicates()
        self.K = K  # (N,N) sparse

        # ---- form A = K + (sigma2 + jitter) I ----
        lam = float(self.sigma2)
        jit = float(self.jitter)
        A = self.K.tocsr(copy=True)
        A = A + sp.diags((lam + jit) * np.ones(self.N), format="csr")
        self._A = A.tocsr(copy=False)

        # ---- solve for S: A S = Y ----
        self.solve_method = str(solve_method)
        self.solve_tol = float(solve_tol)
        self.solve_maxiter = int(solve_maxiter)
        self.solve_verbose = bool(solve_verbose)

        self.S_iX = _solve_multi_rhs(
            self._A,
            Y,
            method=self.solve_method,  # type: ignore
            tol=self.solve_tol,
            maxiter=self.solve_maxiter,
            verbose=self.solve_verbose,
        ).astype(np.float64)

        # store float32 copy for fast inference
        self.S_iX_f32 = self.S_iX.astype(np.float32, copy=False)

        # ---- inference neighbor count ----
        if pred_k is None:
            self.pred_k = int(min(128, self.N - 1))
        else:
            self.pred_k = int(min(max(1, int(pred_k)), self.N - 1))
        self.pred_κ = self.pred_k  # unicode alias

    # ------------------------------------------------------------
    # Convenience alternate constructor
    # ------------------------------------------------------------
    @classmethod
    def fit(cls, R_ix: np.ndarray, R_iX: np.ndarray, **kwargs: Any) -> "GPLM":
        return cls(R_ix, R_iX, **kwargs)

    # ------------------------------------------------------------
    # Internal prediction helpers
    # ------------------------------------------------------------
    def _whiten_query(self, R_ax: np.ndarray) -> np.ndarray:
        R_ax = np.asarray(R_ax, dtype=np.float64)
        if self.whiten_latent:
            return (R_ax - self.lat_mean_x[None, :]) / self.lat_std_x[None, :]
        return R_ax

    def _predict_mean(self, R_ax: np.ndarray) -> np.ndarray:
        """
        Mean prediction using only pred_k nearest training points.
        """
        R_ax = np.asarray(R_ax)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]
        R_ax = np.ascontiguousarray(R_ax.astype(self.dtype, copy=False))

        # whiten query for ANN
        Rw = self._whiten_query(R_ax.astype(np.float64, copy=False)).astype(self.dtype, copy=False)

        # neighbors in training set
        j_aK, D2_aK = self.ann_train.search(Rw, self.pred_k)

        # weights (A,k)
        W = _rbf_weights_from_d2(D2_aK.astype(np.float64, copy=False), beta=self.beta, eps=self.eps)

        # gather S for neighbors -> (A,k,D)
        S = self.S_iX_f32  # (N,D)
        Sj = S[j_aK]       # (A,k,D)

        # weighted sum -> (A,D)
        Yc = np.einsum("ak,akd->ad", W.astype(np.float32, copy=False), Sj, optimize=True).astype(np.float64)

        # add mean
        Y = Yc + self.mean_X[None, :]
        return Y[0] if single else Y

    def _predict_mean_var(self, R_ax: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Mean + cheap scalar uncertainty proxy based on kernel mass on neighbors.
        This is *not* exact GP posterior variance.
        """
        R_ax = np.asarray(R_ax)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]
        R_ax = np.ascontiguousarray(R_ax.astype(self.dtype, copy=False))

        Rw = self._whiten_query(R_ax.astype(np.float64, copy=False)).astype(self.dtype, copy=False)
        j_aK, D2_aK = self.ann_train.search(Rw, self.pred_k)

        W = _rbf_weights_from_d2(D2_aK.astype(np.float64, copy=False), beta=self.beta, eps=self.eps)  # (A,k)
        mass = np.sum(W, axis=1)  # (A,)

        S = self.S_iX_f32
        Sj = S[j_aK]  # (A,k,D)
        Yc = np.einsum("ak,akd->ad", W.astype(np.float32, copy=False), Sj, optimize=True).astype(np.float64)
        mean = Yc + self.mean_X[None, :]

        # heuristic: low mass => off-support => higher uncertainty
        var = (self.sigma2 / (mass + 1e-12)).astype(np.float64)

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

    # ------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------
    def __call__(self, R_ax: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
        """
        Mean prediction only. For uncertainty, use predict(..., return_var=True).
        """
        R_ax = np.asarray(R_ax)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]

        if batch_size is None:
            Y = self._predict_mean(R_ax)
        else:
            bs = int(batch_size)
            out = []
            for s in range(0, R_ax.shape[0], bs):
                out.append(self._predict_mean(R_ax[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,
    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
        """
        Predict mean and (optional) scalar variance proxy per query.
        """
        R_ax = np.asarray(R_ax)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]

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

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

        if single:
            return mean, float(var[0]) if np.ndim(var) > 0 else float(var)
        return mean, var

    def kernel_mass(self, R_ax: Union[np.ndarray, list]) -> np.ndarray:
        """
        Support diagnostic:
          mass(x) = sum_{j in kNN(x)} exp(-beta ||x-x_j||^2 / eps)
        """
        R_ax = np.asarray(R_ax, dtype=np.float64)
        single = (R_ax.ndim == 1)
        if single:
            R_ax = R_ax[None, :]

        Rw = self._whiten_query(R_ax).astype(self.dtype, copy=False)
        _, D2_aK = self.ann_train.search(Rw, self.pred_k)
        W = _rbf_weights_from_d2(D2_aK.astype(np.float64, copy=False), beta=self.beta, eps=self.eps)
        mass = np.sum(W, axis=1)
        return float(mass[0]) if single else mass

    # ------------------------------------------------------------
    # Geometry / flow (kept for API compatibility)
    # ------------------------------------------------------------
    def flow(
        self,
        R_ax: Union[np.ndarray, list],
        v_ax: Union[np.ndarray, list],
        *,
        dt: float = 1e-2,
        K_p: int = 5,
        K_q: int = 5,
        D_block: int = 8192,
        lam: float = 1e-8,
        metric_solver: str = "auto",
        eig_clip: float = 1e-12,
        k_tangent: Optional[int] = None,
        force_fn: Optional[Any] = None,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Placeholder for compatibility with the original GPLM API.

        Sparse-kernel KRR decoder has a well-defined Jacobian via kernel gradients,
        but a robust geodesic integrator is non-trivial and model-specific.

        If you truly need flow() (geodesic-ish latent motion), you can:
          - implement it using local neighbor gradients, OR
          - keep the original GPLM for geometry tasks.

        For now: not implemented.
        """
        raise NotImplementedError(
            "flow() is not implemented in sparse-kernel GPLM. "
            "Use the original Nyström GPLM for geodesic flow, or implement "
            "a local-neighborhood gradient-based flow."
        )


__all__ = ["GPLM"]