"""Puits quantique"""

import numpy as np
from numpy import loadtxt
from scipy import *

import matplotlib as mpl
import matplotlib.pyplot as plt

titre = """Puits quantique"""


#===========================================================
# --- Variables globales et paramètres ---------------------
#===========================================================

N = 1001                  # Discretisation du puits

Vo = 50               # Hauteur du puits quantique
L = 1                     # largeur du puits
b =L                      # Point en dehors du puits pour verifier si la fonction diverge
x = np.linspace(0, b, N)    # abscisses
dx = x[1] - x[0]

#===========================================================
# --- Réalisation du plot ----------------------------------
#===========================================================
def V(x, L=L, Vo=Vo):
    """
    Potentiel du puits quantique. L est la largeur du puits, et Vo la hauteur
    """
    if 0<x and x < L:
        return 0
    else:
        return Vo


listen=np.array([1,2,3,4, 5, 6])
En= listen**2
def psi_inf(n):
    return np.sin(n*pi/L*x)

# La fonction plot_data est appelée à chaque modification des paramètres
def plot_data(ax1, ax2, L=L, Vo=Vo):
    for E in En:
        ax2.plot(np.linspace(0, L, 50), E*np.ones(50), label="E = %.2f"%E)
    #ax2.legend()

    ## Fonctions d'onde
    for E,n in zip(En,listen):
        psi = psi_inf(n)
        norm2 = np.sum(psi**2*dx)
        psi = psi/np.sqrt(norm2)
        ax1.plot(x, psi+ E, label="E = %.2f"%E)


#===========================================================
# --- Création de la figure et mise en page ----------------
#===========================================================

fig = plt.figure()
#fig.suptitle(titre)
#fig.text(0.5, .93, description, multialignment='left', verticalalignment='top', horizontalalignment='center')

ax2, ax1 = fig.subplots(1, 2, sharex=True) # La figure sera composee de deux sous-figures

## ax2 : Energies
ax2.set_title(r'Énergies propres')
ax2.set_ylim(-1,40)
ax2.set_ylabel('$E_n$')
#ax2.set_xlim(-2,2)
ax2.set_xlabel('$x/L$')

## Dessin du puits
ax2.plot(x, np.vectorize(V)(x, L, Vo), linewidth=2, color='k')
ax1.plot(x, np.vectorize(V)(x, L, Vo), linewidth=2, color='k')
## ax1 : Fonctions d'onde
ax1.set_title("Fonctions d'onde propres")
ax1.set_ylim(-1,40)
ax1.set_ylabel('$\Psi_n(x)$')
ax1.set_xlabel('$x/L$')
ax1.yaxis.set_visible(False)
ax1.grid()
ax2.grid()



plot_data(ax1, ax2)

if __name__=="__main__":
    plt.show()