#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
import random
import sys

# Initialize
plt.close('all')

# calculs de trajectoire parabolique
def parabole_fc(xx, angled, vini, y0):
    # g = 5.8  # Acceleration due to gravity in m/s^2
    g = 9.8  # Acceleration due to gravity in m/s^2
    yy = y0 + np.tan(np.radians(angled)) * xx - (g / 2) * ((xx / (vini * np.cos(np.radians(angled))))**2)
    return yy

# Implémentation du maillage
xmax = 400
dx = 2
x = np.arange(0, xmax + dx, dx)
X = -x
Nx = len(x) # xmax/dx +1
# print(Nx)
# sys.exit()
n_iterations = 10000
h = np.zeros((Nx, n_iterations+1))

# Building the topography of the cinder cone
n_plots = 50
i_plot = n_iterations // n_plots

# Initialisation des paramètres
dh = .8
theta_min = 75 # 75
theta_max = 90  # 90
V0 = 50.0 # vitesse initiale 50 +- 10
dV0 = 10.0
hmax = 200
landsliding = True
crit_slope = np.tan(np.radians(30)) # pente critique
phi = 0.7 # probabilite de réabsorption des scories dans la cheminee 0.7

# test taille particules
if dh/dx > crit_slope:
    print('particules trop grande ou dx trop faible')
    sys.exit()

# Create figures
fig = plt.figure(0, figsize=(10,8))
ax = fig.add_subplot()
ax.set_xlim(-xmax, xmax)
ax.set_ylim(0, hmax)
ax.set_title('Growth of a Cinder Cone', fontsize=18)
ax.set_xlabel('Distance from vent (m)', fontsize=18)
ax.set_ylabel('Height of cinder cone (m)', fontsize=18)

for i in range(1, n_iterations + 1):
    v0 = random.gauss(V0, dV0)
    angle0 = random.uniform(theta_min, theta_max)

    # Calculate trajectory
    y = parabole_fc(x, angle0, v0, h[0,i])

    # Find where the bomb hits
    # print(np.argmax(y < h[:,i]))
    # print(np.where(y < h[:,i]))
    # sys.exit(0)
    landing_point = np.argmax(y < h[:,i])

    if 0 < landing_point < Nx - 1:
      h[landing_point, i:] += dh
      if landsliding:
        slope_right = (h[landing_point, i] - h[landing_point + 1, i]) / dx
        slope_left = (h[landing_point, i] - h[landing_point - 1, i]) / dx

        gonna_move_r = slope_right > crit_slope
        gonna_move_l = slope_left > crit_slope

        if gonna_move_r and gonna_move_l:
            if random.random() < 0.5:
                gonna_move_r = False
            else:
                gonna_move_l = False

        while gonna_move_r and landing_point < Nx - 1:
            h[landing_point, i:] -= dh
            landing_point += 1
            h[landing_point, i:] += dh
            slope = (h[landing_point, i] - h[landing_point + 1, i]) / dx
            gonna_move_r = slope > crit_slope

        while gonna_move_l and landing_point > 0:
            h[landing_point, i:] -= dh
            landing_point -= 1
            if landing_point == 0:
                h[0, i:] += dh * (1 - phi)
            else:
                h[landing_point, i:] += dh
                slope = (h[landing_point, i] - h[landing_point - 1, i]) / dx
                gonna_move_l = slope > crit_slope

    elif landing_point == 0:
        h[0, i:] += dh * (1 - phi)

    elif landing_point == Nx:
        h[Nx, i:] += dh

    if i % i_plot == 0:
        print('Iteration: {}, h(x=0): {}'.format(i, h[0,i]))
        color = plt.cm.hsv(i / n_iterations)  # Use a colormap for changing
        #ax.plot(x, y, c='red', alpha=0.1)
        ax.fill_between(x, h[:,i-i_plot], h[:,i], interpolate=False, color=color)
        ax.fill_between(X, h[:,i-i_plot], h[:,i], interpolate=False, color=color)
        plt.pause(0.05)
        plt.draw()

plt.show()
# Clean up
plt.close()