# pico_sieve_ulab.py

# Raspberry Pi Pico - Sieve of Eratosthenes demo using ulab native types

# Calculate all the prime numbers within a range of low integers.  Note that
# this is purely computational.  This version uses the CircuitPython ulab module
# to represent the flag array using native machine data types.  This speeds up
# operation but also uses less memory, allowing the array to be larger.  The
# ulab module implements a subset of numpy.

import time
import math

# documentation: https://circuitpython.readthedocs.io/en/latest/shared-bindings/ulab/
import ulab

size = 40000

# Capture an initial time stamp to compute runtime.
start = time.monotonic()

# Array of boolean flags, one per integer.  A True represents a possible prime number,
# a False a composite number.
flags = ulab.ones((size), dtype=ulab.bool)

# Walk up through the array, identifying all prime numbers and removing their
# multiples from the set of integers by setting the corresponding flag to False.
# Note that the scan can stop once values are reached for which no unique
# multiples are included in the array.
for i in range(2, int(math.sqrt(size))):
    if flags[i]:
        # This value is a prime, now remove all the multiples from the set.
        # Note that this marking may begin at i*i since all lower multiples will
        # have already been removed.
        multiple = i * i
        
        if multiple < size:
            # Set a slice of the array to zeros in one step using fast iteration in native code.
            flags[multiple::i] = 0        

# Capture an final time stamp to compute runtime.
end = time.monotonic()

# Any remaining True values are prime
print("Prime numbers:")

for i in range(2, size):
    if flags[i]:
        print(i)

print("Sieve running time:", end - start, "seconds.")