# examples/videomamba.py
"""Deterministic VideoMamba arithmetic reference using only the standard library.

Run: python3 examples/videomamba.py --smoke
These scalar helpers do not train VideoMamba or reproduce checkpoint outputs.
"""

import json
import math


# region: tokens
def token_geometry(frames, height, width, tubelet=1, patch=16):
    """Count valid non-overlapping tubelets and the one global class token."""
    dimensions = (frames, height, width, tubelet, patch)
    if any(type(value) is not int or value <= 0 for value in dimensions):
        raise ValueError("Dimensions must be positive integers")
    time, rows, columns = frames // tubelet, height // patch, width // patch
    if not min(time, rows, columns):
        raise ValueError("Kernel must fit input")
    tokens = time * rows * columns
    return {"time": time, "rows": rows, "columns": columns,
            "tokens": tokens, "sequence": tokens + 1}
# endregion: tokens


# region: scan
def scan_order(time, rows, columns, temporal_first=False):
    """List spatial-first IDs in the selected scan's traversal order."""
    if any(type(value) is not int or value <= 0
           for value in (time, rows, columns)):
        raise ValueError("Grid must be positive")
    spatial = rows * columns
    if temporal_first:
        return [position + frame * spatial
                for position in range(spatial) for frame in range(time)]
    return list(range(time * spatial))
# endregion: scan


# region: bidirectional
def recurrence(values, decay, write):
    """Evaluate a scalar recurrence with explicitly supplied coefficients."""
    if not len(values) == len(decay) == len(write):
        raise ValueError("Equal-length arrays required")
    if not all(math.isfinite(v) for v in [*values, *decay, *write]):
        raise ValueError("Finite arrays required")
    state, output = 0.0, []
    for value, retention, gain in zip(values, decay, write):
        state = retention * state + gain * value
        output.append(state)
    return output


def bidirectional(values, decay, write):
    """Align a reversed scan before summing; weights are tied for this demo."""
    forward = recurrence(values, decay, write)
    backward = recurrence(values[::-1], decay[::-1], write[::-1])[::-1]
    return {"forward": forward, "backward": backward,
            "merged": [left + right for left, right in zip(forward, backward)]}
# endregion: bidirectional


# region: distillation
def feature_loss(student, teacher):
    """Squared L2 distance after normalization; zero vectors are rejected."""
    if not student or len(student) != len(teacher):
        raise ValueError("Matching nonempty features required")
    if not all(math.isfinite(v) for v in [*student, *teacher]):
        raise ValueError("Finite features required")
    student_norm, teacher_norm = math.hypot(*student), math.hypot(*teacher)
    if not student_norm or not teacher_norm:
        raise ValueError("Feature vectors must be nonzero")
    return sum((left / student_norm - right / teacher_norm) ** 2
               for left, right in zip(student, teacher))
# endregion: distillation


# region: masking
def retained_mask(time, rows, columns, mode="row"):
    """Make authored 50% retention fixtures, not upstream sampled masks."""
    length = len(scan_order(time, rows, columns))
    if mode not in ("row", "scattered"):
        raise ValueError("Unknown mask")
    return [(index // columns) % 2 == 0 if mode == "row" else index % 2 == 0
            for index in range(length)]


def adjacency(mask):
    """Count neighbors in the compacted stream that were already adjacent."""
    kept = [index for index, retain in enumerate(mask) if retain]
    pairs = max(0, len(kept) - 1)
    original = sum(right == left + 1 for left, right in zip(kept, kept[1:]))
    return {"kept": kept, "pairs": pairs, "original": original,
            "fraction": original / pairs if pairs else 0.0}
# endregion: masking


def main():
    """Exercise numeric invariants and emit a reproducible JSON receipt."""
    geometry = token_geometry(16, 224, 224)
    assert geometry["tokens"] == 3136
    assert geometry["sequence"] == 3137
    assert scan_order(2, 1, 3, True) == [0, 3, 1, 4, 2, 5]
    values = [0.0, 0.0, 1.0, 0.0, 0.0]
    result = bidirectional(values, [0.5] * 5, [1.0] * 5)
    assert result["merged"] == [0.25, 0.5, 2.0, 0.5, 0.25]
    assert math.isclose(feature_loss([1, 0], [-1, 0]), 4)
    assert math.isclose(feature_loss([2, 0], [3, 0]), 0)
    assert adjacency(retained_mask(2, 2, 4))["original"] == 6
    print(json.dumps({"geometry": geometry, "impulse": result,
                      "distillation": feature_loss([1, 0], [0, 1]),
                      "adjacency": adjacency(retained_mask(2, 2, 4))},
                     sort_keys=True))


if __name__ == "__main__":
    main()