#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt

GRIDSIZE_X = 35
GRIDSIZE_Y = 35
MAX_TIME = 500
NUM_PEDS = 200

CELL_PED = 1   # cell state: pedestrian
CELL_EMP = 0   # cell state: empty
CELL_OBS = -1  # cell state: obstacle

EXIT_X = GRIDSIZE_X+1       # x-coordinate of exit
EXIT_Y = int(GRIDSIZE_Y/2)  # y-coordinate of exit

VIS_PAUSE = 0.2  # time [s] between two visual updates
VIS_STEPS = 2    # stride [steps] between two visual updates

# count pedestrians left in domain
def count_peds(grid):
    nped = 0
    for x in range(1, GRIDSIZE_X+1):
        for y in range(1, GRIDSIZE_Y+1):
            if grid[x,y] == CELL_PED:
                nped = nped + 1
    return nped

# compute average density [P/m2]
def comp_density(grid):
    dens = 0
    for x in range(1, GRIDSIZE_X+1):
        for y in range(1, GRIDSIZE_Y+1):
            if grid[x,y] == CELL_PED:
                npeds = 0
                for i in range(-1,2):
                    for j in range(-1,2):
                        if grid[x+i,y+j] == CELL_PED:
                            npeds = npeds + 1
                dens = dens + (npeds/1.44)
    if count_peds(grid) != 0:
        return (dens/count_peds(grid))
    else:
        return 0

# state transition t -> t + dt
def update(old, new):
    for x in range(1, GRIDSIZE_X+1):
        for y in range(1, GRIDSIZE_Y+1):
            #
            # transition functions
            #
            if old[x,y] == CELL_PED:
                delta_x = EXIT_X - x
                delta_y = EXIT_Y - y
                    
                terrain = dict(N = {}, E = {}, S = {}, W = {})

                terrain["N"]["x"], terrain["N"]["y"], terrain["N"]["cell"], terrain["N"]["vector"] = set_terrain(old, x, y+1)
                terrain["E"]["x"], terrain["E"]["y"], terrain["E"]["cell"], terrain["E"]["vector"] = set_terrain(old, x+1, y)
                terrain["S"]["x"], terrain["S"]["y"], terrain["S"]["cell"], terrain["S"]["vector"] = set_terrain(old, x, y-1)
                terrain["W"]["x"], terrain["W"]["y"], terrain["W"]["cell"], terrain["W"]["vector"] = set_terrain(old, x-1, y)

                min = np.inf
                move_to = None
                
                # Alternative move_to when optimal move_to is blocked
                for key in terrain:
                    if terrain[key]["cell"] == CELL_EMP:
                        if terrain[key]["vector"] < min:
                            min = terrain[key]["vector"]
                            move_to = terrain[key]
                    
                if move_to != None:
                    new[move_to["x"], move_to["y"]] = CELL_PED
                    old[move_to["x"], move_to["y"]] = CELL_OBS
                    old[x,y] = CELL_OBS
                else:
                    new[x,y] = old[x,y]
                
                # Stay still when optimal move_to is blocked            
                '''for key in terrain:
                    if terrain[key]["vector"] < min:
                        min = terrain[key]["vector"]
                        move_to = terrain[key]
                
                if move_to != None and move_to["cell"] == CELL_EMP:
                    new[move_to["x"], move_to["y"]] = CELL_PED
                    old[move_to["x"], move_to["y"]] = CELL_OBS
                    old[x,y] = CELL_OBS
                else:
                    new[x,y] = old[x,y]'''
                

def set_terrain(grid, x, y):
    cell = grid[x,y]
    vector = np.sqrt((EXIT_X - x)**2 + (EXIT_Y - y)**2)
    return x, y, cell, vector

# allocate memory and initialise grids
old = np.zeros((GRIDSIZE_X+2, GRIDSIZE_Y+2), dtype=np.int32)
new = np.zeros((GRIDSIZE_X+2, GRIDSIZE_Y+2), dtype=np.int32)
old[:,0] = CELL_OBS   # boundary: south
old[:,-1] = CELL_OBS  # boundary: north
old[0,:] = CELL_OBS   # boundary: west
old[-1,:] = CELL_OBS  # boundary: east
old[EXIT_X,EXIT_Y-1] = CELL_EMP # exit
old[EXIT_X,EXIT_Y] = CELL_EMP   # exit
old[EXIT_X,EXIT_Y+1] = CELL_EMP # exit
new = old.copy()

# set random starting points for pedestrians
for i in range(NUM_PEDS):
    while True:
        x = int(float(GRIDSIZE_X)*np.random.random())+1
        y = int(float(GRIDSIZE_Y)*np.random.random())+1
        if old[x,y] == CELL_EMP:
            old[x,y] = CELL_PED
            break

# run updates
time = 0
dens = []
plt.ion()
plt.imshow(np.rot90(old, 1))
plt.pause(VIS_PAUSE)
while count_peds(old) > 0 and time < MAX_TIME:
    new[1:GRIDSIZE_X+1,1:GRIDSIZE_Y+1] = CELL_EMP
    update(old, new)
    new[EXIT_X,EXIT_Y-1] = CELL_EMP  # clear exit
    new[EXIT_X,EXIT_Y] = CELL_EMP    # clear exit
    new[EXIT_X,EXIT_Y+1] = CELL_EMP  # clear exit
    old = new.copy()
    numpeds = count_peds(old)
    dens.append(comp_density(old))
    time = time + 1
    if time%VIS_STEPS == 0:
      plt.clf()
      plt.imshow(np.rot90(old, 1))
      plt.pause(VIS_PAUSE)
plt.ioff()
print(f'Evacuation of {NUM_PEDS} persons done in {time*0.3:.2f} seconds')

# plot density diagramm
x = np.linspace(0,time*.3,time)
plt.clf()
plt.xlabel('Zeit [s]')
plt.ylabel('Dichte [P/m2]')
plt.plot(x, dens, 'r-')
plt.show()