# examples/efficientvmamba.py
"""Deterministic scalar ES2D reference, using only the Python standard library.

Run `python examples/efficientvmamba.py`. This demonstrates routing, recurrence,
scatter and local mixing. It is not the trained network or a speed benchmark.
Routing follows EfficientVMamba commit 0bc5ee288b402d648641f5494b73e9d152e0c62b,
classification/lib/models/mamba/efficient_scan.py with step_size=2.
"""

import math


# region: routing
def routes(height, width):
    """Return padded coordinates in the four released p=2 traversal orders."""
    if any(not isinstance(v, int) or isinstance(v, bool) or v <= 0
           for v in (height, width)):
        raise ValueError("Spatial dimensions must be positive integers")
    rows, cols = (height + 1) // 2, (width + 1) // 2
    paths = []
    for group, (dy, dx) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
        path = []
        outer, inner = (cols, rows) if group % 2 else (rows, cols)
        for a in range(outer):
            for b in range(inner):
                row, col = (b, a) if group % 2 else (a, b)
                path.append((2 * row + dy, 2 * col + dx))
        paths.append(path)
    return paths


def gather(values, height, width):
    """Gather valid features, inserting zero for right/bottom padding."""
    if len(values) != height * width:
        raise ValueError("Input shape mismatch")
    return [[values[y * width + x] if y < height and x < width else 0.0
             for y, x in path] for path in routes(height, width)]
# endregion: routing


# region: recurrence
def selective_scan(values, delta=0.5, a=-1.0, b=1.0, c=1.0, d=0.0):
    """Compute a scalar selective recurrence with supplied coefficients.

    Delta may be a per-token list. Learned projections are intentionally absent.
    The input update uses delta * B * u, as in the released scan kernel family,
    rather than the exact zero-order-hold input integral.
    """
    steps = delta if isinstance(delta, (list, tuple)) else [delta] * len(values)
    if len(steps) != len(values) or a > 0:
        raise ValueError("Invalid time scales or unstable A")
    if not all(math.isfinite(v) for v in [a, b, c, d, *values, *steps]):
        raise ValueError("All values must be finite")
    if any(v < 0 for v in steps):
        raise ValueError("Time scales must be nonnegative")
    state, outputs = 0.0, []
    for value, dt in zip(values, steps):
        state = math.exp(dt * a) * state + dt * b * value
        outputs.append(c * state + d * value)
    return outputs
# endregion: recurrence


# region: merge
def merge(streams, height, width):
    """Scatter one output per coordinate and crop padding; do not sum streams."""
    paths = routes(height, width)
    if len(streams) != 4 or any(len(s) != len(p)
                               for s, p in zip(streams, paths)):
        raise ValueError("Expected four padded streams")
    output = [0.0] * (height * width)
    for stream, path in zip(streams, paths):
        for value, (y, x) in zip(stream, path):
            if y < height and x < width:
                output[y * width + x] = value
    return output
# endregion: merge


# region: fusion
def local_mean(values, height, width):
    """Apply a zero-padded 3x3 mean, standing in for learned local mixing."""
    routes(height, width)
    if len(values) != height * width:
        raise ValueError("Input shape mismatch")
    output = []
    for y in range(height):
        for x in range(width):
            neighbors = [values[(y + dy) * width + x + dx]
                         for dy in (-1, 0, 1) for dx in (-1, 0, 1)
                         if 0 <= y + dy < height and 0 <= x + dx < width]
            output.append(sum(neighbors) / 9.0)
    return output


def illustrative_fusion(values, height, width, scan_gate=0.5, local_gate=0.5):
    """Combine branches with supplied scalar gates, without learned SE or MLP."""
    scanned = [selective_scan(s) for s in gather(values, height, width)]
    global_values = merge(scanned, height, width)
    local_values = local_mean(values, height, width)
    return [scan_gate * g + local_gate * loc
            for g, loc in zip(global_values, local_values)]
# endregion: fusion


# region: cost
def scan_work(height, width):
    """Count visits, including padding, not FLOPs or measured latency."""
    paths = routes(height, width)
    dense = 4 * height * width
    sparse = sum(len(path) for path in paths)
    return {"dense": dense, "sparse": sparse,
            "ratio": dense / sparse, "padding": sparse - height * width}
# endregion: cost


def main():
    """Check deterministic identities and print a small inspectable example."""
    for height, width in ((1, 1), (1, 7), (3, 5), (4, 6), (7, 7)):
        values = list(range(1, height * width + 1))
        assert merge(gather(values, height, width), height, width) == values
    values = [1.0] + [0.0] * 23
    global_values = merge([selective_scan(s) for s in gather(values, 4, 6)], 4, 6)
    assert global_values[1] == 0.0
    assert global_values[2] > 0.0
    assert local_mean(values, 4, 6)[1] == 1 / 9
    assert selective_scan([1, 3], delta=0, d=2) == [2.0, 6.0]
    assert scan_work(3, 5)["ratio"] == 2.5
    assert scan_work(4, 6)["ratio"] == 4.0
    # Independent closed-form impulse decay checks recurrence indexing.
    output = selective_scan([1.0, 0.0, 0.0], delta=0.5)
    assert all(math.isclose(v, 0.5 * math.exp(-0.5 * i))
               for i, v in enumerate(output))
    print("Routing round-trips, recurrence, impulse and cost checks passed.")
    print("3x5 visit counts:", scan_work(3, 5))
    print("4x6 fused impulse:", [round(v, 4) for v in illustrative_fusion(values, 4, 6)])


if __name__ == "__main__":
    main()