# examples/s4.py
"""Deterministic, standard-library S4 arithmetic reference.

Run: python3 examples/s4.py
This is a dense teaching reference, not a trainable S4 implementation, optimized
Cauchy backend, checkpoint runner, or benchmark reproduction. Regions are the
canonical source for the review's annotated code panels.
"""

import cmath
import json
import math


def dot(left, right):
    """Return a bilinear dot product (no implicit complex conjugation)."""
    return sum(x * y for x, y in zip(left, right))


def transpose(matrix):
    """Transpose a nonempty rectangular matrix."""
    return list(map(list, zip(*matrix)))


def matvec(matrix, vector):
    """Apply a dense matrix to a vector."""
    return [dot(row, vector) for row in matrix]


def solve(matrix, rhs):
    """Solve multiple right-hand sides using partial-pivot elimination."""
    size = len(matrix)
    rows = [list(row) + list(right) for row, right in zip(matrix, rhs)]
    for col in range(size):
        pivot = max(range(col, size), key=lambda i: abs(rows[i][col]))
        if abs(rows[pivot][col]) < 1e-13:
            raise ValueError("Singular system")
        rows[col], rows[pivot] = rows[pivot], rows[col]
        scale = rows[col][col]
        rows[col] = [value / scale for value in rows[col]]
        for i in range(size):
            if i == col:
                continue
            factor = rows[i][col]
            rows[i] = [x - factor * y for x, y in zip(rows[i], rows[col])]
    return [row[size:] for row in rows]


# region: hippo
def hippo(size):
    """Construct the real HiPPO-LegS matrix of S4 equation (2)."""
    if not isinstance(size, int) or not 1 <= size <= 32:
        raise ValueError("State size must be 1 through 32")
    return [
        [
            0.0 if i < j else -(i + 1.0) if i == j
            else -math.sqrt((2 * i + 1) * (2 * j + 1))
            for j in range(size)
        ]
        for i in range(size)
    ]
# endregion: hippo


# region: discretize
def discretize(matrix, b, delta):
    """Bilinear discretization; solve instead of forming an inverse."""
    if not math.isfinite(delta) or delta <= 0:
        raise ValueError("Positive finite step required")
    size = len(matrix)
    minus = [[float(i == j) - delta * matrix[i][j] / 2
              for j in range(size)] for i in range(size)]
    plus = [[float(i == j) + delta * matrix[i][j] / 2
             for j in range(size)] + [delta * b[i]] for i in range(size)]
    result = solve(minus, plus)
    return [row[:size] for row in result], [row[-1] for row in result]
# endregion: discretize


# region: recurrence
def recurrence(matrix, b, c, inputs, initial=None, skip=0.0):
    """Update before reading; the initial state is x[-1]."""
    state = [0.0] * len(b) if initial is None else list(initial)
    output = []
    for value in inputs:
        state = [x + injection * value
                 for x, injection in zip(matvec(matrix, state), b)]
        output.append(dot(c, state) + skip * value)
    return output


def impulse_kernel(matrix, b, c, length):
    """Return K[j] = C Abar**j Bbar, starting with K[0] = C Bbar."""
    state = list(b)
    kernel = []
    for _ in range(length):
        kernel.append(dot(c, state))
        state = matvec(matrix, state)
    return kernel


def convolve(kernel, inputs):
    """Direct causal convolution; no circular wraparound."""
    return [sum(kernel[t - j] * inputs[j] for j in range(t + 1))
            for t in range(len(inputs))]
# endregion: recurrence


# region: structure
def corrected_hippo(size):
    """A + pp^T = -I/2 + skew-symmetric, hence normal."""
    matrix = hippo(size)
    p = [math.sqrt(i + 0.5) for i in range(size)]
    normal = [[matrix[i][j] + p[i] * p[j]
               for j in range(size)] for i in range(size)]
    return normal, p


