# This line configures matplotlib to show figures embedded in the notebook, 
# instead of opening a new window for each figure. More about that later. 
# If you are using an old version of IPython, try using '%pylab inline' instead.
%matplotlib inline
%load_ext snakeviz

import numpy as np
from scipy.optimize import minimize
from scipy.special import expit
import matplotlib.pyplot as plt

from matplotlib.lines import Line2D

import timeit

import pandas as pd

import plotly.plotly as py
from plotly.graph_objs import *
import plotly.tools as tls

# hbar=1.054571726*10**(-34)
hbar=1.0
delm2E=1.0
# lamb=1.0  ## lambda for neutrinos
# lambb=1.0 ## lambda for anti neutrinos
# gF=1.0
# nd=1.0  ## number density
# ndb=1.0   ## number density
omega=1.0
omegab=-1.0

## Here are some matrices to be used

elM = np.array([[1.0,0.0],[0.0,0.0]])
bM = 1.0/2*np.array( [ [ -np.cos(0.4),np.sin(0.4)] , [np.sin(0.4),np.cos(0.4)] ] )
#bM = 1.0/2*np.array( [ [ 1,0] , [0,1] ] )
print bM
## sqareroot of 2
sqrt2=np.sqrt(2.0)

"""
Explanation: Vacuum Neutrino Oscillations
Here is a notebook for homogeneous gas model.
Here we are talking about a homogeneous gas bulk of neutrinos with single energy. The EoM is
$$
i \partial_t \rho_E = \left[ \frac{\delta m^2}{2E}B ,\rho_E \right]
$$
while the EoM for antineutrinos is
$$
i \partial_t \bar\rho_E = \left[- \frac{\delta m^2}{2E}B ,\bar\rho_E \right]
$$
Initial:
Homogeneous, Isotropic, Monoenergetic $\nu_e$ and $\bar\nu_e$
The equations becomes
$$
i \partial_t \rho_E = \left[ \frac{\delta m^2}{2E} B  ,\rho_E \right]
$$
$$
i \partial_t \bar\rho_E = \left[- \frac{\delta m^2}{2E}B,\bar\rho_E \right]
$$
Define $\omega=\frac{\delta m^2}{2E}$,  $\omega = \frac{\delta m^2}{-2E}$, $\mu=\sqrt{2}G_F n_\nu$
$$
i \partial_t \rho_E = \left[ \omega B ,\rho_E \right]
$$
$$
i \partial_t \bar\rho_E = \left[\bar\omega B,\bar\rho_E \right]
$$
where
$$
B = \frac{1}{2} \begin{pmatrix} 
-\cos 2\theta_v & \sin 2\theta_v \
\sin 2\theta_v & \cos 2\theta_v
\end{pmatrix}
$$
or just use theta =0.2rad
$$
L = \begin{pmatrix}
1 & 0 \
0 & 0
\end{pmatrix}
$$
Initial condition 
$$
\rho(t=0) = \begin{pmatrix}
1 & 0 \
0 & 0
\end{pmatrix}
$$
$$
\bar\rho(t=0) =\begin{pmatrix}
1 & 0 \
0 & 0
\end{pmatrix}
$$
define the following quantities

hbar$=\hbar$
%2. delm2E$= \delta m^2/2E$
%3. lamb $= \lambda$, lambb $= \bar\lambda$
%4. gF $= G_F$
%5. mu $=\mu$
omega $=\omega$, omegab $=-\bar\omega$

Numerical
End of explanation
"""


def trigf(x):
    #return 1/(1+np.exp(-x)) # It's not bad to define this function here for people could use other functions other than expit(x).
    return expit(x)

## The time derivative part

### Here are the initial conditions

init = np.array( [[1,0],[0,0]] )
#init = np.array( [[1,2],[3,4]] )

### For neutrinos

def rho(x,ti,initialCondition): # x is the input structure arrays, ti is a time point
    
    elem=np.ones(4)
    
    for i in np.linspace(0,3,4):
        elem[i] = np.sum(ti * x[i*3] * trigf( ti*x[i*3+1] + x[i*3+2] ) )
    
    return initialCondition + np.array([[ elem[0] , elem[1] ],[elem[2], elem[3] ]])


## Test
xtemp=np.ones(120)
rho(xtemp,0,init)

## Define Hamiltonians for both

def hamilv():
    return delm2E*bM

## The commutator

def commv(x,ti,initialCondition):
    
    return np.dot(hamilv(), rho(x,ti,initialCondition) ) - np.dot(rho(x,ti,initialCondition), hamilv() )
    #return rho(x,ti,initialCondition)


## Test

print bM

print hamilv()

print "neutrino\n",commv(xtemp,0,init)

nplinspace(0,9,4)

## The COST of the eqn set

regularization = 0.0001
npsum = np.sum
npravel = np.ravel
nplinspace = np.linspace
nparray = np.array

