# examples/mamba3.py
"""Deterministic Mamba-3 arithmetic references using only the standard library.

These scalar/two-dimensional examples explain paper equations, not a trained
language model. They omit learned projections, normalization, gates and kernels.
Run: python3 examples/mamba3.py --smoke
"""

import argparse
import json
import math


# region: discretization
def coefficients(delta, a, weight):
    """Returns alpha, beta, gamma for the exponential-trapezoidal rule.

    Args:
        delta: Nonnegative time interval.
        a: Nonpositive real decay generator.
        weight: Current-endpoint weight in [0, 1].

    Returns:
        Three scalar coefficients from paper Eq. 6.
    """
    if not all(map(math.isfinite, (delta, a, weight))):
        raise ValueError("Finite parameters required")
    if delta < 0 or a > 0 or not 0 <= weight <= 1:
        raise ValueError("Require delta >= 0, a <= 0, weight in [0, 1]")
    alpha = math.exp(delta * a)
    return alpha, (1 - weight) * delta * alpha, weight * delta


def scan(inputs, delta, a, weight):
    """Scans scalar state-inputs v=B*x, with h[-1]=v[-1]=0."""
    alpha, beta, gamma = coefficients(delta, a, weight)
    state, previous = 0.0, 0.0
    outputs = []
    for value in inputs:
        state = alpha * state + beta * previous + gamma * value
        outputs.append(state)
        previous = value
    return outputs
# endregion: discretization


# region: mask
def causal_mask(length, delta, a, weight):
    """Expands the recurrence into a lower triangular scalar input map."""
    if not isinstance(length, int) or length < 1:
        raise ValueError("Positive length required")
    alpha, beta, gamma = coefficients(delta, a, weight)
    return [
        [0.0 if s > t else gamma * alpha ** (t - s)
         + (beta * alpha ** (t - s - 1) if s < t else 0.0)
         for s in range(length)]
        for t in range(length)
    ]
# endregion: mask


# region: rotation
def rotate(vector, angle):
    """Applies one real 2-by-2 rotation, equivalent to complex phase."""
    x, y = vector
    return (math.cos(angle) * x - math.sin(angle) * y,
            math.sin(angle) * x + math.cos(angle) * y)


def rotating_scan(inputs, angles, delta, a, weight, transformed=False):
    """Computes physical or Eq. 11 transformed-frame outputs for B=C=(1,0)."""
    if len(inputs) != len(angles):
        raise ValueError("One angle per input required")
    alpha, beta, gamma = coefficients(delta, a, weight)
    state, previous, phase = (0.0, 0.0), (0.0, 0.0), 0.0
    outputs = []
    for value, angle in zip(inputs, angles):
        phase += angle
        write = rotate((value, 0), -phase) if transformed else (value, 0)
        retained = state if transformed else rotate(state, angle)
        prior = previous if transformed else rotate(previous, angle)
        state = tuple(alpha * h + beta * old + gamma * new
                      for h, old, new in zip(retained, prior, write))
        previous = write
        physical = rotate(state, phase) if transformed else state
        outputs.append(physical[0])
    return outputs
# endregion: rotation


# region: mimo
def mimo_write(b, x):
    """Computes B @ X.T with shapes [N,R], [P,R] -> [N,P]."""
    if not b or not x or not b[0]:
        raise ValueError("Nonempty matrices required")
    rank = len(b[0])
    if any(len(row) != rank for row in b + x):
        raise ValueError("Matching rectangular rank required")
    return [[sum(br * xr for br, xr in zip(row, column))
             for column in x] for row in b]
# endregion: mimo


# region: intensity
def arithmetic_intensity(n, p, rank):
    """Returns Table 2's modeled FLOPs/byte for 2-byte values, not latency."""
    if any(not isinstance(v, int) or v < 1 for v in (n, p, rank)):
        raise ValueError("Positive integer dimensions required")
    flops = 4 * n * p * rank + n * p - p * rank
    byte_count = 2 * (1 + 2 * n * rank + p * rank + n * p)
    return flops / byte_count
# endregion: intensity


def smoke():
    """Checks independent formulations and emits deterministic calculated output."""
    inputs = [1, -2, 0, 3, 1]
    angles = [0.2, -0.7, 0.4, 1.2, -0.1]
    direct = rotating_scan(inputs, angles, 0.4, -0.7, 0.3)
    rotary = rotating_scan(inputs, angles, 0.4, -0.7, 0.3, True)
    assert max(abs(a - b) for a, b in zip(direct, rotary)) < 1e-12
    matrix = causal_mask(len(inputs), 0.4, -0.7, 0.3)
    parallel = [sum(a * b for a, b in zip(row, inputs)) for row in matrix]
    sequential = scan(inputs, 0.4, -0.7, 0.3)
    assert max(abs(a - b) for a, b in zip(parallel, sequential)) < 1e-12
    assert mimo_write([[1, 2], [3, 4]], [[1, 0], [0, 1]]) == [[1, 2], [3, 4]]
    assert coefficients(0, -1, 0.5) == (1, 0, 0)
    print(json.dumps({"scalar": sequential, "rotating": direct,
                      "table2_intensity_r4": arithmetic_intensity(128, 64, 4)},
                     sort_keys=True))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", help="Run deterministic checks")
    parser.parse_args()
    smoke()