#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Flore Nabet et Thomas Wick
# Ecole Polytechnique
# MAP 411 de Grégoire Allaire
# Hiver 2016/2017

# Execution of *.py files
# Possiblity 1: in terminal
#  terminal> python3 file.py  # here file.py = convection.py

# Possiblity 2: executing in an python environment
# such as spyder.

# Resolution numerique de l'equation de la chaleur
# d_t u - u''(x) = 0

########################################################
# Load packages (need to be installed first if
# not yet done - but is not difficult)
import numpy as np
import matplotlib.pyplot as plt # pour plot functons
plt.switch_backend('tkAgg')  # necessary for OS SUSE 13.1 version,
# otherwise, the plt.show() function will not display any window

import pylab
import scipy as sp
from scipy import sparse
from scipy.sparse import linalg


########################################################
# Parametres du probleme
lg = 10        # intervalle en x=[-lg,lg]
dx = 0.1       # dx = pas d'espace
cfl = 0.55     # cfl = dt/dx^2
dt = dx*dx*cfl # dt = pas de temps
Tfinal = 0.5   # Temps final souhaite


########################################################
# Initialisation

nx = 2*lg/dx + 1 # nx = nombre d'intervals
x = np.linspace(-lg,lg,nx)

#print(int(nx))

# Initialize u0
u0 = np.zeros(len(x))
#print(len(u0))

# Set specific u0 values (same as in the scilab program)
for k in range (len(x)):
    if (1.0 - x[k]**2) < 0:
       u0[k] = 0
    else:
       u0[k] = 1.0 - x[k]**2 # donnee initiale


# Sanity check if everything works correctly
#print (x)
#print (u0)

# Plot initial condition
plt.ion()
plt.plot(x,u0,'o')
#plt.show()

# We also write the intial condition into a file
# of *.png or *.pdf format
fig, ax = plt.subplots(nrows=1,ncols=1)
ax.plot(x,u0) # trace u0 en fonction de x
fig.savefig('u0.png')
plt.close(fig)


########################################################
# Schemas numeriques

# Initialize u by the initial data u0
uimp = u0.copy() # il faut faire une copie sinon on va changer u0 en changeant u
uexp = u0.copy()

# Construction de la matrice (creuse) pour le schema explicite
# u(-10,t)=u(10,t)=0 on peut donc enlever ces deux coordonnees du systeme
#u_j^{n+1}=nu*cfl*u_{j-1}^n+(1-2*nu*cfl)u_j^n+nu*cfl*u_{j+1}^n
Aexp=sp.sparse.diags([cfl, 1-2*cfl, cfl], [-1, 0, 1], shape=(nx-2, nx-2))

# Construction de la matrice (creuse) pour le schema implicite
#-nu*cfl*u_{j-1}^{n+1}+(1+2*nu*cfl)u_j^{n+1}-nu*cfl*u_{j+1}^{n+1}=u_j^n
Aimp=sp.sparse.diags([-cfl, 1+2*cfl, -cfl], [-1, 0, 1], shape=(nx, nx))

# Sanity check to see how matrix looks like
#print(Aexp.toarray())
#print(Aimp.toarray())


# Nombre de pas de temps effectues
nt = int(Tfinal/dt)
Tfinal = nt*dt # on corrige le temps final (si Tfinal/dt n'est pas entier)


# Time loop
for n in range(1,nt+1):

    # Solution 1: Solve explicit system
    uexp[1:len(uexp)-1]=Aexp*uexp[1:len(uexp)-1]

    # Solution 2: schema implicite
    uimp = sp.sparse.linalg.spsolve(Aimp,uimp)


   # Print solution
    if n%5 == 0:
        plt.figure(2)
        plt.clf()

        plt.subplot(121)
        plt.plot(x,u0,'b',x,uexp,'r')
        plt.xlabel('$x$')
        plt.title('Schema explicite')

        plt.subplot(122)
        plt.plot(x,u0,'b',x,uimp,'m')
        plt.xlabel('$x$')
        plt.title('Schema implicite')

        plt.pause(0.1)
 #       plt.show()

        # Print solution into seperate files
#        fig, ax = plt.subplots(nrows=1,ncols=1)
#        ax.plot(x,uexp,'b',x,uimp,'r') # trace uexp et uimp en fonction de x
#        # On peux utiliser *.png ou *.pdf
#        fig.savefig(str(n) + 'u.png')
#        plt.close(fig)

####################################################################
# comparaison solution exacte avec solution numerique au temps final

uexacte = np.zeros(len(u0))

def noyauc(x,t):
    return np.exp(-x**2/(4*t))/np.sqrt(4*np.pi*t)

# on calcule la solution exacte en utilisant la formule des rectangles
# uexacte(xi,Tfinal)=\int u0(y) noyauc(xi-y,Tfinal)
# uexacte(xi,Tfinal)~sum_{j=0}^{2*lg/dx} dx*u0(xj) noyauc(xi-xj,Tfinal)
# et xi=-lg+i*dx donc xi-xj=(i-j)*dx
for i in range(int(nx)):
    for j in range(int(nx)-1):
        uexacte[i] = uexacte[i] + u0[j]*dx*noyauc((i-j)*dx,Tfinal)


plt.figure(3)
plt.clf()
plt.plot(x,u0,'b',x,uexp,'or',x,uimp,'vm',x,uexacte,'k')
plt.legend(['Donnee initiale','Schema explicite','Schema implicite','Solution exacte'],loc='best')
plt.title('Comparaison entre la solution exacte et la solution numerique au temps final')