# examples/mambavision.py
"""Offline scalar mechanics reference for MambaVision, using Python's stdlib.

Run: python examples/mambavision.py
Source: NVlabs/MambaVision revision 7860a506b2eb844eaaae676f08461ce8c3c26f43.
Authored calculations, not the full backbone, trained outputs, or benchmarks.
"""

import json
import math


# region: hierarchy
def hierarchy(height, width, dim=80, depths=(1, 3, 8, 4)):
    """Returns pre-downsample stage shapes and released block ordering."""
    if min(height, width, dim, *depths) < 1 or len(depths) != 4:
        raise ValueError("Positive dimensions and four stage depths required")
    height, width = (height + 3) // 4, (width + 3) // 4
    stages = []
    for stage, depth in enumerate(depths):
        split = (depth + 1) // 2
        blocks = "C" * depth if stage < 2 else "M" * split + "A" * (depth - split)
        stages.append((height, width, dim * 2 ** stage, blocks))
        height, width = (height + 1) // 2, (width + 1) // 2
    return stages
# endregion: hierarchy


# region: convolution
def same_conv(sequence, kernel, causal=False):
    """Computes zero-padded cross-correlation, as Conv1d does.

    Args:
        sequence: Scalar token sequence.
        kernel: Ordered tap weights; even sizes pad more on the right.
        causal: Whether to place all padding on the left instead.

    Returns:
        A list with the same length as the input.
    """
    if not sequence or not kernel:
        raise ValueError("Nonempty input and kernel required")
    left = len(kernel) - 1 if causal else (len(kernel) - 1) // 2
    return [sum(weight * sequence[t + k - left]
                for k, weight in enumerate(kernel)
                if 0 <= t + k - left < len(sequence))
            for t in range(len(sequence))]
# endregion: convolution


# region: scan
def selective_scan(inputs, deltas, a=-1.0, b=1.0, c=1.0, direct=0.0):
    """Runs one channel, one state with supplied positive effective deltas.

    Uses the selective kernel's delta*B*u injection, not exact ZOH B.
    Learned delta projection/softplus and input-dependent B/C are omitted.
    """
    if len(inputs) != len(deltas) or a > 0:
        raise ValueError("Matching sequence lengths and nonpositive A required")
    state, outputs = 0.0, []
    for value, delta in zip(inputs, deltas):
        if not math.isfinite(delta) or delta < 0:
            raise ValueError("Finite nonnegative deltas required")
        state = math.exp(delta * a) * state + delta * b * value
        outputs.append(c * state + direct * value)
    return outputs
# endregion: scan


# region: mixer
def mixer(sequence, expansion=2):
    """Illustrates two branches with fixed scalar weights and a 2-to-1 output.

    Expansion documents width accounting only. This scalar demo is not a
    learned channel projection or an implementation of the full MLP block.
    """
    if expansion not in (1, 2):
        raise ValueError("Choose paper width 1 or released default 2")
    def silu(value):
        return value / (1 + math.exp(-value))
    x = [silu(v) for v in same_conv(sequence, (0.2, 0.6, 0.2))]
    z = [silu(v) for v in same_conv(sequence, (-0.1, 0.8, 0.3))]
    scanned = selective_scan(x, [0.5] * len(x))
    concatenated = list(zip(scanned, z))
    return [0.7 * state + 0.3 * local for state, local in concatenated]
# endregion: mixer


# region: windows
def window_layout(height, width, size):
    """Partitions row-major IDs, retaining None for right/bottom padding."""
    if min(height, width, size) < 1:
        raise ValueError("Positive sizes required")
    hp = math.ceil(height / size) * size
    wp = math.ceil(width / size) * size
    windows = []
    for top in range(0, hp, size):
        for left in range(0, wp, size):
            windows.append([
                row * width + column if row < height and column < width else None
                for row in range(top, top + size)
                for column in range(left, left + size)])
    return windows
# endregion: windows


def smoke_test():
    """Checks numeric boundaries and prints deterministic teaching output."""
    assert hierarchy(224, 224)[2] == (14, 14, 320, "MMMMAAAA")
    assert hierarchy(225, 193)[-1][:2] == (8, 7)
    assert same_conv([1, 2, 3], [1, 1, 1]) == [3, 6, 5]
    assert same_conv([1, 2, 3], [1, 1, 1], True) == [1, 3, 6]
    assert same_conv([1, 2, 3], [1, 1, 1, 1]) == [6, 6, 5]
    assert selective_scan([1, 2], [0, 0], direct=1) == [1, 2]
    assert selective_scan([1, 0], [1, 1])[1] == math.exp(-1)
    windows = window_layout(3, 5, 2)
    ids = [value for window in windows for value in window if value is not None]
    assert sorted(ids) == list(range(15)) and len(windows) == 6
    assert mixer([0, 0, 0]) == [0, 0, 0]
    print(json.dumps({"stages": hierarchy(224, 224),
                      "scan": selective_scan([1, 0, 0, 0], [0.5] * 4),
                      "mixer": mixer([1, 0, 0, 0])}, indent=2))


if __name__ == "__main__":
    smoke_test()