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


def generate_forest(density, grid_size):
	lattice = np.zeros((grid_size))
	number_of_trees = int(np.around(density*grid_size))
	#print("number of trees in artificial forest: ", number_of_trees)
	tree_sites = np.random.randint(0, grid_size, (number_of_trees))
	lattice[tree_sites] = 1
	return lattice


def induce_artificial_fire(lattice):
	tree_sites = np.where(lattice == 1)[0]
	ignition_location_index = np.random.randint(0, len(tree_sites))
	ignition_location = tree_sites[ignition_location_index]
	lattice, size_of_fire = burn_down_trees(lattice, ignition_location)
	return lattice, size_of_fire


def burn_down_trees(lattice, starting_point):
	size_of_fire = 0
	lattice[starting_point] = 2
	is_fire_active = True
	while (is_fire_active):

		ignited_sites = np.where(lattice == 2)[0]
		for site in ignited_sites:
			if (site < GRID_SIDE_LENGTH):
				site_above = (GRID_SIZE - GRID_SIDE_LENGTH) + site%GRID_SIDE_LENGTH
			else:
				site_above = site - GRID_SIDE_LENGTH
			if (site >= (GRID_SIZE - GRID_SIDE_LENGTH)):
				site_below = site%GRID_SIDE_LENGTH
			else:
				site_below = site + GRID_SIDE_LENGTH

			if (site%GRID_SIDE_LENGTH == 0):
				site_left = site + (GRID_SIDE_LENGTH - 1)
			else:
				site_left = site - 1
			if (site%GRID_SIDE_LENGTH == (GRID_SIDE_LENGTH-1) ):
				site_right = site - (GRID_SIDE_LENGTH - 1)
			else:
				site_right = site + 1

			if (lattice[site_above] == 1):
				lattice[site_above] = 2

			if (lattice[site_below] == 1):
				lattice[site_below] = 2

			if (lattice[site_left] == 1):
				lattice[site_left] = 2

			if (lattice[site_right] == 1):
				lattice[site_right] = 2

			lattice[site] = 3
			size_of_fire = size_of_fire + 1
		ignited_sites = (np.where(lattice == 2))[0]

		is_fire_active = (len(ignited_sites) > 0)
	return lattice, size_of_fire



GRID_SIDE_LENGTH = 128
GRID_SIZE = GRID_SIDE_LENGTH*GRID_SIDE_LENGTH
SIMULATION_LIMIT = 1000

pop_up_probability = 0.01
lightning_probability = 0.1

lattice = np.zeros((GRID_SIZE))


#Set up plot environment

"""
plt.ion()
fig = plt.figure()
ax1 = fig.add_subplot(111)

fig3 = plt.figure(3)
ax3 = fig3.add_subplot(111)

"""
total_number_of_fires = 0
data1 = np.zeros((SIMULATION_LIMIT))
data2 = np.zeros((SIMULATION_LIMIT))


while (total_number_of_fires < SIMULATION_LIMIT):
	#Creation of new trees
	for iLattice in range(len(lattice)):
		if (lattice[iLattice] != 1):
			r = np.random.rand()
			if (r < pop_up_probability):
				lattice[iLattice] = 1


	tree_locations = np.where(lattice == 1)[0]
	tree_density = len(tree_locations)/GRID_SIZE

	#Lightning strikes
	lightning_strike_occured = (np.random.rand() < lightning_probability)
	size_of_fire = 0

	if (lightning_strike_occured):
		random_site = np.random.randint(0, GRID_SIZE)
		has_ignited_tree = (lattice[random_site] == 1)
		if (has_ignited_tree):
			#print("Number of trees in original forest: ", len(tree_locations))
			lattice, size_of_fire = burn_down_trees(lattice, random_site)
			comp_forest = generate_forest(tree_density, GRID_SIZE)
			comp_lattice, comp_size_of_fire = induce_artificial_fire(comp_forest)

			#print("Size of original fire: ", size_of_fire)
			#print("Size of artificial fire: ", comp_size_of_fire)
			data1[total_number_of_fires] = size_of_fire/GRID_SIZE
			data2[total_number_of_fires] = comp_size_of_fire/GRID_SIZE
			total_number_of_fires = total_number_of_fires + 1

			print("Total number of fires ", total_number_of_fires)
			artificial_tree_locations = np.where(comp_lattice == 1)[0]
			artificial_fire_locations = np.where(comp_lattice == 3)[0]

			y_coordinates = artificial_tree_locations%GRID_SIDE_LENGTH
			x_coordinates = (artificial_tree_locations - y_coordinates)/GRID_SIDE_LENGTH

			y_coordinates_fire = artificial_fire_locations%GRID_SIDE_LENGTH
			x_coordinates_fire = (artificial_fire_locations - y_coordinates_fire)/GRID_SIDE_LENGTH

			"""
			ax3.clear()
			ax3.scatter(x_coordinates, y_coordinates, s = 10, c = 'g', marker = 's')
			ax3.scatter(x_coordinates_fire, y_coordinates_fire, s = 10, c = '#FFA500', marker = 's')
			plt.xlim(-1, GRID_SIDE_LENGTH+1)
			plt.ylim(-1, GRID_SIDE_LENGTH+1)

			fig3.suptitle('Density: ' + str(tree_density))

			fig3.canvas.draw()
			"""


	tree_locations = np.where(lattice == 1)[0]





	y_coordinates = tree_locations%GRID_SIDE_LENGTH
	x_coordinates = (tree_locations - y_coordinates)/GRID_SIDE_LENGTH

	fire_locations = np.where(lattice == 3)[0]
	y_coordinates_fire = fire_locations%GRID_SIDE_LENGTH
	x_coordinates_fire = (fire_locations - y_coordinates_fire)/GRID_SIDE_LENGTH

	#ax1.clear()
	#ax1.scatter(x_coordinates, y_coordinates, s = 10, c = 'g', marker = 's')
	#c = ax1.scatter(x_coordinates_fire, y_coordinates_fire, s = 10, c = '#FFA500', marker = 's')
	#plt.xlim(-1, GRID_SIDE_LENGTH+1)
	#plt.ylim(-1, GRID_SIDE_LENGTH+1)
	#fig.canvas.draw()
	#ax1.scatter(x_coordinates_fire, y_coordinates_fire, s = 10, c = '#ffffff')
	#c.remove()
	lattice[fire_locations] = 0

	#fig.canvas.draw()





	#plt.pause(0.001)


#plt.waitforbuttonpress()

original_data = np.sort(data1)[::-1]
artificial_data = np.sort(data2)[::-1]

ranking = np.zeros((SIMULATION_LIMIT))

for i in range(SIMULATION_LIMIT):
	ranking[i] = (i+1)/SIMULATION_LIMIT


np.save('fire_data.npy', original_data)
np.save('ranking.npy', ranking)
np.save('artificial_data.npy', artificial_data)
#print("original: ", original_data)
#print("artificial: ", artificial_data)
#print("ranking: ", ranking)

fig2 = plt.figure(2)
ax2 = fig2.add_subplot(111)
axes = fig2.gca()
ax2.scatter(original_data, ranking, s = 8, c = 'b')
ax2.scatter(artificial_data, ranking, s= 8,  c = 'r')
axes.set_xscale('log')
axes.set_yscale('log')
plt.xlim(0.00001, 2)
plt.ylim(0.0000001, 2)

fig2.canvas.draw()

plt.waitforbuttonpress()