___author___ = "Endre Jacobsen, Wilhelm A. Mangersnes"
import numpy as np
import matplotlib
from scipy import integrate
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
"""
CONSTANTS
"""
H_BAR = 1       # Planck's constant (Atomic units)
MASS = 1        # Mass of an electron (Atomic units)
K_0 = 20        # Wave number
L = 20          # Length of box
OMEGA = (H_BAR * K_0 ** 2)/(2*MASS)     # Wave speed
E = OMEGA * H_BAR       # Energy
V_G = H_BAR*K_0/MASS        # Group velocity
T = L/(2*V_G)       # Time it will take for the propagation
N_X = 2000       # Steps in x-space
N_T = 10000   # Number of iterations in time
DELTA_X = L/(N_X-1)     # x-step

'''For problem 2. Change these values.'''
X_S = 5  # Start position.

SIGMA_S = 1.3       # Sigma values to test: 0.5, 1.0, 1.3, 1.5, 2.0


def main():
    delta_t = (T/N_T)    # 1 time-step
    """
    Arrays
    """
    x_list = np.linspace(0, L, N_X)      # List containing all N_X steps in x-space
    pot_list = np.zeros(N_X)        # Empty array for storing potential values
    psi = np.zeros(N_X, dtype=complex)     # Complex array

    '''PROBLEMS. COMMENT IN/OUT AS NEEDED'''
    problem1(delta_t, x_list, psi)
    # problem2(delta_t, x_list, pot_list, psi)
    # problem3(delta_t, x_list, pot_list, psi)
    # problem4(delta_t)
    # problem5(delta_t)


def problem1(delta_t, x_list, psi):
    """
    Problem 1: Visualize the initial values of both the complex and real parts of Psi
    :param delta_t: Time-step
    :param x_list: List of length N_X with step-length DELTA_X
    :param psi: Empty complex array for Psi values
    :return: Nothing
    """
    init_psi(psi, delta_t)      # Initialize Psi
    prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate the probability density

    '''PLOT IMAGINARY AND REAL PARTS'''
    real_plot, = plt.plot(x_list, np.real(psi), label='Real')
    im_plot, = plt.plot(x_list, np.imag(psi), label='Imaginary')
    prob_plot, = plt.plot(x_list, prob_dens, label='Prob. density')
    plt.xlabel('x')
    plt.ylabel('Psi')
    plt.title('Initial wave-packet')
    plt.legend(handles=[real_plot, im_plot, prob_plot])
    plt.show()


def problem2(delta_t, x_list, pot_list, psi):
    """
    Problem 2: Propagate the wave a distance L/2
    :param delta_t: Timestep
    :param x_list: List of length N_X with step-length DELTA_X
    :param pot_list: Empty array for potential values
    :param psi: Empty complex array for Psi values
    :return: Noting
    """
    '''Initialize and plot Psi values (Problem 1)'''
    init_psi(psi, delta_t)      # Initialize psi
    prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate probability density of initial wave-packet
    in_psi, = plt.plot(x_list, psi, label='Initialized wave')       # Plot wave-packet
    in_prob, = plt.plot(x_list, prob_dens, label='Prob. density for initialized wave')     # Plot prob. density
    '''Propagate the wave packet'''
    psi = draw_psi(psi, pot_list, delta_t)      # Calls draw_psi() to propagate the wave
    prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate the prob. density of the new wave-packet
    new_psi, = plt.plot(x_list, psi, label='New wave')       # Plot new wave-packet
    new_prob, = plt.plot(x_list, prob_dens, label='Prob. density for new wave')     # Plot prob. density
    plt.xlabel('x')
    plt.ylabel('Psi')
    plt.title('Propagating a wave-packet, no barrier, sigma_x = 1.3')
    plt.legend(handles=[in_psi, in_prob, new_psi, new_prob])
    plt.show()


