%matplotlib inline
%load_ext snakeviz

import numpy as np
from scipy.integrate import odeint
import matplotlib.pylab as plt

#import sys
#sys.path.insert(0, '../../module')
import neuosc as no

"""
Explanation: Understanding Periodic Matter Driven Neutrino Oscillations
blablablabla
I have an analytical expression using two folds of approximations.
To check out how good the expression is doing, I need to verify at three levels.

Compare the solution to the equation I have after the first RWA approximation and the final result;
Compare the solution to the equation I have after the first RWA approximation and the solution to the original equation;
IF the solution works well in the previous two comparisons, compare the solution to the original equation and the final analytical expression.

Before doing anything, check if it is possible to find an analytical solution to the RWA approximated equation!
Prep
End of explanation
"""


beta = 0.1 # beta = \omega_\lambda/\omega
alpha = 0.1 # alpha = lambda_0 /\omega
sin2thetav = no.mixing2()[0,1]
cos2thetav = no.mixing2()[0,0]

sin2thetav**2+cos2thetav**2

def prob3(alp,bet,x):
    
    # alp = alpha, bet = beta
    
    aOverOmega = (alp * sin2thetav / 2 )
    deltaOverOmega = 1.0 - bet - 0.5*alp*cos2thetav*np.cos(bet*x)
    
    omegaRHat = np.sqrt(deltaOverOmega**2 + aOverOmega**2 )
    omegaRHatsquare = deltaOverOmega**2 + aOverOmega**2 
    
    return (alp*sin2thetav)**2*(np.sin(omegaRHat*x/2))**2/(4*(omegaRHatsquare**2))
    

prob3(alpha,beta,0.5)

pltend = 800
pltlinsize = pltend*10 + 1
x = np.linspace(0, pltend, pltlinsize)
plt.figure(figsize = (16,9.36))
plt.plot(x, prob3(alpha,beta,x))
plt.xlabel('$\omega x$')
plt.ylabel('$|C_2|^2$')
plt.axis('tight')
plt.title('Transition Probility for $\lambda_0/\lambda = \omega_\lambda /\omega = 0.1$"')
plt.show()

"""
Explanation: Analytical Expressioin
The final approximated result is
End of explanation
"""


def kMatrix(x):
    k0 = alpha*np.cos(beta*x)/2
    k11, k22, k33, k44 = 0,0,0,0
    k12, k21, k34, k43 = 1.0*cos2thetav, -1.0*cos2thetav, -1.0*cos2thetav, 1.0*cos2thetav
    k13, k31, k14, k41 = - sin2thetav*np.sin(x), sin2thetav*np.sin(x),  sin2thetav*np.cos(x), -sin2thetav*np.cos(x)
    k23, k32, k24, k42 =  - sin2thetav*np.cos(x), sin2thetav*np.cos(x), -sin2thetav*np.sin(x),sin2thetav*np.sin(x)
    
    return k0*np.array([[k11,k12,k13,k14],[k21,k22,k23,k24],[k31,k32,k33,k34],[k41,k42,k43,k44]])

def derivlist(c,x):
    return np.dot(kMatrix(x),np.array([c[0],c[1],c[2],c[3]]))

xlist = np.linspace(0,800,10000)
cinit = np.array([1.0,0.0,0.0,0.0])
c = odeint(derivlist,cinit,xlist)

"""
Explanation: Solving Completely
Define the coefficient matrix K for the equations to be solved, which comes in to the equation as $\frac{dC}{dt} = K C$, where C is the vector of unknown, $C = \begin{pmatrix} C_{1,R} & C_{1,I} & C_{2,R} & C_{2,I} \end{pmatrix}$
End of explanation
"""


def probOriginal(x):
    return c[:,2]**2+c[:,3]**2

def probOriginal2(x):
    return c[:,0]**2 + c[:,1]**2

def normOriginal(x):
    return  c[:,0]**2+c[:,1]**2 + c[:,2]**2+c[:,3]**2

plt.figure(figsize = (16,9.36))
plt.plot(xlist,probOriginal(xlist)/normOriginal(xlist),'b-',label='$| C_2 |^2$')
plt.xlabel("$\omega x$")
plt.ylabel("$| C_2 |^2$")
plt.title("Transition Probility for $\lambda_0/\lambda = \omega_\lambda /\omega = 0.1$")
plt.legend()
plt.show()

plt.figure(figsize = (16,9.36))
plt.plot(xlist,probOriginal2(xlist)/normOriginal(xlist),'r-',label='$| C_1 |^2$')
plt.xlabel("$\omega x$")
plt.ylabel("$| C_1 |^2$")
plt.title("Transition Probility for $\lambda_0/\lambda = \omega_\lambda /\omega = 0.1$")
plt.legend()
plt.show()

plt.figure(figsize = (16,9.36))
plt.plot(xlist,normOriginal(xlist))
plt.show()

"""
Explanation: Probability is $\lvert C_2 \rvert^2$
End of explanation
"""


def rwac2derivlist(rwac2,x):
    
    tired = alpha*np.cos(beta*x)*cos2thetav/2- ((alpha**2)*( np.cos(beta*x) )**2)/4
    
    dc2R = rwac2[2]
    dc2I = rwac2[3]
    
    dc2pR = tired*rwac2[0] - beta*np.tan(beta*x)*rwac2[2] - rwac2[3]
    dc2pI = tired*rwac2[1] + rwac2[2] - beta*np.tan(beta*x)*rwac2[3]
    
    return np.array([dc2R,dc2I,dc2pR,dc2pI])

