import math
import numpy as np


def concentration_matrix(length: int, bandwidth: float) -> np.ndarray:
    """Finite discrete time-frequency concentration matrix."""
    if length < 1 or not 0.0 < bandwidth < 0.5:
        raise ValueError("length >= 1 and bandwidth in (0, .5) required")
    i = np.arange(length)
    d = i[:, None] - i[None, :]
    C = np.empty((length, length), dtype=np.float64)
    off = d != 0
    C[off] = np.sin(2.0 * np.pi * bandwidth * d[off]) / (np.pi * d[off])
    C[~off] = 2.0 * bandwidth
    return (C + C.T) * 0.5


def asymptotic_rank(length: int, bandwidth: float, delta: float) -> float:
    """Paper-inspired rank estimate with c=2*bandwidth*length."""
    if not 0.0 < delta < 0.5:
        raise ValueError("delta must be in (0, .5)")
    c = 2.0 * bandwidth * length
    log_odds = math.log((1.0 - delta) / delta)
    return c + log_odds / math.pi**2 * math.log(
        max(4.0 * math.pi**2 * c / log_odds, 1.000001))


def fourier_basis(length: int, rank: int) -> np.ndarray:
    """Real orthonormal low-frequency Fourier basis (cos/sin pairs)."""
    if not 1 <= rank <= length:
        raise ValueError("rank must be between 1 and length")
    n = np.arange(length)
    cols = [np.ones(length) / np.sqrt(length)]
    k = 1
    while len(cols) < rank:
        cols.append(np.sqrt(2.0 / length) * np.cos(2 * np.pi * k * n / length))
        if len(cols) < rank:
            cols.append(np.sqrt(2.0 / length) * np.sin(2 * np.pi * k * n / length))
        k += 1
    return np.stack(cols[:rank], axis=1)


class ProlateBottleneck:
    """Fixed orthogonal DPSS projection along axis 1 of [B,T,D] arrays."""
    def __init__(self, length: int, bandwidth: float, delta: float = 0.1,
                 rank: int | None = None):
        C = concentration_matrix(length, bandwidth)
        vals, vecs = np.linalg.eigh(C)
        order = np.argsort(vals)[::-1]
        self.eigenvalues = vals[order]
        self.full_basis = vecs[:, order]
        proposed = math.ceil(asymptotic_rank(length, bandwidth, delta))
        self.empirical_rank = int(np.sum(self.eigenvalues > delta))
        chosen = proposed if rank is None else rank
        self.rank = int(np.clip(chosen, 1, length))
        self.basis = self.full_basis[:, :self.rank]

    def encode(self, x: np.ndarray) -> np.ndarray:
        if x.ndim != 3 or x.shape[1] != self.basis.shape[0]:
            raise ValueError("x must have shape [batch, matching_length, channels]")
        return np.einsum("btd,tr->brd", x, self.basis)

    def decode(self, z: np.ndarray) -> np.ndarray:
        if z.ndim != 3 or z.shape[1] != self.rank:
            raise ValueError("z must have shape [batch, rank, channels]")
        return np.einsum("brd,tr->btd", z, self.basis)

    def reconstruct(self, x: np.ndarray) -> np.ndarray:
        return self.decode(self.encode(x))

    @property
    def attention_work_fraction(self) -> float:
        return (self.rank / self.basis.shape[0]) ** 2