import numpy as np
import matplotlib.pyplot as plt
import time, sys
from matplotlib import cm

"""
Simulation comments :
As we increase the number of grid points, the solution tends to a hat that
have move to the right. 
For nx<50, we have a bell not fully developed.
For nx=50, we have something ike a perfect bell.
For nx>50, we go back to the hat function.

For big dt, ~0.25,we have divergence of data.
Neverthe less, if we increase the timestep x2, we obtain a better solution
even if the number of grid points stays lower.

There seems to be some numerical diffusion problem, for high number of grid points 
the solution seems more like a transitory phenomena that diverges greatly !
For more information about this read about courant number.
"""

nx = 41  # Number of grid points.
ny = 41
nt = 50	# Number of timesteps we want to calculate.
c = 1 # Assume wavespeed of c = 1.
dx = 2 / (nx-1)	# Distance between any pair of adjacent grid points.
dy = 2 / (ny-1)	# Distance between any pair of adjacent grid points.
sigma = 0.2
dt = sigma*dx	# Amount of time each timesteps covers (delta t)

x = np.linspace(0,2,nx)
y = np.linspace(0,2,ny)

# Boundary conditions.
# Starting all values as 1 m/s.
u = np.ones((nx,ny))

# Setting u = 2 between 0.5 and 1 as per out initial conditions.
u[int(0.5 / dy):int(1 / dy + 1),int(0.5 / dx):int(1 / dx + 1)] = 2

# Generating two grids one for the values of x, another for the values of y. This will be needed for the plotting function.
X,Y = np.meshgrid(x,y)

# Plot initial conditions.
# Initialize the figure window. figsize and dpi are optional and specify size and resolution respectively.
fig, ax = plt.subplots(figsize=(11,7), dpi=100)
# Assigns the plot window  and the axes label 'ax' also specifies 3d projection plot.
ax = plt.subplot(1,2,1, projection='3d')
#ax = fig.gca(projection='3d')
# Equivalent to the regular plot command, but it takes grid of X,Y values for data point positions.
surf = ax.plot_surface(X, Y, u[:], cmap=cm.viridis)


# Temporary array for the solutions.
un = np.ones(nx)

for n in range(nt+1):   # Loop for values of n from 0 to nt.
    un = u.copy()	# Copy the existing values of u into un.
    row, col = u.shape
    for i in range(1, row): # Loop for values in u from 1 to nx.
        for j in range(1,col):
            u[j,i] = un[j,i] - c * dt / dx * (un[j,i] - un[j,i-1]) - c * dt / dx * (un[j,i] - un[j-1,i])
            
            # Boundary conditions
            u[0,:] = 1
            u[-1,:] = 1
            u[:,0] = 1
            u[:,-1] = 1

# Assigns the plot window  and the axes label 'ax' also specifies 3d projection plot.
ax = plt.subplot(1,2,2, projection='3d')
#ax = fig.gca(projection='3d')
# Equivalent to the regular plot command, but it takes grid of X,Y values for data point positions.
surf = ax.plot_surface(X, Y, u[:], cmap=cm.viridis)
# rotate the axes and update
for angle in range(0, 360):
    ax.view_init(30, angle)
    plt.draw()
    plt.pause(.001)