# examples/s5.py
"""Deterministic S5 mechanics using only Python's standard library.

Run: python examples/s5.py
This is a causal linear SSM reference, not a training or checkpoint reproduction.
Source: arXiv:2208.04933v3, equations 5-6 and Appendix A.
Official code: lindermanlab/S5 at 3c18fdb6b06414da35e77b94b9cd855f6a95ef17.
"""

import cmath
import json
import math


# region: discretization
def discretize(poles, b_matrix, steps):
    """Return diagonal ZOH A and dense B for per-mode positive steps."""
    if not poles or len(poles) != len(b_matrix) or len(poles) != len(steps):
        raise ValueError("Expected one B row and interval per pole")
    if any(not math.isfinite(dt) or dt < 0 for dt in steps):
        raise ValueError("Intervals must be finite and nonnegative")
    a_bar, b_bar = [], []
    for pole, row, dt in zip(poles, b_matrix, steps):
        z = pole * dt
        a_bar.append(cmath.exp(z))
        if abs(z) < 1e-4:
            term, phi = 1.0, 1.0
            for n in range(1, 9):
                term *= z / (n + 1)
                phi += term
            gain = dt * phi
        else:
            gain = (cmath.exp(z) - 1) / pole
        b_bar.append([gain * value for value in row])
    return a_bar, b_bar
# endregion: discretization


# region: projection
def project(matrix, vector):
    """Multiply a dense matrix by a vector with explicit shape checks."""
    if any(len(row) != len(vector) for row in matrix):
        raise ValueError("Matrix width must equal vector length")
    return [sum(weight * value for weight, value in zip(row, vector))
            for row in matrix]


def elements(a_bar, b_bar, inputs):
    """Build (diagonal transition, dense input projection) at each time."""
    return [(list(a_bar), project(b_bar, u)) for u in inputs]
# endregion: projection


# region: scan
def compose(left, right):
    """Compose earlier left then later right; multiplication is elementwise."""
    a_left, b_left = left
    a_right, b_right = right
    return ([r * l for l, r in zip(a_left, a_right)],
            [r * l + b for r, l, b in zip(a_right, b_left, b_right)])


def prefix_scan(items):
    """Recursive work-efficient scan; Python executes it serially."""
    if len(items) <= 1:
        return list(items)
    paired = [compose(items[i], items[i + 1])
              for i in range(0, len(items) - 1, 2)]
    prefixes = prefix_scan(paired)
    return [prefixes[i // 2] if i % 2 else
            item if i == 0 else compose(prefixes[i // 2 - 1], item)
            for i, item in enumerate(items)]
# endregion: scan


# region: readout
def readout(states, inputs, c_matrix, direct, half_spectrum=False):
    """Read real outputs; factor two is only for paired half-spectrum storage."""
    factor = 2 if half_spectrum else 1
    outputs = []
    for state, u in zip(states, inputs):
        hidden = project(c_matrix, state)
        outputs.append([factor * value.real + d * source
                        for value, d, source in zip(hidden, direct, u)])
    return outputs
# endregion: readout


# region: initialization
def hippo_normal(size):
    """Construct HiPPO-N = HiPPO-LegS + p p^T, before diagonalization."""
    if size < 1:
        raise ValueError("State size must be positive")
    return [[-0.5 if n == k else
             (1 if n < k else -1) * math.sqrt((n + 0.5) * (k + 0.5))
             for k in range(size)] for n in range(size)]


def log_spaced_steps(count, low=0.001, high=0.1):
    """Deterministic log grid for teaching, not upstream random initialization."""
    if count < 1 or not 0 < low <= high:
        raise ValueError("Expected positive count and ordered positive bounds")
    return [math.exp(math.log(low) + i / max(count - 1, 1)
                     * math.log(high / low)) for i in range(count)]
# endregion: initialization


def main():
    """Check scan/loop equivalence, ZOH limits and conjugate reconstruction."""
    poles = [-0.5 + 0.8j, -0.5 - 0.8j]
    b_matrix = [[1.0, 0.25], [1.0, 0.25]]
    a_bar, b_bar = discretize(poles, b_matrix, [0.1, 0.1])
    inputs = [[1.0, -0.5], [0.0, 1.0], [0.5, 0.0], [-1.0, 0.5], [0.0, 0.0]]
    updates = elements(a_bar, b_bar, inputs)
    states = [b for _, b in prefix_scan(updates)]
    current = [0j, 0j]
    error = 0.0
    for (a, b), state in zip(updates, states):
        current = [ai * xi + bi for ai, xi, bi in zip(a, current, b)]
        error = max(error, max(abs(x - y) for x, y in zip(current, state)))
    assert error < 1e-12
    full = readout(states, inputs, [[0.5, 0.5], [0.25j, -0.25j]], [0.1, 0.0])
    half = readout([[x[0]] for x in states], inputs, [[0.5], [0.25j]],
                   [0.1, 0.0], half_spectrum=True)
    assert max(abs(x - y) for row, other in zip(full, half)
               for x, y in zip(row, other)) < 1e-12
    zero_a, zero_b = discretize([0j], [[2.0]], [0.3])
    assert zero_a == [1 + 0j] and abs(zero_b[0][0] - 0.6) < 1e-12
    normal = hippo_normal(4)
    assert all(abs(normal[n][k] + normal[k][n]) < 1e-12
               for n in range(4) for k in range(4) if n != k)
    assert prefix_scan([]) == []
    print(json.dumps({"scan_max_error": error, "outputs": full,
                      "checks": ["scan-loop", "conjugate-readout", "zero-pole", "normal-matrix", "empty-scan"]}, indent=2))


if __name__ == "__main__":
    main()