def costvTi(x,ti,initialCondition): # l is total length of x
    
    list = nplinspace(0,3,4)

    fvec = []
    costi = np.zeros(4) + 1.0j* np.zeros(4)
    commvi = npravel(commv(x,ti,initialCondition))
    
    fvecappend = fvec.append
    
    for i in list:
        fvecappend(np.asarray(trigf(ti*x[i*3+1] + x[i*3+2]) ) )
        
    fvec = nparray(fvec)
    
    for i in list:
        costi[i] = ( npsum (x[i*3]*fvec[i] + ti * x[i*3]* fvec[i] * ( 1 -  fvec[i]  ) * x[i*3+1] ) +   ( commvi[i] )  )  

    costiTemp = 0.0
    
    for i in list:
        costiTemp = costiTemp + (np.real(costi[i]))**2 + (np.imag(costi[i]))**2

    
    return costiTemp
    
 

print costvTi(xtemp,2,init), costvTi2(xtemp,2,init)

## Calculate the total cost

def costv(x,t,initialCondition):

    t = np.array(t)
    
    costvTotal = np.sum( costvTList(x,t,initialCondition)  )
        
    return costvTotal
    

def costvTList(x,t,initialCondition):  ## This is the function WITHOUT the square!!! 
        
    t = np.array(t)
    
    costvList = np.asarray([])
    
    for temp in t:
        tempElement = costvTi(x,temp,initialCondition)
        costvList = np.append(costvList, tempElement)
        
    return np.array(costvList)

    

ttemp = np.linspace(0,10)
print ttemp

ttemp = np.linspace(0,10)
print costvTList(xtemp,ttemp,init)
print costv(xtemp,ttemp,init)

"""
Explanation: I am going to substitute all density matrix elements using their corrosponding network expressions.
So first of all, I need the network expression for the unknown functions.
A function is written as
$$ y_i= 1+t_i v_k f(t_i w_k+u_k) ,$$
while it's derivative is
$$v_k f(t w_k+u_k) + t v_k f(tw_k+u_k) (1-f(tw_k+u_k)) w_k .$$
Now I can write down the equations using these two forms.
End of explanation
"""


tlin = np.linspace(0,2,10)
# tlinTest = np.linspace(0,14,10) + 0.5
# initGuess = np.ones(120)
# initGuess = np.asarray(np.split(np.random.rand(1,720)[0],12))
initGuess = np.asarray(np.split(np.random.rand(1,60)[0],12))
    


costvF = lambda x: costv(x,tlin,init)
#costvFTest = lambda x: costv(x,tlinTest,init)

print costv(initGuess,tlin,init)#, costv(initGuess,tlinTest,init)

## %%snakeviz
# startCG = timeit.default_timer()
#costvFResultCG = minimize(costvF,initGuess,method="CG")
#stopCG = timeit.default_timer()

#print stopCG - startCG

#print costvFResultCG

"""
Explanation: Minimization
Here is the minimization
End of explanation
"""


#initGuessIter = costvFResultNM.get("x")

#for TOLERANCE in np.asarray([1e-8,1e-9,1e-12,1e-14,1e-16,1e-19,1e-20,1e-21,1e-22]):
#    costvFResultNMIter = minimize(costvF,initGuessIter,method="Nelder-Mead",tol=TOLERANCE,options={"maxfun":90000,"maxiter":900000})
#    initGuessIter = costvFResultNMIter.get("x")
#    print TOLERANCE,costvFResultNMIter

#%%snakeviz
#startBFGS = timeit.default_timer()
#costvFResultBFGS = minimize(costvF,initGuess,method="BFGS")
#stopBFGS = timeit.default_timer()

#print stopBFGS - startBFGS

#print costvFResultBFGS

"""
Explanation: %%snakeviz
startNM = timeit.default_timer()
costvFResultNM = minimize(costvF,initGuess,method="Nelder-Mead",options={"maxfun":20000})
stopNM = timeit.default_timer()
print stopNM - startNM
print costvFResultNM
End of explanation
"""


%%snakeviz
startSLSQP = timeit.default_timer()
costvFResultSLSQP = minimize(costvF,initGuess,method="SLSQP")
stopSLSQP = timeit.default_timer()

print stopSLSQP - startSLSQP

print costvFResultSLSQP

#%%snakeviz
#startSLSQPTest = timeit.default_timer()
#costvFResultSLSQPTest = minimize(costvFTest,initGuess,method="SLSQP")
#stopSLSQPTest = timeit.default_timer()

#print stopSLSQPTest - startSLSQPTest

#print costvFResultSLSQPTest

print costvFResultSLSQP.get('x')
print np.asarray(np.split(costvFResultSLSQP.get('x'),12))
print CostOfStructure(np.asarray(np.split(costvFResultSLSQP.get('x'),12)))

#np.savetxt('./assets/homogen/optimize_ResultSLSQPT2120_Vac.txt', costvFResultSLSQP.get('x'), delimiter = ',')

