import sys
sys.path.append("../../")

import numpy as np
from collections import deque
import matplotlib.pyplot as plt

from util.io import readvalue

def positive_int(s):
	i = int(s)
	if(not i > 0):
		raise ValueError("{} <= 0".format(i))
	return i
def float_0_1(s):
	i = float(s)
	if(not i > 0):
		raise ValueError("{} <= 0".format(i))
	if(i > 1):
		raise ValueError("{} > 1".format(i))
	return i

IS_TREE = 1
TREE_IS_ON_FIRE = 2
TREE_WILL_BE_ON_FIRE = 4

def prepare_model(N, p):
	model = np.zeros((N, N), dtype=np.int8)
	model[np.random.uniform(size=(N,N)) < p] = IS_TREE
	return model

def index_plus(i, j, add_i, add_j, model):
	new_i = i + add_i
	new_j = j + add_j

	if(new_i >= model.shape[0]):
		new_i -= model.shape[0]
	if(new_i < 0):
		new_i += model.shape[0]
	if(new_j >= model.shape[1]):
		new_j -= model.shape[1]
	if(new_j < 0):
		new_j += model.shape[1]

	return new_i, new_j
	

def do_tick(model):
	for i in range(model.shape[0]):
		for j in range(model.shape[1]):
			if((not model[i][j] & IS_TREE) or model[i][j] & TREE_IS_ON_FIRE):
				continue


			if(
				(model[index_plus(i, j, 1, 0, model)] & TREE_IS_ON_FIRE)
				or (model[index_plus(i, j, 0, 1, model)] & TREE_IS_ON_FIRE)
				or (model[index_plus(i, j, -1, 0, model)] & TREE_IS_ON_FIRE)
				or (model[index_plus(i, j, 0, -1, model)] & TREE_IS_ON_FIRE)
				):

				model[i][j] |= TREE_WILL_BE_ON_FIRE
			
	model[(model & TREE_WILL_BE_ON_FIRE) != 0] |= TREE_IS_ON_FIRE
	model[(model & TREE_WILL_BE_ON_FIRE) != 0] ^= TREE_WILL_BE_ON_FIRE

def trees_on_fire(model):
	unique, counts = np.unique((model & TREE_IS_ON_FIRE) != 0, return_counts=True)
	if(True not in unique):
		return 0
	return counts[np.where(unique == True)[0]][0]
	
def set_initial_on_fire(model):
	i = np.random.randint(0, model.shape[0])
	j = np.random.randint(0, model.shape[1])
	while True:
		if(model[i][j] & IS_TREE):
			model[i][j] |= TREE_IS_ON_FIRE
			break
		i = (i + 1) % model.shape[0]
		j = (j + 1) % model.shape[1]


N = readvalue("N > ", positive_int)
p = readvalue("p > ", float_0_1)

model = prepare_model(N, p)
set_initial_on_fire(model)

on_fire = deque()
last_on_fire = 1
while True:
	do_tick(model)
	now_on_fire = trees_on_fire(model)
	on_fire.append(now_on_fire)
	if(now_on_fire == last_on_fire):
		break
	last_on_fire = now_on_fire

plt.plot(*zip(*enumerate(on_fire)))
plt.show()