def problem3(delta_t, x_list, pot_list, psi):
    """
    Problem 3: Introduce a potential barrier at L/2 with height 0.5E and calculate the prob. density of the reflected
    and the transmitted wave-packet.
    :param delta_t: Timestep
    :param x_list: List of length N_X with step-length DELTA_X
    :param pot_list: Empty array for potential values
    :param psi: Empty complex array for Psi values
    :return: Nothing
    """
    '''Create barrier and propagate the wave past it. Plot results'''
    pot_list = make_barrier(pot_list, 0.5*E, L/50)       # Make a barrier at L/2
    init_psi(psi, delta_t)       # Initialize Psi
    psi = draw_psi(psi, pot_list, delta_t)      # Propagate the wave a distance L/2
    prob_dens = calc_prob_density(psi, 0, N_X)      # Calculate the prob. density of the resulting wave
    psi_plot, = plt.plot(x_list, psi, label='Wave-packet')       # Plot the resulting wave
    prob_plot, = plt.plot(x_list, prob_dens, label='Prob. density')     # Plot the prob. density
    v_plot, = plt.plot(x_list, pot_list / (2*E), label='Potential barrier')        # Plot the barrier. Scaled down with a factor 2*E
    plt.xlabel('x')
    plt.ylabel('Psi')
    plt.title('Wave propagated through a square barrier')
    plt.legend(handles=[psi_plot, prob_plot, v_plot])
    plt.show()

    '''Calculate the probability of reflection and transmission'''
    prob_t_dens_ = calc_prob_density(psi, int(np.ceil((L/2)/DELTA_X)), N_X)     # Prob. density of transmitted wave
    prob_t = ratio(prob_t_dens_) * 100      # Probability of transmission
    print('There is a', prob_t, '% chance of transmission.')

    prob_r_dens = calc_prob_density(psi, 0, int(np.floor((L / 2) / DELTA_X)))   # Prob. density of reflected wave
    prob_r = ratio(prob_r_dens) * 100       # Probability of reflection
    print('There is a', prob_r, '% chance of reflection.')

    print('prob_t + prob_r =', prob_t + prob_r, '% . This should be close to 100%.')


def problem4(delta_t):
    """
    Problem 4: Calculate the probabilities of reflection and transmission with 50 different barrier heights
    :param delta_t: Time-step
    :return: Nothing
    """
    v_max = (3 / 2) * E     # Maximum value of the
    total_barriers = 50     # Number of barriers to be tested
    barrier_heights = np.linspace(0, v_max, total_barriers)         # Create an array with all the barrier heights up to v_max

    prob_t_list = np.zeros(total_barriers)      # Empty array to hold the probability of transmission at each barrier height
    prob_r_list = np.zeros(total_barriers)      # Empty array to hold the probability of reflection at each barrier height

    start_psi = init_psi(np.zeros(N_X, dtype=complex), delta_t)  # Initialize Psi

    for i in range(len(barrier_heights)):
        psi = np.copy(start_psi)        # Reset psi for every new barrier
        pot_list = make_barrier(np.zeros(N_X), barrier_heights[i], L/50)        # Make the barrier
        psi = draw_psi(psi, pot_list, delta_t)      # Propagate the wave
        prob_dens_trans = calc_prob_density(psi, int(np.ceil((L / 2) / DELTA_X)), N_X)      # Calculate the prob. density for transmission
        prob_t_list[i] = ratio(prob_dens_trans)     # Calculate the probability for transmission
        prob_r_list[i] = 1 - prob_t_list[i]         # Calculate the probability for reflection

    '''Plot the probabilities'''
    r_plot, = plt.plot(barrier_heights, prob_r_list, label='Probability for reflection', color="blue")
    t_plot, = plt.plot(barrier_heights, prob_t_list, label='Probability for transmission', color="gold")
    plt.xlabel('Potential')
    plt.ylabel('Probability')
    plt.title('P(Reflected) vs. P(Transmitted)')
    plt.legend(handles=[r_plot, t_plot])
    plt.show()


def problem5(delta_t):
    """
    Problem 5: Calculate the probabilities of reflection and transmission with 50 different barrier widths
    :param delta_t: Time-step
    :return: Nothing
    """
    v_max = (9/10)*E        # Height of the barrier
    max_width = L/20        # Maximum width of the barrier
    total_barriers = 50     # Number of barriers to test

    barrier_widths = np.linspace(0, max_width, total_barriers)      # Create an array of barrier widths

    prob_t_list = np.zeros(total_barriers)      # Empty array to hold the probability of transmission at each barrier height
    prob_r_list = np.zeros(total_barriers)      # Empty array to hold the probability of reflection at each barrier height

    start_psi = init_psi(np.zeros(N_X, dtype=complex), delta_t)  # Initialize Psi

    for i in range(len(barrier_widths)):
        psi = np.copy(start_psi)        # Reset psi for every new barrier width
        pot_list = make_barrier(np.zeros(N_X), v_max, barrier_widths[i])        # Create a new barrier
        psi = draw_psi(psi, pot_list, delta_t)      # Propagate the wave-packet
        prob_t_dens = calc_prob_density(psi, int(np.ceil((L/2)/DELTA_X)), N_X)      # Calculate the prob. density for transmission
        prob_t_list[i] = ratio(prob_t_dens)     # Calculate the probability for transmission
        prob_r_list[i] = 1 - prob_t_list[i]     # Calculate the probability for reflection

    '''Plot the probabilities'''
    r_plot, = plt.plot(barrier_widths, prob_r_list, label='Probability for reflection', color="blue")
    t_plot, = plt.plot(barrier_widths, prob_t_list, label='Probability for transmission', color ="gold")
    plt.xlabel('Barrier widths')
    plt.ylabel('Probability')
    plt.title('P(Reflected) vs. P(Transmitted)')
    plt.legend(handles=[r_plot, t_plot])
    plt.show()


