from itertools import count, cycle
from math import ceil, sqrt
from numpy import arange, array, full, ones, insert, where
from numpy import append as ap
from time import perf_counter as pf
from time import sleep
from tkinter import Canvas, Tk, constants

root = Tk()
# root.attributes('-fullscreen', True) # make fullscreen window
n, w, h = 1, 1300, 1300 # n=2 because 1 messes with the algo, starting number
canvas = Canvas(root, bg="black", width=w, height=h)
ct = canvas.create_text

text_size, text_font = 2, "Hack" # text size
pixels_to_jump = text_size/0.9 # not sure why this value works!
upto = 50_000 # 100M will take ~2.3GB RAM and 20s on powerful machine
ANIMATION_DELAY = 0.0125 # def: 0.05
PRINT_CHAR = '❖' # '×'

primes = arange(2,upto+1)
isprime = ones((upto-1),dtype=bool)
R, U, L, D = ((pixels_to_jump, 0), 0), ((0, -pixels_to_jump), 90), \
        ((-pixels_to_jump, 0), 180), ((0, pixels_to_jump), 270)

run_times = {}

# print_me = arange(1,upto)
# ^ uncomment for numbers, also switch the ct() in draw loop
print_me = full((len(primes)) + 1, ' ')

def prime6(): # new prime generation algo
    """Prime number generation algo"""

    time_it(1)
    for factor in primes[:int(sqrt(upto))]:
        if isprime[(factor)//2]: # if it is an even number
            isprime[((factor*3)-2)::factor]=0 # find a factor amongst the existing primes
    insert(primes[isprime],0,1) # no factors found, add to list of primes
    time_it(1)

def direction():
    """Legacy function, will be removed next version"""
    for _ in range(1, int(sqrt(upto))):
        # return cycle((R, U, L, D)) # movement sequence
        return cycle((L, D, R, U)) # movement sequence
# ((-pixels_to_jump, 0), 180), ((0, pixels_to_jump), 270)...

def distance():
    """Legacy function, will be removed next version"""
    # loop increments x every 2 loops
    for x in range(1, int(sqrt(upto))+1):
        for _ in (0, 1):
            yield x
# 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6...

def transform_array():
    """Star of the show. Matches, folds and trims"""
    global print_me, primes, isprime

    time_it(2)
    # list comprehension to match PRINT_CHAR if prime, otherwise ' '
    time_it(3)
    print_me = [PRINT_CHAR if isprime[_] else ' ' for _ in range(len(primes))]
    time_it(3)

    # Pre-made list of what used to be distance(), named so for readability later, sorry!
    # should be 'distance_list'
    d_l = [x for x in range(1, int(sqrt(upto))+1) for y in (0, 1)]

    time_it(5)
    # fold matrix into the required shape ([[1],[2],[3, 4],[5, 6]...])
    # although print_me[X:Y] looks like a mess, the maths is simple:
    # X = sum(d_l[0:i] = WHAT HAS BEEN READ ALREADY
    # Y = sum(d_l[0:i]) + d_l[i] = HOW FAR NEEDS TO BE READ FROM X 
    # putting these two together, we can cobble together the required output
    temp_array = [print_me[(sum(d_l[0:i])) : sum(d_l[0:i]) + d_l[i]] \
                            for i in range(0, len(d_l))]
    time_it(5)

    # trim empty elements from array
    time_it(6)
    print_me = [_ for _ in temp_array if len(_) != 0]
    time_it(6)

    temp_array = []
    time_it(2)

def time_it(index):
    """Timing function, takes int for index,
       retrieve values from array run_times[n]"""
    if index in run_times: # number exists, stop timer
        run_times[index] = (pf() - run_times[index]) * 1000
        return run_times[index] # return milliseconds

    else: # index does not exist
        if (0 > index) or (index > 10): # value error
            return "Enter value 1-10."
        else: # start timer
            run_times[index]=pf()

def draw_spiral():
    """Draw the spiral on screen"""
    text_x = w // 2; text_y = h // 2;

    time_it(7)
    for _, c, coords in zip(range(len(print_me)), distance(), direction()):
        delta_x, delta_y = coords[0][0], coords[0][1] 
        offset_x = ((len(print_me[_]) * delta_x) / 0.75)
        offset_y = ((len(print_me[_]) * delta_y) / 0.75)

        text_x += delta_x + offset_x
        text_y += delta_y + offset_y
        # this expression works for int array, DO NOT DELETE
        # ct(text_x, text_y, fill="white", text=(str(print_me[_])[1:-1]),\
        #                         font=(text_font, text_size), angle = coords[1])
        # this expression works for str arrays, DO NOT DELETE
        ct(text_x, text_y, fill="white", text=(' '.join(print_me[_])),\
            font=(text_font, text_size), angle = coords[1], justify=constants.CENTER)
        
        text_x += delta_x + offset_x
        text_y += delta_y + offset_y

        root.update_idletasks() # used to update canvas in realtime
        canvas.pack()
        sleep(ANIMATION_DELAY) # recom: (upto / (upto * 20)) i.e 0.05
    time_it(7)

    # canvas.pack()
    # canvas.update()
    # canvas.postscript(file="ulam.ps", colormode='color')
    root.mainloop()

def print_stats():
    """Print time stats on console"""
    ANIMATION_TIME = (len(print_me) * (ANIMATION_DELAY * 1000))

    print(f"\nTimes for {len(primes)+1:,} numbers;\n\t  {(len(primes[isprime])):,} primes:")
    # str(primes[isprime][-1:])[1:-1] is a ridiculous convolution,
    # to print the value without brackets, wrapped in int() to comma separate
    print(f"[1] primes  : {(run_times[1] * 1000):.2f}µs (max: {int(str(primes[isprime][-1:])[1:-1]):,})")
    print(f"[2] xform   : {run_times[2]:.2f}ms")
    print(f"[3]  --match: {run_times[3]:.2f}ms")
    print(f"[4]  --list : n/a") # {run_times[4] * 1000:.2f}µs")
    print(f"[5]  --fold : {run_times[5]:.2f}ms")
    print(f"[6]  --trim : {run_times[6]:.2f}ms")
    print(f"[7]  draw   : {(run_times[7] - ANIMATION_TIME):,.2f}ms (+{(run_times[7]/1000):.2f}s animation)")
    print(f"    TOTAL   : {(run_times[1] + run_times[2] + run_times[7]) - ANIMATION_TIME:,.2f}ms\n")

def main():
    prime6()
    transform_array()
    draw_spiral()

    ## Stats:
    # if foo:
    print_stats()

if __name__ == "__main__":
    main()

########################
##### DOWN & DIRTY #####
########################
## Records (50K, delay=0)
## v1.0-
# primes: 295.13µs
#  xform: 11840.98ms    80%
#   --match: 4702.20ms  --40%
#   -- fold: 6112.60ms  --52%
#   -- trim: 624.96ms
#   draw: 2.87s         20%
#  total: 14.71s
# ------------------------------
## v1.1- 207x faster
# [1] primes  : 270.85µs
# [2] xform   : 4.59ms      2574x
# [3]  --match: 2.91ms      1621x
# [4]  -- list: 14.25µs
# [5]  -- fold: 1.56ms      3918x
# [6]  -- trim: 0.09ms      6933x
# [7]  draw   : 66.30ms     43x (not even optimised yet)
#        TOTAL: 71.16ms     207x
# ------------------------------
## v1.1- a real challenge (upto=100M, delay=0)
# [1] primes  : 1,267,025.34µs
# [2] xform   : 8,951.02ms
# [3]  --match: 5,496.00ms
# [4]  -- list: 514.88µs
# [5]  -- fold: 3,298.45ms
# [6]  -- trim: 155.93ms
# [7]  draw   : 8,631.62ms
#     TOTAL   : 18,849.67ms
# -------------------------------