#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt

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


# query distance between two cars
def get_distance(grid, i):
    dist = 0
    print(grid)
    while grid[(i+dist+1)%GRIDSIZE] == EMPTY_CELL and dist < 5:
        dist = dist + 1
    return dist

# state transition t -> t + dt
def update(grid_old, grid_new):
    for i in range(GRIDSIZE):

        if grid_old[i] != EMPTY_CELL:
            v = np.min((grid_old[i]+1, 5))    # accelerate
            dist = get_distance(grid_old, i)
            if i == 295:
                print(grid_old[i], " ", dist)
                quit()
            if v > dist: v = dist             # break
            p = np.random.random()
            if p == DALLY_LIMIT: v = max(grid_old[i-1], 0)
            grid_new[(i+v)%GRIDSIZE] = v      # move

# 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()