#!/usr/bin/env python
"""
Created on Sun 3 Mar 2019/Magne Guttormsen/Fabio Zeiser
"""
import numpy as np
from numpy import sqrt, sin, cos, exp, pi
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import sys
import scipy.integrate as integrate

# some constants

mc2 = 939.56563     # MeV
hbarc = 197.3       #MeV fm
a = 1.2 * 235**(1/3.)  # fm
V0 = 10.            # MeV


def square_well(E):
    """ Potential and solutions to the finite square well
    Args:
        E : Energy in MeV

    Returns:
        psi (function): wave function in [1/fm^0.5]
        V (function): potential in MeV
    """
    k  = sqrt(2*mc2*E) / hbarc
    l  = sqrt(2*mc2*(E+V0)) / hbarc

    A = 1.
    F = exp(-2j*k*a) * A \
        / (cos(2*l*a) - 1j*sin(2*l*a)/(2*k*l) * (k**2+l**2))
    B = 1j*sin(2*l*a)/(2*k*l) * (l**2-k**2) * F
    C = (sin(l*a) + 1j*k/l*cos(l*a)) * exp(1j*k*a) * F
    D = (cos(l*a) - 1j*k/l*sin(l*a)) * exp(1j*k*a) * F

    # print a
    # print A
    # print B
    # print C
    # print D
    # print F
    Trans = np.abs(F)**2 / np.abs(A)**2

    # Definer boks potensialet
    def boks1(x):
        return 0.0
    def boks2(x):
        return 0.0
    def boks3(x):
        return -V0
    def V(x):
        conds = [x<-a, x>a, (x>-a) & (x<a)]
        funcs = [boks1, boks2, boks3]
        return np.piecewise(x, conds, funcs)

    def psi1(x):
        return A*exp(1j*k*x) + B*exp(-1j*k*x)
    def psi2(x):
        return F*exp(1j*k*x)
    def psi3(x):
        return C*sin(l*x) + D*cos(l*x)
    def psi(x):
        conds = [x<-a, x>a, (x>-a) & (x<a)]
        funcs = [psi1, psi2, psi3]
        return np.piecewise(x+0j, conds, funcs)

    # def Asq1(x):
    #     return A*exp(1j*k*x)
    # def Asq2(x):
    #     return 0.
    # def Asq3(x):
    #     return 0.
    # def Asq(x):
    #     conds = [x<-a, x>a, (x>-a) & (x<a)]
    #     funcs = [Asq1, Asq2, Asq3]
    #     return np.piecewise(x+0j, conds, funcs)

    return V, psi, Trans


E_arr= np.linspace(0.1,20)
T = np.zeros(len(E_arr))
for i, E in enumerate(E_arr):
    V, psi, Trans = square_well(E)
    T[i] = Trans


x = np.linspace(-30, 30, num=500) # x grid for plotting
fig, (ax1, ax2) = plt.subplots(2,1, sharex=True, sharey=True)

xmax = 30
xmin = -xmax

ax1.set_xlim(xmin, xmax)
ax1.set_ylim(-2, 2.5)

def make_plots(ax, E):
    V, psi, Trans = square_well(E)
    def norm2(x):
        return np.abs(psi(x))**2
    integral1 = integrate.quad(norm2,xmin,a)
    integral2 = integrate.quad(norm2,a,xmax)
    integral3 = integrate.quad(norm2,-a,a)
    # print(integral3[0]/(integral1[0]+integral2[0]))
    print(integrate.quad(norm2,xmin,xmax))

    ax.set_xlabel('$x$ [fm]')
    ax.set_ylabel('$\Psi (x, t)$ [1/fm]$^{1/2}$')
    ax.text(13,-1.6, "E = {:.1f} MeV".format(E))
    ax.set_ylim(-3,3)
    ticks = ax.get_xticks()


    ax2 = ax.twinx()
    ax2.set_ylim(-3,3)
    ax2.set_ylabel('$|\Psi (x, t)|^2$ [1/fm]', color="g")
    locator=MaxNLocator(nbins=3)
    ax.yaxis.set_major_locator(locator)
    ax2.yaxis.set_major_locator(locator)

    #  plot
    ax.plot( x, V(x),'r',alpha=0.7, label='V(x)')
    ax.plot( x, psi(x).real,'k',alpha=0.7,label='Real' )
    ax.plot( x, psi(x).imag,'b',alpha=0.7,label='Imaginary' )
    ax.plot(x,np.abs(psi(x))**2,'g',label='Probability density')
    #Asq, = plt.plot(x,np.abs(Asq(x))**2,'k',label='$|A|^2$')


make_plots(ax1, E=0.5)
make_plots(ax2, E=3)
ax1.legend(loc='upper right')

# Fine-tune figure; make subplots close to each other and hide x ticks for
# all but bottom plot.
fig.subplots_adjust(hspace=0)
# plt.setp([a.get_xticklabels() for a in fig.axes[:-1]], visible=False)

# plt.figure()
# plt.plot(E_arr,1-T)

plt.show()