# ALSA.py
# ============================================================
# ALSA: Analog-Laplacian Spectrum Analysis (nonparametric baseline)
#
# Purely nonparametric analog forecaster built ONLY from a contextual
# time series prefix R_{aX} (shape (ell, D)).
#
# Core idea (Takens + kernel analog forecasting):
#   - Hankelize prefix into L-delay windows W_A = R[A:A+L]  (A=0..K-1)
#   - For each window with an observed next sample, define target y_A = R[A+L]
#     (so context/library size is C = ell - L)
#   - Given a query window (initially the last L points of the prefix),
#     predict the next sample by a row-normalized kernel average:
#
#       k(q,A) = exp( -beta * ||w_q - w_A||^2 / eps )
#       w_A    = k(q,A) / sum_A k(q,A)
#       y_hat  = sum_A w_A * y_A
#
#   - Rollout autoregressively for 'steps' by appending predictions.
#
# Notes
# -----
# - This is "NLSA-like" in the sense of Laplacian-kernel / Markov-row
#   analog forecasting on Takens windows, but WITHOUT a trained global
#   parametric decoder and WITHOUT computing diffusion eigenfunctions.
# - You can optionally sparsify with top_k nearest windows per step.
# - Complexity per forecast step is O(C * (L*D)) due to a single dense
#   dot-product against the library. For large C, use top_k.
#
# ============================================================

from __future__ import annotations

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

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


