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

from __future__ import division
from tridiagonal import solveLinearTridiagonalEquation
from memo import memoized
from cmath import exp, sqrt
import sys
sys.setrecursionlimit(10000)

@memoized
def d(minimum, maximum, resolution):
    return (maximum - minimum) / resolution
def psi0(x):
    pi = 3.141592653589793
    sigma = 2.0
    k0 = 1.0
    return  1.0/(sqrt(sigma*sqrt(pi)))*exp(-x**2/(2*sigma**2))*exp(1j*k0*x)
def V(x):
    if (10 < x and x < 14):
        return 10.0
    else:
        return 0.0
def normInline(dx, psixs):
    s = sum([(psi.conjugate()*psi).real*dx for psi in psixs])
    l = sqrt(s)
    for i in range(len(psixs)):
        psixs[i] /= l
    # fix rounding error
    psixs[0] = 0.0
    psixs[-1] = 0.0
    return(psixs)
def main():
    xminimum = -22.0
    xmaximum = 22.0
    xresolution = 140
    dx = d(xminimum, xmaximum, xresolution)
    dxsquare = dx*dx
    psixs = [psi0(xminimum + i*dx) for i in range(xresolution)]
    tminimum = 0.0
    tmaximum = 220.0
    tresolution = 220
    dt = d(tminimum, tmaximum, tresolution)
    tau = dt/dxsquare
    #print("%d\t%f\%20.5f" % (0, 10.0, 10.0))
    diagMain = [1] + [(1 + (1j*dt*(1/dxsquare + V(xminimum + i*dx)))) for i in range(1, xresolution - 1)] + [1]
    diagLower = [0] + [(-1/2*1j*tau) for i in range(1, xresolution - 1)] + [0] # dummy component first. Make sure to hard-code the boundary condition to stay as it was.
    diagUpper = [0] + [(-1/2*1j*tau) for i in range(1, xresolution - 1)] # make sure to hard-code the boundary condition to stay as it was.
    for tindex in range(tresolution):
        psixs = normInline(dx, psixs)
        # solve (1 + I dt ((1/2)∆ + V[x]1)) psi2 = psi for psi2
        for i in range(xresolution):
            print("%d\t%f\t%20.5f" % (i, xminimum + i*dx, (psixs[i].conjugate() * psixs[i]).real))
        print("")
        print("")
        psixs = solveLinearTridiagonalEquation(diagUpper, diagMain, diagLower, psixs)
if __name__ == "__main__":
    main()