def draw_psi(psi, pot_list, delta_t):
    """
    Function for propagating the wave-packet contained in psi
    :param psi: A complex array containing the wave-packet
    :param pot_list: Potential
    :param delta_t: Time-step
    :return: The propagated wave-packet psi
    """
    for t in range(1, N_T):
        psi = calc_psi_i(psi, pot_list, delta_t)        # Calculate Im(Psi(x,t))
        psi = calc_psi_r(psi, pot_list, delta_t)        # Calculate Re(Psi(x,t+delta_t/2
    return psi


def init_psi(psi, delta_t):
    """
    Initializes Psi at time zero.
    :param psi:     Wave equation
    :param delta_t:     Time-step
    :return:        Initialized array of wave equation values
    """
    c = calc_c()        # Calculate normalisation constant
    for x in range(1, N_X-1):
        '''
        Equation (8)
        t = 0 for Psi_I and t = delta_t/2 for Psi_R
        '''
        psi[x] = complex(c[0] * np.exp(calc_alpha(x*DELTA_X)) * np.cos(K_0 * x*DELTA_X - OMEGA*(delta_t/2)),
                         c[0] * np.exp(calc_alpha(x*DELTA_X)) * np.sin(K_0 * x*DELTA_X))
    return psi


def calc_alpha(x):
    """
    Calculates alpha (the first exponent) from equation (8).
    :param x:   Position
    :return:    Alpha
    """
    alpha = -((x-X_S) ** 2)/(2*(SIGMA_S ** 2))
    return alpha


def calc_prob_density(psi, start, stop):
    """
    Calculate the probability density for a complex array
    :param psi: Complex array
    :param start: Where to start the calculation
    :param stop: Where to stop the calculation
    :return: Prob list, a list of probability densities
    """
    prob_list = np.zeros(stop-start)        # Create empty array containing number of elements to be calculated
    for i in range(stop-start):
        real = np.real(psi[i+start])        # Get the real value of the complex number
        imaginary = np.imag(psi[i+start])   # Get the imaginary value of the complex number
        prob_list[i] = real**2 + imaginary ** 2     # Calculate the probability density
    return prob_list


def calc_psi_i(psi, pot_list, delta_t):
    """ Calculate Psi at position x and time t.
    :param psi:    The list of wave equation values
    :param pot_list:    The list containing the potential values
    :param delta_t:     Time-step
    :return:    New value of psi(x,t)
    """
    '''Equation (7a)'''
    val_i = np.imag(psi[1:N_X-1]) - delta_t * ((pot_list[1:N_X-1]/H_BAR) * np.real(psi[1:N_X-1]) - H_BAR/(2*MASS) *
                                               (np.real(psi[2:N_X]) - 2*np.real(psi[1:N_X-1]) +
                                                np.real(psi[0:N_X-2]))/(DELTA_X ** 2))

    psi[1:N_X-1] = np.real(psi[1:N_X-1]) + val_i*1j
    return psi


def calc_psi_r(psi, pot_list, delta_t):
    """ Calculate Psi at position x and time t. The imaginary part is calculated at t + delta_t/2
    :param psi: The list of wave equation values
    :param pot_list:    Potential values
    :param delta_t:     Time-step
    :return: The new value of psi
    """
    '''Equation (7b)'''
    val_r = np.real(psi[1:N_X-1]) + delta_t *((pot_list[1:N_X-1]/H_BAR) * np.imag(psi[1:N_X-1]) - H_BAR/(2*MASS) *
                                              (np.imag(psi[2:N_X]) - 2*np.imag(psi[1:N_X-1]) +
                                               np.imag(psi[0:N_X-2]))/(DELTA_X ** 2))

    psi[1:N_X-1] = val_r + np.imag(psi[1:N_X-1])*1j
    return psi


def calc_c():
    """
    Calculate the normalisation constant for equation (8)
    :return: Normalisation constant
    """
    def normalize_function():
        return lambda x: np.exp(-((x-X_S)**2)/SIGMA_S**2)
    expr = normalize_function()
    ans = integrate.quadrature(expr, 0, L, args=(), tol=1.49e-08)
    c = 1/np.sqrt(ans)
    return c


def make_barrier(pot_list, height, width):
    """
    Create a barrier of potential value 'height'
    :param pot_list: The empty potential values
    :param height: The height that the barrier is to be set to
    :param width: The width of the potential barrier
    :return: The potential (barrier)
    """
    pot_list[int(np.floor((L/2 - width/2)/DELTA_X)):int(np.ceil((L/2 + width/2)/DELTA_X))] = height
    return pot_list


def ratio(prob_dens):
    """
    Calculate the probability for a given prob. density
    :param prob_dens: Array of probability density
    :return: The probability for the given prob. density
    """
    return np.sum(prob_dens) * DELTA_X      # With DELTA_X small enough this is approximately the same as integrating.


main()