rwaxlist = np.linspace(0,800,10000)
rwac2init = np.array([0.0,0.0,0.0,-alpha*sin2thetav/2])
rwac2 = odeint(rwac2derivlist,rwac2init,rwaxlist)

def probRWA(x):
    return rwac2[:,0]**2 + rwac2[:,1]**2

plt.figure(figsize = (16,9.36))
plt.plot(rwaxlist,probRWA(xlist),'b-',label='$| C_2 |^2$')
plt.xlabel("$\omega x$")
plt.ylabel("$| C_1 |^2$")
plt.title("Transition Probility for $\lambda_0/\lambda = \omega_\lambda /\omega = 0.1$")
plt.legend()
plt.show()

"""
Explanation: Solving Equation for $C_2$ Numerically (Correct) NOT RWA!!!
This is not RWA! NO approximation has been made! 
The first order derivatives of $RWAC_2 = (C_{2,R},C_{2,I}, C_{2p,R},C_{2p,I})$ where $C_{2,p}=\frac{d C_2}{d\bar x}$, is given by this list.
End of explanation
"""


def probOrigMinusRWA(x):
    
    return probOriginal(x) - probRWA(x)

plt.figure(figsize = (16,9.36))
plt.plot(rwaxlist,probOrigMinusRWA(xlist),'b:',label='$\Delta | C_2 |^2$')
plt.xlabel("$\omega x$")
plt.ylabel("$| C_1 |^2$")
plt.title("Transition Probility Difference Between Two Methods for $\lambda_0/\lambda = \omega_\lambda /\omega = 0.1$")
plt.legend()
plt.show()

"""
Explanation: Compare this with the original result
End of explanation
"""


def kMatrixRWA(x):
    tired = alpha*np.cos(beta*x)*cos2thetav/2- ((alpha**2)*( np.cos(beta*x) )**2)/4
    k11, k12, k13, k14 = 0,0,1.0,0
    k21, k22, k23, k24 = 0, 0, 0, 1.0
    k31, k32, k33, k34 = -tired, 0, beta*np.tan(beta*x), 1.0
    k41, k42, k43, k44 = 0, -tired, -1.0, beta*np.tan(beta*x)

    return np.array([[k11,k12,k13,k14],[k21,k22,k23,k24],[k31,k32,k33,k34],[k41,k42,k43,k44]])

def c2derivlistRWA(c2,x):
    return np.dot(kMatrixRWA(x),np.array([c2[0],c2[1],c2[2],c2[3]]))

xlistRWA = np.linspace(0,15,1000)
cinitRWA = np.array([0.0,0.0,0.0,-alpha*cos2thetav/2])
c2 = odeint(c2derivlistRWA,cinitRWA,xlistRWA)

def probRWA(x):
    return c2[:,0]**2+c2[:,1]**2

def normc2RWA(x):
    return  c2[:,0]**2+c2[:,1]**2 + c2[:,2]**2+c2[:,3]**2

plt.figure(figsize = (16,9.36))
plt.plot(xlistRWA,probRWA(xlist),'b-',label='RWA $| C_2 |^2$ ')
plt.xlabel("$\omega x$")
plt.ylabel("$| C_2 |^2$")
plt.title("RWA Transition Probility for $\lambda_0/\lambda = \omega_\lambda /\omega = 0.1$")
plt.legend()
plt.show()

plt.figure(figsize = (16,9.36))
plt.plot(xlistRWA,normc2RWA(xlist))
plt.show()

"""
Explanation: Solving RWA Equation Numerically (Not Correct)
End of explanation
"""


from scipy.integrate import ode

def rwac2ODEFunc(x,rwac2):
    
    dc2R = rwac2[2]
    dc2I = rwac2[3]
    
    dc2pR = (-a + (a**2 - b**2)*np.cos(beta*x) )*np.cos(beta*x)*rwac2[0] + a*beta*np.sin(beta*x)*rwac2[1]+ beta*np.tan(beta*x)*rwac2[2] + ( 1 - 2*a*np.cos(beta*x) )*rwac2[3]      
    dc2pI = (- a * beta*np.sin(beta*x)  )*rwac2[0] + ( -a + (a**2 - b**2)*np.cos(beta*x))*np.cos(beta*x)*rwac2[1] - (1 - 2*a*np.cos(beta*x))*rwac2[2] + beta*np.tan(beta*x)*rwac2[3]
    
    return np.array([dc2R,dc2I,dc2pR,dc2pI])

rwac2ODE = ode(rwac2ODEFunc).set_integrator('dopri5', method='bdf')
rwainitxODE = 0.0
rwac2initODE = np.array([0.0,0.0,0.0,-a])
rwac2ODE.set_initial_value(rwac2initODE, rwainitxODE)
rwaendatODE = 10
rwadxODE = 0.0001
rwacountODE = 0
rwac2Vec = np.array([np.zeros(rwaendatODE/rwadxODE),np.zeros(rwaendatODE/rwadxODE),np.zeros(rwaendatODE/rwadxODE),np.zeros(rwaendatODE/rwadxODE)])
while rwac2ODE.successful() and rwac2ODE.t < rwaendatODE:
    rwac2ODE.integrate(rwac2ODE.t+rwadxODE)
    for i in np.array([0,1,2,3]):
        rwac2Vec[i,rwacountODE] = rwac2ODE.y[i]
    rwacountODE += 1
    
rwaxlistODE = np.linspace(rwainitxODE,rwaendatODE,rwaendatODE/rwadxODE)

plt.figure(figsize = (16,9.36))
plt.plot(rwaxlistODE,rwac2Vec[2])
plt.show()

"""
Explanation: Using ode
End of explanation
"""


