from enum import Enum
import math
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from math import cos, sin, exp
import cmath

planc_const = 1e1
class Grid:
    h = 0
    tau = 0
    l = 0

    def __init__(self, h, tau, l):
        self.h, self.tau, self.l = h, tau, l

    def getH(self):
        return self.h

    def getL(self):
        return self.l

    def getTau(self):
        return self.tau

    def getStepCount(self):
        return math.ceil(self.l / self.h) + 1


class ConditionMode(Enum):
    NormalFunction = 0,
    DeriviativeFunction = 1


class Condition:
    def __init__(self, Condition, mode=ConditionMode):
        self.Condition, self.mode = Condition, mode


class Task:
    def __init__(self, l, a, StartCondition, BoundaryConditions):
        assert isinstance(StartCondition, Condition)
        self.l, self.a = l, a
        self.StartCondition, self.BoundaryConditions = StartCondition, BoundaryConditions

    def getL(self):
        return self.l

    def getA(self):
        return self.a

    def getStartConditions(self):
        return self.StartCondition

    def getBoundaryConditions(self):
        return self.BoundaryConditions


def buildGenerator(condition, stepCount, step):
    assert isinstance(condition, Condition)
    prevElement = 0
    for i in range(stepCount):
        if (condition.mode == ConditionMode.DeriviativeFunction):
            currentElement = prevElement + step * condition.Condition(i * step)
            yield currentElement
            prevElement = currentElement
        else:
            yield condition.Condition(i * step)


def primeCalc(prevU, h, tau, a, BoundaryConditions, PotentialFunc, tick=0):
    ##u(x,t+1)-u(x,t)/tau = u(x-1,t) -2u(x,t) + u(x+1,t)/(h^2) + PotentialFunc
    ##assert a * tau / (h * h) < 0.5
    currentU = [0] * len(prevU)
    for j in range(1, len(prevU) - 1):
        currentU[j] = prevU[j] + a * tau * (prevU[j - 1] - 2 * prevU[j] + prevU[j + 1]) / (h * h) + tau  * PotentialFunc(j * h, tick * tau) * prevU[j]

    # u0-u1/h = u`(0)
    if (BoundaryConditions[0].mode == ConditionMode.DeriviativeFunction):
        currentU[0] = currentU[1] + h * BoundaryConditions[0].Condition(tick * tau)
    else:
        currentU[0] = BoundaryConditions[0].Condition(tick * tau)
    sc = len(prevU)


    # u(l-1)-ul/h=u`(l)
    if (BoundaryConditions[1].mode == ConditionMode.DeriviativeFunction):
        currentU[sc - 1] = currentU[sc - 2] - h * BoundaryConditions[1].Condition(tick * tau)
    else:
        currentU[sc - 1] = BoundaryConditions[1].Condition(tick * tau)
    return currentU


def Solve(task, grid, iterationCount, recalcFunc, PotentialFunc):
    assert isinstance(task, Task)
    assert (task.getL() == grid.getL())

    prevU = [x for x in buildGenerator(task.getStartConditions(), grid.getStepCount(), grid.getH())]

    xx = []
    for i in range(grid.getStepCount()):
        xx.append(i * grid.getH())

    currentU = prevU
    for i in range(iterationCount):
        #print(currentU)
        yield xx, currentU
        prevU = [x for x in currentU]
        h = grid.getH()
        tau = grid.getTau()
        a = task.getA()
        currentU = recalcFunc(prevU, h, tau, a, task.getBoundaryConditions(), PotentialFunc, i)



