import random
import pylab
import pygame
import numpy as np

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREY = (128, 128, 128)
YELLOW = (255, 255, 0)
WIDTH = 12
HEIGHT = 12
MARGIN = 3


def initialize_model_matrix():
    matrix = np.zeros((30, 50))
    matrix[0, 0:50] = 2
    matrix[-1, 0:50] = 2
    return matrix


def set_people(matrix, count):
    people_count = count
    random_dict = {}
    for i in range(1, 29):
        random_dict[i] = []
    while people_count > 0:
        random_row = random.randint(1, 28)
        random_col = random.randint(0, 47)
        if random_col not in random_dict[random_row]:
            matrix[random_row, random_col] = 1
            random_dict[random_row].append(random_col)
            people_count -= 1


def set_walls(matrix, offset):
    matrix[offset, 48:50] = 2
    matrix[29 - offset, 48:50] = 2


def check_people(matrix):
    for i in range(1, 29, 1):
        for j in range(50):
            if matrix[i][j] == 1:
                return True
    return False


def move_people(matrix, screen, clock):
    iterations_result = 0
    while check_people(matrix):
        iterations_result += 1

        for is_move in [False, True]:
            for i in range(1, 29, 1):
                for j in range(49, -1, -1):
                    if is_move:
                        if matrix[i][j] == 3:
                            matrix[i][j] = 0
                        elif matrix[i][j] in [4, 5, 7]:
                            matrix[i][j] = 1
                        elif matrix[i][j] == 1:
                            if j + 1 == 50:
                                matrix[i][j] = 0
                                continue
                    else:
                        if (j + 1 < 50) and matrix[i][j] == 1:
                            if matrix[i][j + 1] in [0, 3]:
                                matrix[i][j] = 3
                                matrix[i][j + 1] = 7
                            elif (matrix[i + 1][j] in [0, 3]) and not (matrix[i - 1][j] in [0, 3]):
                                matrix[i][j] = 3
                                matrix[i + 1][j] = 5
                            elif not (matrix[i + 1][j] in [0, 3]) and (matrix[i - 1][j] in [0, 3]):
                                matrix[i][j] = 3
                                matrix[i - 1][j] = 4
                            elif (matrix[i + 1][j] in [0, 3]) and (matrix[i - 1][j] in [0, 3]):
                                matrix[i][j] = 3
                                random_index = random.choice([-1, 1])
                                if random_index == -1:
                                    matrix[i + random_index][j] = 4
                                else:
                                    matrix[i + random_index][j] = 5
        screen.fill(BLACK)
        for i in range(30):
            for j in range(50):
                color = WHITE
                if matrix[i][j] == 1:
                    color = YELLOW
                if matrix[i][j] == 2:
                    color = GREY
                pygame.draw.rect(screen, color,
                                 [(MARGIN + WIDTH) * j + MARGIN, (MARGIN + HEIGHT) * i + MARGIN, WIDTH, HEIGHT])
        clock.tick(10)
        pygame.display.flip()
    return iterations_result

offset = []
iterations = []
for i in range(1, 15, 1):
    pygame.init()
    WINDOW_SIZE = [750, 450]
    screen = pygame.display.set_mode(WINDOW_SIZE)
    pygame.display.set_caption("Crowd simulation")
    clock = pygame.time.Clock()
    model_matrix = initialize_model_matrix()
    set_people(model_matrix, 500)
    current_matrix = np.copy(model_matrix)
    set_walls(current_matrix, i)
    iterations.append(move_people(current_matrix, screen, clock))
    offset.append(i)
pylab.plot(offset, iterations)
pylab.show()
pygame.quit()