import numpy as np
import matplotlib.pyplot as plt
import time, sys
import sympy
from sympy.utilities.lambdify import lambdify


# Style and other non mathematical imports.
from sympy import init_printing
init_printing()
"""
Simulation comments :
"""

# Setting up symbolic variables.
x, nu, t = sympy.symbols('x nu t')
phi = (sympy.exp(-(x - 4 * t)**2 / (4 * nu * (t + 1))) +
       sympy.exp(-(x - 4 * t - 2 * sympy.pi)**2 / (4 * nu * (t + 1))))
# sympy.pprint(phi)

phiprime = phi.diff(x)
# sympy.pprint(phiprime)

u = -2 * nu * (phiprime / phi) + 4
# sympy.pprint(u)

# Lambdify : takes a sympy symbolic equation and turns it into a callable function
ufunc = lambdify((t, x, nu), u)


## Variable declaration for the physical problem.
# Physical parameters.
nu = 0.07

# Spatial parameters.
nx = 251 # Number of grid points.
dx = 2 * np.pi / (nx-1) # Distance between any pair of adjacent grid points.

#Time parameters.
nt = 250 # Number of timesteps we want to calculate.
sigma = 0.3 # Courant number.
dt = dx*nu  # Amount of time each timesteps covers (delta t)
#dt = sigma * dx**2 / nu # dt is defined using sigma, for convergence.

# Initialize time and space (creation of the universe.)
x = np.linspace(0, 2*np.pi, nx)
un = np.empty(nx)
t = 0

# Use lambdi(sci)fy function to start u variable.
u = np.asarray([ufunc(t, x0, nu) for x0 in x])

for n in range(nt):
    un = u.copy()
    for i in range(1, nx-1):
        u[i] = un[i] + nu * dt / (dx)**2 * (un[i+1]-2*un[i] + un[i-1]) - u[i] * dt / dx * (un[i] - un[i-1])
    
    u[0] = un[0] + nu * dt / (dx)**2 * (un[1]-2*un[0] + un[-2]) - u[0] * dt / dx * (un[0] - un[-2])
    u[-1] = u[0]

u_analytical = np.asarray([ufunc(nt*dt, xi, nu) for xi in x])

# Plotting and seeing :)
fig, ax = plt.subplots(nrows=1, ncols =2, sharey=True)

# Initial conditions.
ax[0].plot(x,u)

# Parameters for the plot.  
ax[0].set_xlim([0, 2 * np.pi])
ax[0].set_ylim([0, 10])

# Plotting the answer
ax[1].plot(x, u, label='Comp', marker='.', linestyle='None')
ax[1].plot(x,u_analytical, label='Anal')

# Parameters for the plot.
ax[1].set_xlim([0, 2 * np.pi])
ax[1].legend()

plt.show()


# A trabajar para que funcione bien sin sympy, défi !
'''
# Initial conditions.
# Starting all values as 1 m/s.
phi = []
for i in range(nx) :
    x = dx*i
    phi.append(np.exp(-x**2/(4 * vis)) + 
        np.exp(-(x-2*np.pi)**2/(4*vis)))

u = []
for i in range(0,nx-1) :
    print(i)
    u.append(-2 * vis * ((phi[i]-phi[i-1])/dx)/phi[i] + 4)

# Boundary conditions.
u.append(u[0])

# Analytical solution
phi_an = []
t = nt*dt
for i in range(nx) :
    x = dx*i
    phi_an.append(np.exp(-(x-4*t)**2/(4 * vis*(t+1))) + 
        np.exp(-(x-4*t-2*np.pi)**2/(4*vis*(t+1))))

u_an = []

for i in range(0,nx-1) :
    u_an.append(-2 * vis * ((phi_an[i]-phi_an[i-1])/dx)/phi_an[i] + 4)
    
u_an.append(u_an[0])
'''