"""
Explanation: initGuessIter = initGuess
for TOLERANCE in np.asarray([1e-19,1e-20,1e-21,1e-22]):
    costvFResultSLSQP = minimize(costvF,initGuessIter,method="SLSQP",tol=TOLERANCE)
    initGuessIter = costvFResultSLSQP.get("x")
    print TOLERANCE,costvFResultSLSQP
End of explanation
"""


# costvFResultSLSQPx = np.genfromtxt('./assets/homogen/optimize_ResultSLSQP.txt', delimiter = ',')

## The first element of neutrino density matrix
xresult = np.asarray(costvFResultSLSQP.get('x'))
#xresult = np.asarray(costvFResultNM.get('x'))
#xresult = np.asarray(costvFResultBFGS.get('x'))

print xresult

plttlin=np.linspace(0,2,100)

pltdata11 = np.array([])
pltdata11Test = np.array([])
pltdata22 = np.array([])
pltdata12 = np.array([])
pltdata21 = np.array([])

for i in plttlin:
    pltdata11 = np.append(pltdata11 ,rho(xresult,i,init)[0,0] )
    
print pltdata11

for i in plttlin:
    pltdata12 = np.append(pltdata12 ,rho(xresult,i,init)[0,1] )
    
print pltdata12

for i in plttlin:
    pltdata21 = np.append(pltdata21 ,rho(xresult,i,init)[1,0] )
    
print pltdata21

#for i in plttlin:
#    pltdata11Test = np.append(pltdata11Test ,rho(xresultTest,i,init)[0,0] )
#    
#print pltdata11Test


for i in plttlin:
    pltdata22 = np.append(pltdata22 ,rho(xresult,i,init)[1,1] )
    
print pltdata22

print rho(xresult,0,init)

rho(xresult,1,init)

#np.savetxt('./assets/homogen/optimize_pltdatar11.txt', pltdata11, delimiter = ',')
#np.savetxt('./assets/homogen/optimize_pltdatar22.txt', pltdata22, delimiter = ',')

plt.figure(figsize=(16,9.36))
plt.ylabel('rho11')
plt.xlabel('Time')
plt11=plt.plot(plttlin,pltdata11,"b4-",label="vac_rho11")
plt.plot(plttlin,np.exp(-plttlin)*1,"r+")
#plt.plot(plttlin,pltdata11Test,"m4-",label="vac_rho11Test")
plt.show()
#py.iplot_mpl(plt.gcf(),filename="vac_HG-rho11")


# tls.embed("https://plot.ly/~emptymalei/73/")

plt.figure(figsize=(16,9.36))
plt.ylabel('rho12')
plt.xlabel('Time')
plt12=plt.plot(plttlin,pltdata12,"b4-",label="vac_rho12")
plt.plot(plttlin,np.exp(-plttlin)*2.0,"r+")
#plt.plot(plttlin,pltdata11Test,"m4-",label="vac_rho11Test")
plt.show()
#py.iplot_mpl(plt.gcf(),filename="vac_HG-rho11")


# tls.embed("https://plot.ly/~emptymalei/73/")

plt.figure(figsize=(16,9.36))
plt.ylabel('rho21')
plt.xlabel('Time')
plt11=plt.plot(plttlin,pltdata21,"b4-",label="vac_rho21")
plt.plot(plttlin,np.exp(-plttlin)*3,"r+")
#plt.plot(plttlin,pltdata11Test,"m4-",label="vac_rho11Test")
plt.show()
#py.iplot_mpl(plt.gcf(),filename="vac_HG-rho11")


# tls.embed("https://plot.ly/~emptymalei/73/")

plt.figure(figsize=(16,9.36))
plt.ylabel('Time')
plt.xlabel('rho22')
plt22=plt.plot(plttlin,pltdata22,"r4-",label="vac_rho22")
plt.plot(plttlin,np.exp(-plttlin)*4,"r+")
plt.show()
#py.iplot_mpl(plt.gcf(),filename="vac_HG-rho22")

MMA_optmize_Vac_pltdata = np.genfromtxt('./assets/homogen/MMA_optmize_Vac_pltdata.txt', delimiter = ',')

plt.figure(figsize=(16,9.36))
plt.ylabel('MMArho11')
plt.xlabel('Time')
plt.plot(np.linspace(0,15,4501),MMA_optmize_Vac_pltdata,"r-",label="MMAVacrho11")
plt.plot(plttlin,pltdata11,"b4-",label="vac_rho11")
plt.show()
#py.iplot_mpl(plt.gcf(),filename="MMA-rho11-Vac-80-60")

"""
Explanation: Functions
Find the solutions to each elements.
End of explanation
"""


xtemp1 = np.arange(4)
xtemp1.shape = (2,2)
print xtemp1
xtemp1[0,1]
np.dot(xtemp1,xtemp1)
xtemp1[0,1]

"""
Explanation: Practice
End of explanation
"""


