import cv2 

#from scipy.ndimage import gaussian_filter, uniform_filter
from cupyx.scipy.ndimage import gaussian_filter, uniform_filter
#import numpy as np
import cupy as np # don't laugh


# inspired by https://www.youtube.com/watch?v=X-iSQQgOd1A

speed = 2
damping = 0.03
steering = 0.9
turn_randomness = 0.3
straight_randomness = 0.1
#sense_angle = 0.33333*np.pi
sense_angle = 0.3*np.pi

img = np.zeros((1080, 2048, 3), dtype=np.float32)

# representation [y, x, angle, c] columns
agents = np.zeros((300_000, 3, 3), dtype=np.float32)
agents[:,0,:] = np.random.randint(0, img.shape[0] - 1, size=(agents.shape[0], 3))
agents[:,1,:] = np.random.randint(0, img.shape[1] - 1, size=(agents.shape[0], 3))
agents[:,2,:] = 2*np.pi*np.random.rand(agents.shape[0], 3)


def step(agents):
    """Move agents forward, bounce on the border"""
    for c in range(3):
        agents[:,2,c] = agents[:,2,c] % (2*np.pi)
        dy = speed*np.sin(agents[:,2,c])
        dx = speed*np.cos(agents[:,2,c])
        bounce_top = (agents[:,0,c] + dy) < 0
        bounce_bottom = (agents[:,0,c] + dy) >= img.shape[0]
        bounce_right = (agents[:,1,c] + dx) >= img.shape[1]
        bounce_left = (agents[:,1,c] + dx) < 0
        agents[bounce_top,2,c] = -agents[bounce_top,2,c]
        agents[bounce_bottom,2,c] = -agents[bounce_bottom,2,c]
        agents[bounce_right,2,c] = (-(agents[bounce_right,2,c] - np.pi/2) + np.pi/2)
        agents[bounce_left,2,c] = (-(agents[bounce_left,2,c] - np.pi/2) + np.pi/2)
        agents[:,2,c] = agents[:,2,c] % (2*np.pi)
        agents[:,0,c] += speed*np.sin(agents[:,2,c])
        agents[:,1,c] += speed*np.cos(agents[:,2,c])


def getVal(img, pos, angle, step, c):
    dir_y = np.round(np.sin(angle))
    dir_x = np.round(np.cos(angle))
    direction = np.vstack((dir_y, dir_x)).T
    new_pos = (pos + step*direction).astype('int16')
    new_pos[:,0] = np.clip(new_pos[:,0], 0, img.shape[0] - 1)
    new_pos[:,1] = np.clip(new_pos[:,1], 0, img.shape[1] - 1)
    return img[new_pos[:,0], new_pos[:,1], c], np.sum(img[new_pos[:,0], new_pos[:,1], :], axis=1)


ta = sense_angle
def sense(img, agents):
    
    for c in range(3):
        pos = agents[:,:2,c].astype('int16')
        front = left = right = 0

        for i in range(4):
            front, front_total = getVal(img, pos, agents[:,2,c], (i+1), c)
            left, left_total = getVal(img, pos, agents[:,2,c] + ta, (i+1), c)
            right, right_total = getVal(img, pos, agents[:,2,c] - ta, (i+1), c)
        
        if c == 2:
            front /= front_total
            left /= left_total
            right /= right_total
        else:
            front = front_total - front
            left = left_total - left
            right = right_total - right
        left_most = np.logical_and(left > right, left > front)
        right_most = np.logical_and(right > left, right > front)

        rot = np.zeros_like(agents[:,2,c])
        rot[left_most] = ta*((1 - turn_randomness) + 2*turn_randomness*np.random.rand(int(left_most.sum().item())))
        rot[right_most] = -ta*((1 - turn_randomness) + 2*turn_randomness*np.random.rand(int(right_most.sum().item())))
        
        center = np.logical_and(~left_most, ~right_most)
        rot[center] = 2*straight_randomness*(np.random.rand(int(center.sum().item())) - 0.5)
        
        agents[:,2,c] = (1-steering)*agents[:,2,c] + steering*(agents[:,2,c] + rot)


print("press q to quit")
overall_step = 0
while True:
    step(agents)
    sense(img, agents)
    img = np.maximum(img - damping, 0)
    #img = np.clip(img - damping, 0, 10)
    e = 1/1000*(-abs(overall_step - 1000) + 1000) + 0.2
    for c in range(3):
        img[agents[:,0,c].astype('int16'),agents[:,1,c].astype('int16'),c] = 1 # + 1 also gives interesting results
        img[:,:,c] = uniform_filter(img[:,:,c], 3)
    if hasattr(np, 'asnumpy'):
        cv2.imshow('image', np.asnumpy(img**e))
    else:
        cv2.imshow('image', img)
 
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

    overall_step = (overall_step + 1) % 2000

cv2.destroyAllWindows()