# -----------------------------
# helpers
# -----------------------------
def _as_2d(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    if X.ndim == 1:
        X = X[:, None]
    if X.ndim != 2:
        raise ValueError("Expected 1D or 2D array.")
    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.
    """
    F = _as_2d(F_tD)
    L = int(L)
    if L <= 0:
        raise ValueError("L must be positive.")
    if F.shape[0] < L:
        raise ValueError(f"Need len(F) >= L. Got {F.shape[0]} < {L}.")

    W = sliding_window_view(F, window_shape=L, axis=0)  # (K, L, D) or (K, D, L)
    # Normalize to (K, L, D)
    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_KLD: np.ndarray) -> np.ndarray:
    """
    Flatten windows (K, L, D) -> (K, L*D).
    """
    W = np.asarray(W_KLD, dtype=np.float64)
    if W.ndim != 3:
        raise ValueError("Expected windows with shape (K, L, D).")
    return W.reshape(W.shape[0], -1)


def _estimate_eps_from_library(
    X_lib: np.ndarray,
    *,
    sample_pairs: int = 4096,
    rng: Optional[np.random.Generator] = None,
    eps_min: float = 1e-12,
) -> float:
    """
    Heuristic bandwidth: median of squared distances over random pairs.
    Works well enough for analog forecasting, avoids O(C^2).
    """
    X = np.asarray(X_lib, dtype=np.float64)
    n = X.shape[0]
    if n <= 1:
        return 1.0

    if rng is None:
        rng = np.random.default_rng(0)

    m = int(min(max(64, sample_pairs), n * (n - 1) // 2))
    i = rng.integers(0, n, size=m)
    j = rng.integers(0, n, size=m)
    neq = (i != j)
    if not np.any(neq):
        j = (i + 1) % n
        neq = np.ones_like(i, dtype=bool)
    i = i[neq]
    j = j[neq]

    diff = X[i] - X[j]
    d2 = np.einsum("ij,ij->i", diff, diff, optimize=True)
    med = float(np.median(d2)) if d2.size else 1.0
    return float(max(med, eps_min))


def _kernel_weights_rbf(
    X_lib: np.ndarray,
    x_q: np.ndarray,
    lib_norm2: np.ndarray,
    *,
    beta: float,
    eps: float,
    top_k: Optional[int],
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute row-normalized RBF weights over library windows.

    Returns:
      idx: indices used (full or top_k subset)
      w:   weights over idx (sum to 1)
    """
    X_lib = np.asarray(X_lib, dtype=np.float64)
    x_q = np.asarray(x_q, dtype=np.float64).reshape(-1)
    lib_norm2 = np.asarray(lib_norm2, dtype=np.float64).reshape(-1)

    q_norm2 = float(np.dot(x_q, x_q))
    # d2 = ||x_i||^2 + ||x_q||^2 - 2 <x_i, x_q>
    dots = X_lib @ x_q
    d2 = lib_norm2 + q_norm2 - 2.0 * dots
    np.maximum(d2, 0.0, out=d2)

    if top_k is not None:
        k = int(top_k)
        k = max(1, min(k, d2.size))
        # indices of k smallest distances
        idx = np.argpartition(d2, kth=k - 1)[:k]
        d2_use = d2[idx]
    else:
        idx = np.arange(d2.size)
        d2_use = d2

    beta = float(beta)
    eps = float(eps)
    if eps <= 0:
        raise ValueError("eps must be > 0.")

    # k = exp(-beta * d2 / eps)
    # subtract min exponent for numerical stability if needed
    ex = -beta * (d2_use / eps)
    ex = ex - float(np.max(ex))  # stabilize
    kvec = np.exp(ex)

    s = float(np.sum(kvec))
    if not np.isfinite(s) or s <= 0:
        # fallback: nearest neighbor
        j = int(idx[np.argmin(d2_use)])
        return np.array([j], dtype=np.int64), np.array([1.0], dtype=np.float64)

    w = (kvec / s).astype(np.float64, copy=False)
    return idx.astype(np.int64, copy=False), w


@dataclass
class ALSAInfo:
    L: int
    D: int
    ell: int
    C: int
    eps: float
    beta: float
    top_k: Optional[int]
    # Optional per-step diagnostics (can be big; disabled by default)
    weights: Optional[list[Tuple[np.ndarray, np.ndarray]]] = None


# ============================================================
# Public API
# ============================================================
def alsa_forecast(
    prefix: np.ndarray,
    *,
    steps: int,
    L: int,
    beta: float = 1.0,
    eps: Optional[float] = None,
    eps_mul: float = 1.0,
    top_k: Optional[int] = None,
    return_info: bool = False,
    store_weights: bool = False,
    seed: int = 0,
) -> Tuple[np.ndarray, ALSAInfo] | np.ndarray:
    """
    ALSA nonparametric autoregressive forecast from a prefix.

    Parameters
    ----------
    prefix : array (ell, D)
        Context time series (observed).
    steps : int
        Forecast horizon H.
    L : int
        Takens window length. Requires ell >= L+1 for at least one context pair.
    beta : float
        Kernel sharpness (larger -> more local / NN-like).
    eps : float or None
        Bandwidth scale in the kernel. If None, estimated from library windows.
    eps_mul : float
        Multiplier applied to estimated eps (or provided eps) to tune locality.
    top_k : int or None
        If set, restrict weights to the top_k nearest library windows (sparse analog).
        This can massively speed up forecasting for large contexts.
    return_info : bool
        Return ALSAInfo diagnostics.
    store_weights : bool
        If True, store (indices, weights) per forecast step in ALSAInfo.
    seed : int
        RNG seed used only for eps estimation when eps is None.

    Returns
    -------
    preds : (steps, D) float64
    info  : ALSAInfo (if return_info=True)

    Notes
    -----
    Library construction:
      - windows W_A = prefix[A:A+L] for A=0..K-1, K=ell-L+1
      - context/library indices A=0..C-1 where C=ell-L
      - targets y_A = prefix[A+L]
      - query window starts at prefix[ell-L:ell]
    """
    prefix = _as_2d(prefix)
    ell, D = prefix.shape
    L = int(L)
    H = int(steps)
    if H <= 0:
        out = np.zeros((0, D), dtype=np.float64)
        if return_info:
            info = ALSAInfo(L=L, D=D, ell=ell, C=max(0, ell - L), eps=float(eps or 1.0), beta=float(beta), top_k=top_k)
            return out, info
        return out

    if ell < L + 1:
        raise ValueError(f"Need prefix length ell >= L+1 for ALSA. Got ell={ell}, L={L}.")

    # Build Hankel windows from prefix
    W_all = _sliding_windows(prefix, L)  # (K, L, D)
    K = W_all.shape[0]
    C = ell - L  # number of context pairs with observed next sample
    # Library windows are W_all[0:C], targets are prefix[L:L+C]
    W_lib = np.ascontiguousarray(W_all[:C])
    Y_lib = np.ascontiguousarray(prefix[L : L + C])

    X_lib = _flatten_windows(W_lib)  # (C, L*D)
    lib_norm2 = np.sum(X_lib * X_lib, axis=1)

    rng = np.random.default_rng(int(seed))
    if eps is None:
        eps0 = _estimate_eps_from_library(X_lib, rng=rng)
    else:
        eps0 = float(eps)
    eps_used = float(max(1e-12, eps0 * float(eps_mul)))

    # Initial query window for rollout
    cur = prefix[-L:, :].copy()
    preds = np.zeros((H, D), dtype=np.float64)

    weights_log = [] if (return_info and store_weights) else None

    for h in range(H):
        x_q = cur.reshape(-1)  # (L*D,)
        idx, w = _kernel_weights_rbf(
            X_lib, x_q, lib_norm2,
            beta=float(beta),
            eps=eps_used,
            top_k=top_k,
        )
        y_hat = (w[:, None] * Y_lib[idx]).sum(axis=0)
        preds[h] = y_hat

        if weights_log is not None:
            # store a copy so it is stable if caller mutates
            weights_log.append((idx.copy(), w.copy()))

        # shift window and append prediction
        if L > 1:
            cur[:-1] = cur[1:]
        cur[-1] = y_hat

    if not return_info:
        return preds

    info = ALSAInfo(
        L=L,
        D=D,
        ell=ell,
        C=C,
        eps=eps_used,
        beta=float(beta),
        top_k=top_k,
        weights=weights_log,
    )
    return preds, info


__all__ = ["alsa_forecast", "ALSAInfo"]