# A small Mojo 1.0 tour of CPU SIMD.
#
# Run with:
#   pixi run simd-demo
#
# The program uses the CPU's native Float32 lane count. It prints each packet
# before and after one vector operation so the data-parallel behaviour remains
# visible instead of disappearing inside a benchmark.

from std.sys.info import simd_width_of


comptime DTYPE = DType.float32
comptime LANES = simd_width_of[DTYPE]()


def show_packet[
    width: SIMDLength,
](label: String, values: SIMD[DTYPE, width]):
    print(label, values)


def expose_pixel(pixel: Float32) -> Float32:
    """Increase one normalised pixel value and clamp it to display white."""
    var exposed = pixel * 1.5
    return 1.0 if exposed > 1.0 else exposed


def expose_packet[
    width: SIMDLength,
](values: SIMD[DTYPE, width]) -> SIMD[DTYPE, width]:
    """Apply the same image operation to every pixel lane."""
    var exposed = values * 1.5
    var over_white = exposed.gt(1.0)
    return over_white.select(SIMD[DTYPE, width](1.0), exposed)


def demonstrate_one_register():
    print("1. One SIMD value contains several independent image pixels")
    print("   Native Float32 lane count:", Int(LANES))

    var input = SIMD[DTYPE, LANES](0.0)
    for lane in range(Int(LANES)):
        input[lane] = Float32(lane + 1) * 0.25

    var exposed = input * 1.5
    var over_white = exposed.gt(1.0)
    var output = over_white.select(SIMD[DTYPE, LANES](1.0), exposed)

    show_packet("   input pixels:  ", input)
    show_packet("   after exposure:", exposed)
    print("   over white?:   ", over_white)
    show_packet("   output pixels: ", output)
    print()


def demonstrate_scanline_packets_and_tail():
    print("2. Process one grayscale image scanline in SIMD packets")

    var count = Int(LANES) * 2 + 3
    var input_pixels = List[Float32]()
    var output_pixels = List[Float32]()

    for index in range(count):
        # Produce normalised grayscale pixels in the inclusive range 0.0...1.0.
        input_pixels.append(Float32((index * 3) % 8 + 1) * 0.125)
        output_pixels.append(0.0)

    var cursor = 0
    var packet_number = 0

    # Process every complete packet. Each lane owns one independent pixel.
    while cursor + Int(LANES) <= count:
        var packet = SIMD[DTYPE, LANES](0.0)
        for lane in range(Int(LANES)):
            packet[lane] = input_pixels[cursor + lane]

        var result = expose_packet(packet)
        print("   packet", packet_number, "covers indices", cursor, "to", cursor + Int(LANES) - 1)
        show_packet("      input:  ", packet)
        show_packet("      output: ", result)

        # Store each transformed lane in its corresponding output pixel.
        for lane in range(Int(LANES)):
            output_pixels[cursor + lane] = result[lane]

        cursor += Int(LANES)
        packet_number += 1

    # Fewer than LANES pixels remain, so finish with the scalar reference.
    print("   scalar tail starts at index", cursor)
    while cursor < count:
        output_pixels[cursor] = expose_pixel(input_pixels[cursor])
        print(
            "      index",
            cursor,
            ":",
            input_pixels[cursor],
            "->",
            output_pixels[cursor],
        )
        cursor += 1

    print()
    print("   input scanline: ", input_pixels)
    print("   output scanline:", output_pixels)
    print()


def main():
    print("SIMD in Mojo 1.0")
    print("================")
    print("One instruction, several values. Each value occupies one lane.")
    print("Example: raise normalised brightness by 50%, then clamp to white (1.0).")
    print()

    demonstrate_one_register()
    demonstrate_scanline_packets_and_tail()

    print("3. The perspective shift")
    print("   Scalar question: what happens to this value?")
    print("   SIMD question:  what happens to this packet of values?")
    print("   Masks preserve lane-by-lane decisions without splitting the packet.")