def resolvent(lambdas, p, q, b, c, s):
    """Return C(sI-A)^-1 B for the real fixture A=diag(lambda)-pq^T.

    Four Cauchy sums and one rank-one Woodbury correction. General complex
    DPLR uses q conjugates; here q is real, so conjugation is identity.
    """
    r = [1 / (s - value) for value in lambdas]

    def pair(left, right):
        return sum(x * y * pole for x, y, pole in zip(left, right, r))

    return pair(c, b) - pair(c, p) * pair(q, b) / (1 + pair(q, p))
# endregion: structure


def dft(values, inverse=False):
    """Naive DFT, used to verify identities, not FFT complexity."""
    size = len(values)
    sign = 1 if inverse else -1
    return [sum(value * cmath.exp(sign * 2j * math.pi * k * j / size)
                for j, value in enumerate(values)) / (size if inverse else 1)
            for k in range(size)]


# region: kernel
def kernel_spectrum(length=17, delta=0.3, truncate=True):
    """Verify Algorithm 1 on an authored real DPLR system.

    Odd length avoids z=-1, where the bilinear change of variable needs its
    removable-singularity limit. Dense preprocessing is intentionally slow.
    """
    if not isinstance(length, int) or length < 3 or length % 2 == 0:
        raise ValueError("Odd length >= 3 required")
    lambdas, p = [-0.5, -2.0, -4.0], [0.4, 0.3, 0.2]
    b, c = [1.0, 0.5, 0.25], [1.0, -0.6, 0.2]
    matrix = [[(value if i == j else 0) - p[i] * p[j]
               for j in range(3)] for i, value in enumerate(lambdas)]
    abar, bbar = discretize(matrix, b, delta)
    c_power = list(c)
    for _ in range(length):
        c_power = matvec(transpose(abar), c_power)
    c_tilde = [x - (tail if truncate else 0) for x, tail in zip(c, c_power)]
    spectrum = []
    for k in range(length):
        z = cmath.exp(-2j * math.pi * k / length)
        s = 2 * (1 - z) / (delta * (1 + z))
        spectrum.append(2 / (1 + z) * resolvent(
            lambdas, p, p, b, c_tilde, s))
    expected = impulse_kernel(abar, bbar, c, length)
    recovered = [value.real for value in dft(spectrum, inverse=True)]
    error = max(abs(x - y) for x, y in zip(expected, recovered))
    return expected, spectrum, recovered, error
# endregion: kernel


def main():
    """Check equivalence, normality, stability, truncation and causality."""
    matrix = hippo(4)
    b = [math.sqrt(2 * i + 1) for i in range(4)]
    c = [1.0, -0.3, 0.2, 0.1]
    abar, bbar = discretize(matrix, b, 0.2)
    inputs = [math.sin(i * 0.7) for i in range(20)]
    kernel = impulse_kernel(abar, bbar, c, len(inputs))
    recurrent = recurrence(abar, bbar, c, inputs)
    convolution = convolve(kernel, inputs)
    error = max(abs(x - y) for x, y in zip(recurrent, convolution))
    assert error < 1e-12
    normal, _ = corrected_hippo(8)
    assert max(abs(normal[i][j] + normal[j][i] + float(i == j))
               for i in range(8) for j in range(8)) < 1e-12
    kernel_error = kernel_spectrum()[-1]
    assert kernel_error < 1e-12
    assert kernel_spectrum(truncate=False)[-1] > 1e-3
    assert recurrence([[0.5]], [1], [1], [1, 0, 0]) == [1, 0.5, 0.25]
    assert recurrence([[0.5]], [1], [1], [0], initial=[2]) == [1]
    for delta in (0.01, 0.3, 2, 10):
        stable, _ = discretize([[-2]], [1], delta)
        assert abs(stable[0][0]) < 1
    print(json.dumps({"recurrence_convolution_error": error,
                      "finite_kernel_error": kernel_error,
                      "evidence": "computed teaching arithmetic; no training"}, indent=2))


if __name__ == "__main__":
    main()