def shuttleCalc(prevU, h, tau, a, BoundaryConditions, PotentialFunc, tick=0):
    # ux,t+1 - ux,t/tau = a*(ux-1,t+1 -2ux,t+1 + ux+1,t+1)/h**2 + PotentialFunc(r, t) * u(x,t+1)
    # -a/h**2(ux-1,t+1 + ux+1,t+1) + (1/tau +2a/h**2) * ux,t+1 = 1/tau * ux,t
    #
    ## -a/h**2 (1/tau +2a/h**2) -a/h**2
    gamma = a * tau / (h ** 2)
    a_orig = -gamma
    b_orig = (1 + 2 * gamma)
    c_orig = a_orig
    n = len(prevU)

    currentU = [0] * n
    ##i*h*psi^ = Hpsi ==
    ##i*h * psi^ = -h^2/2*m d2psi/dx2
    #psi^=-h/2mi
    ##psi^ = -h/i^3 d2psi/dx2

    e = [a_orig] * n
    c = [c_orig] * n
    d = [b_orig] * n

    ## u(x, t + 1) - u(x, t)/tau = a * (ux-1,t+1 - 2u(x,t+1) + u(x+1,t+1)/h**2  + p(r,t)*u(x,t +1)
    ## -gamma * u(x-1,t+1) -gamma * u(x + 1, t + 1) + (1 + 2gamma - p(r,t) * tau))*u(x,t + 1) = u(x,tau)
    for i in range(n):
        d[i] -= tau  * complex(0, 1)  / planc_const * PotentialFunc(i * h, tick * tau)
        #assert(abs(1 - tau * complex(0, 1) * PotentialFunc(i * h, tick * tau)) > 0)

    ##e[0] = BoundaryConditions[0].Condition(tick * tau)
    ##e[n - 1] = BoundaryConditions[1].Condition(tick * tau)
    #psi1 = alpha * u(0, t + 1) + beta * ( u(0, t + 1) - u(1, t + 1)) /h
    #psi2 = alpha * u(l, t + 1) + beta * ( u(l, t + 1) - u(l - 1, t + 1)) /h

    c[0] = complex(0, 0)
    if (BoundaryConditions[0].mode == ConditionMode.DeriviativeFunction):
        d[0] = complex(1.0 / h, 0)
        e[0] = complex(-1.0 / h, 0)
    else:
        d[0] = complex(1.0, 0)
        e[0] = complex(0.0, 0)

    e[n - 1] = complex(0, 0)
    if (BoundaryConditions[1].mode == ConditionMode.DeriviativeFunction):
        d[n - 1] = complex(-1.0 / h, 0)
        c[n - 1] = complex(1.0 / h, 0)
    else:
        d[n - 1] = complex(1.0, 0)
        c[n - 1] = complex(0.0, 0)

    prevU[0] = BoundaryConditions[0].Condition(tick * tau)
    prevU[n - 1] = BoundaryConditions[1].Condition(tick * tau)

    # a_c[1] = -e[0] / d[0]
    # b_c[1] = prevU[0] / d[0]
    # for i in range(2, n):
    #     a_c[i] = -e[i] / (d[i] + c[i] * a_c[i - 1])
    #     b_c[i] = (-c[i] * b_c[i - 1] + prevU[i - 1]) / (d[i] + c[i] * a_c[i - 1])
    #
    # currentU[n - 1] = (-c[n - 1] * b_c[n - 1] + prevU[n - 1]) / (d[n - 1] + c[n - 1] * a_c[n - 1])
    # for i in range(n - 2, -1, -1):
    #     currentU[i] = a_c[i + 1] * currentU[i + 1] + b_c[i + 1]

    alpha = [0] * n
    beta = [0] * n
    ##CDE
    alpha[1] = -e[0]/ d[0]
    beta[1] = prevU[0] / d[0]
    for i in range(2, n):
        alpha[i] = (-e[i - 1])/(c[i - 1] * alpha[i - 1] + d[i - 1])
        beta[i] = (prevU[i - 1] - c[i - 1] * beta[i - 1]) / (c[i - 1] * alpha[i - 1] + d[i - 1])

    currentU[n - 1] = (prevU[n - 1] - c[n - 1] * beta[n - 1]) / (c[n - 1] * alpha[n - 1] + d[n - 1])
    for i in range(n - 2, -1, -1):
        currentU[i] = alpha[i + 1] * currentU[i + 1] + beta[i + 1]

    b2 = prevU[2]
    b1 = prevU[1]

    g = gamma
    x1 = (b1 + g * b2) / (1 + 2 * g + g*g/(1+2 *g))
    x2=(b2 + g * x1)/(1 + 2 *g)


    return currentU



## -h^2/2m * delta + Ep = ih|psi>/dt
##  h/2m  -i/h Ep = |psi>/dt


#def osscilator
#

#-h^2/2m delta  + 0.5 m * omega *
mass = 1.5*1e3
#a = -planc_const * planc_const / (2 * mass)
#
a = complex(0, 1) * planc_const / (2 * mass)


def origSolve(task, grid, iter, func, ep):
    prevU = [x for x in buildGenerator(task.getStartConditions(), grid.getStepCount(), grid.getH())]

    xx = []
    for i in range(grid.getStepCount()):
        xx.append(i * grid.getH())

    currentU = prevU
    for i in range(iter):
        #print(currentU)
        yield xx, currentU
        prevU = [x for x in currentU]
        h = grid.getH()
        tau = grid.getTau()
        a = task.getA()
        for p in range(len(currentU)):
            currentU[p] = cmath.exp(-math.pi * math.pi *  task.getA() *  i * tau) * sin(math.pi * p * grid.getH())





