"""
Part of take-home exam for FYS2140, spring of 2020.
This script plots real and imaginary part of  wavefunction
as well as probability density.

author: Sebastian G. Winther-Larsen
date: March 3, 2020
email: sebastwi@uio.no
"""

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation

hbarc = 197.4  # eVfm
m_p = 938.3  # MeV
m_n = 939.6  # Mev
mu = m_p * m_n / (m_p + m_n)  # MeV
c = 299_792_458  # m / s = fm / fs

def radial_wfn(r, E=-2.72, V0=35, R=2.1):
    """
    We use Heaviside functions in this function 
    to be able to preform a vectorized computation (=fast)
    """

    # Normalization factors
    A = 1
    C = 1.645

    k = np.sqrt(2 * mu * (E + V0) / hbarc ** 2)
    kappa = np.sqrt(-2 * mu * E / hbarc ** 2)

    radial_wfn = np.zeros_like(r)
    # For 0 < r < R
    radial_wfn += np.heaviside(R - r, 1) * A * np.sin(k * r)
    # For r > R
    radial_wfn += np.heaviside(r - R, 0) * C * np.exp(-kappa * r)

    return radial_wfn


def temporal_wfn(r, t, E=-2.72, V0=35, R=2.1):

    u = radial_wfn(r, E=E, V0=35, R=2.1)
    u /= r

    wfn = u * np.exp(-1j * E * t * c / hbarc)
    return wfn / np.linalg.norm(wfn)


r = np.linspace(0.1, 10, 10001)

for t in [0, 5, 10]:
    psi = temporal_wfn(r, t)
    psi_real = np.real(psi)
    plt.plot(
        r,
        psi_real,
        color="blue",
        alpha=(15 - t) / 15,
        label=f"Re, t={t}",
    )
    psi_imag = np.imag(psi)
    plt.plot(
        r,
        psi_imag,
        color="green",
        alpha=(15 - t) / 15,
        label=f"Im, t={t}",
    )

plt.plot(
    r,
    np.abs(temporal_wfn(r, 0)),
    color="black",
    linestyle="dashed",
    label=r"$|\Psi(x,t)|^2$",
)

plt.title("Wavefunction time development")
plt.xlabel(r"$r$ [fm]")
plt.ylabel(r"$\Psi$ or $|\Psi|^2$ (arb. units)")
plt.yticks([], [])
plt.legend()
plt.show()