"""Residual-screened Koopman latent bottleneck utilities.

Small NumPy implementation of the empirical residual criterion from the idea.
"""
import numpy as np


def fit_koopman(z):
    """Fit K in z[t+1] ~= K z[t], with z shaped [T, m]."""
    zm, zp = z[:-1].T, z[1:].T
    return zp @ np.linalg.pinv(zm)


def eig_residuals(K, z, eps=1e-12):
    """Return eigenvalues, right eigenvectors, and normalized residuals.

    This follows the displayed idea formula using u_i^* z_t. For the
    diagonal/normal systems used in the sanity check this is also the usual
    modal coordinate. Complex arithmetic is retained until final metrics.
    """
    vals, vecs = np.linalg.eig(K)
    residuals = []
    for i, lam in enumerate(vals):
        u = vecs[:, i]
        a0 = np.conjugate(u) @ z[:-1].T
        a1 = np.conjugate(u) @ z[1:].T
        residuals.append(np.sqrt(np.sum(np.abs(a1 - lam * a0)**2) /
                                 (np.sum(np.abs(a0)**2) + eps)))
    return vals, vecs, np.asarray(residuals).real


def screen_mask(residuals, quantile=0.75):
    tau = float(np.quantile(residuals, quantile))
    return residuals <= tau + 1e-12, tau


def spectral_forecast(K, z0, steps, mask=None):
    """Forecast with K's spectral modes, optionally zeroing screened modes."""
    vals, vecs = np.linalg.eig(K)
    inv = np.linalg.pinv(vecs)
    if mask is None:
        mask = np.ones(len(vals), dtype=bool)
    amps = inv @ z0
    amps = amps * mask
    out = []
    for h in range(1, steps + 1):
        out.append(np.real(vecs @ ((vals ** h) * amps)))
    return np.asarray(out)