omega = math.pi  / 16
def Ep(r, tau):
   #return (0.5 * mass * r * r * omega) * complex(0, 1) / planc_const
   return complex(0, 0)
#def Ep(r, tau):
#    return complex(0, -1 / planc_const)

l = 1

def startCondition(x):

    return 0.5*complex(exp(-(x+1.5) * (x+1.5) / 2.0), 0)
    #return complex(sin(3 * math.pi * x / l), 0)
    if (abs(x) < eps):
        return complex(0.05 * exp(-15 *(x - 0.5) * (x - 0.5)), 0)
    else:
        return complex(0, 0)


def boundaryCondition0(x):
    #return exp(-a * math.pi * math.pi * x)
    return complex(0, 0)


def boundaryConditionL(x):
    #return -exp(-a * math.pi * math.pi * x)
    return complex(0, 0)

maxPlots=3
gxmin, gxmax = 0, 1
gymin, gymax = -1, 1
fig, (axRe, axIm, axAbs) = plt.subplots(3, sharex=True, sharey= True)
axRe.text("R")
lineRe, = axRe.plot([], [], lw=2)
lineIm, = axIm.plot([], [], lw=2)
lineAbs, = axAbs.plot([], [], lw=2)

for i in [axRe, axIm, axAbs]:
    i.set_ylim(gymin, gymax)
    i.set_xlim(gxmin, gxmax)


rescale = False

def updateAx(xdata, ydata, ax, line):
    xmin, xmax = ax.get_xlim()
    if rescale:
        if xdata[-1] >= xmax:
            ax.set_xlim(xmin, 2 * xmax)
            ax.figure.canvas.draw()

        ymin, ymax = ax.get_ylim()


        curYmin = ydata[0]
        curYmax = ydata[0]

        for i in range(len(ydata)):
            curYmax = max(curYmax, ydata[i])
            curYmin = min(curYmin, ydata[i])
        if curYmax >= ymax or curYmin <= ymin:
            if (curYmin < 0):
                curYmin *= 2
            else:
                curYmin /= 2

            if (curYmax < 0):
                curYmax /= 2
            else:
                curYmax *= 2

            ax.set_ylim(curYmin, curYmax)
            ax.figure.canvas.draw()

    line.set_data(xdata, ydata)
    #ax.figure.canvas.draw()



def run(data):
    # update the data
    t, y = data
    xdata = t
    yRe = [el.real for el in y]
    yIm = [el.imag for el in y]
    yAbs = [abs(el) for el in y]

    updateAx(xdata, yRe, axRe, lineRe)
    updateAx(xdata, yIm, axIm, lineIm)
    updateAx(xdata, yAbs, axAbs, lineAbs)



def potentialFunc(r, t):
    return (r+0.5)*(r+ 0.5) * 10
    #return complex(0, 0)/planc_const
    #return cmath.exp(complex(-vr * vr)/2.0)
    return 0

def Re(z):
    return z.real

def Im(z):
    return z.imag

def Abs(z):
    return abs(z)


def main():
    startConditionWrapper = Condition(startCondition, ConditionMode.NormalFunction)
    boundCondition0Wrapper = Condition(boundaryCondition0, ConditionMode.NormalFunction)
    boundConditionLWrapper = Condition(boundaryConditionL, ConditionMode.NormalFunction)

    l = 1
    h = 0.01
    #h(
    #tau = h * h / 4 *a
    tau = 0.08
    grid = Grid(h, tau, l)

    stCond = [startCondition(x * h) for x in range(0, grid.getStepCount())]
    f = [Re, Im, Abs]
    axA = [axRe, axIm,axAbs]
    for i in range(len(axA)):
        maxHC = max(f[i](x) for x in stCond)
        axA[i].set_ylim(-1.2 * maxHC, 1.2 * maxHC)
        axA[i].set_xlim(0, l)



    ## if we use primeCalc (a * tau) /h*h should be < 1/2
    task = Task(l, a, startConditionWrapper, [boundCondition0Wrapper, boundConditionLWrapper])

    TimePeriod = 10

    iter  = 5000

    #ani = animation.FuncAnimation(fig, run, Solve(task, grid, iter, shuttleCalc, Ep), blit=False, interval=TimePeriod,
    #                              repeat=False)

    ani = animation.FuncAnimation(fig, run, Solve(task, grid, iter, shuttleCalc, potentialFunc), blit=False, interval=TimePeriod,
                                  repeat=True)


    plt.show()


if __name__ == "__main__":
    main()