#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt
import time

GRIDSIZE = 300
INIT_CELLS = int(GRIDSIZE*0.1)
MAX_TIME = 300
EMPTY_CELL = -1

# query distance between two cars
def get_distance(grid):
    """ Compute distance between current position and next car """
    """
    Insert your code here
    """
    for n in grid:
        if grid[n] != EMPTY_CELL:
            return n
    return 5

# state transition t -> t + dt
def update(grid_old, grid_new):
    for i in range(GRIDSIZE):
        if grid_old[i] != EMPTY_CELL:
            dist = get_distance(grid_old[i+1:i+6])
            """ Update cars according rules using value 'dist' """
            """
            Insert your code here
            """
            if dist != -1:
                print(dist)

# allocate memory and initialise grids
grid_old = np.full((GRIDSIZE), EMPTY_CELL, dtype=np.int32)
grid_new = np.full((GRIDSIZE), EMPTY_CELL, dtype=np.int32)
traffic = np.zeros((MAX_TIME, GRIDSIZE), dtype=np.int32)

# set intial car positions and velocities
for k in range(INIT_CELLS):
    while True:
        i = int(float(GRIDSIZE)*np.random.random())
        if grid_old[i] == EMPTY_CELL:
            grid_old[i] = int(float(6)*np.random.random())
            break

# run updates
for t in range(MAX_TIME):
    traffic[t,:] = grid_old[:]
    update(grid_old, grid_new)
    for i in range(GRIDSIZE):
        grid_old[i] = grid_new[i]
        grid_new[i] = EMPTY_CELL
plt.xlabel('Cells')
plt.ylabel('Timesteps')
plt.imshow(traffic, cmap='Blues')
plt.show()
time.sleep(0.1)