# -*- coding: utf-8 -*-
"""
Optical conductance and Lifshitz transition on graphite
FUNCTIONS
Created on Thu Apr  2 18:39:48 2020

@author: 216366
"""
import numpy as np
from scipy.constants import e, pi,c, Boltzmann as kb
import matplotlib.pyplot as plt
from sympy.solvers import solve
from scipy.integrate import odeint
from scipy import optimize
from mpl_toolkits import mplot3d
from scipy.signal import argrelextrema #finding local minima and maxima


global kb,hbar
kb=kb/e  #Boltzmann constant in eV/K
hbar=6.582119624E-16 #Reduced Planck constant in eV*s


def f(E,Ef,T):
    b=1/kb/T
    return 1/(1+np.exp(b*(E-Ef)))

def dfdE(E,Ef,T):
    b=1/kb/T
    return -b*np.exp(b*(E-Ef))/(1+np.exp(b*(E-Ef)))**2
    

def energy(kx,ky,kz,s):
    '''Energy in graphite inthe two-band approximation. s= +1 or -1 (band index)
    values of kx,ky,kz in units 2*pi/az
    WARNING: As per the article supplementary information, 
    https://journals.aps.org/prl/supplemental/10.1103/PhysRevLett.108.017602/Supplementary.pdf
    This model is only valid for |kz| not close to the Brillouin zone boundary, 
    i.e., kz < 0.242 (in 2*pi/az units). Here we leave kz up to 0.245.'''
    
    assert s==1 or s==-1, 's must be either 1 or -1'  
    
    # gamma parameters of SWM model in eV
    # gamma parameters from Schneider's thesis
#    g0 = 3.37;
#    g1 = 0.363;
#    g2 = -0.0243;
#    g3 = 0.31;
#    g4 = 0.07;
#    g5 = 0.05;
#    Delta = -0.007;
    
    # gamma parameters from Basko's paper
    g0 = 3.15;
    g1 = 0.375;
    g2 = -0.020;
    g3 = 0.315;
    g4 = 0.044;
    g5 = 0.038;
    Delta = -0.008;
    
    a0=2.46e-10;    # lattice parameter in single graphite lattice in m
    a= 1.42E-10     # Distance between neighboring carbon atoms in the same layer
    az= 3.35E-10    # Distance between graphene layers
    
    kx,ky,kz= kx*2*pi/az, ky*2*pi/az, kz*2*pi/az
    
    px,py= hbar*kx, hbar*ky
    p= np.sqrt(px**2 +py**2)
    phi= np.arctan2(py,px)
    
    C=np.cos(kz*az)
    
    v=3/2 *g0*a/hbar
    G1=2*g1*C
    G2=2*g2*C**2
    alp3=2*g3/g0*C #adimensional
    alp4=2*g4/g0*C #adimensional
    G5=2*g5*C**2 +Delta
    zeta= alp4 + (G5-G2)/(2*G1)
    
    square = (v*p)**4/G1**2 -2*alp3*(v*p)**3/G1 *np.cos(3*phi) +alp3**2*(v*p)**2
    
    return G2+ 2*zeta*(v*p)**2/G1 + s*np.sqrt(square)
#    return np.where(kz<=0.245, G2+ 2*zeta*(v*p)**2/G1 + s*np.sqrt(square), 0) #return 0 if kz>0.245

def vel(kx,ky,kz,s):
    '''Energy in graphite inthe two-band approximation. s= +1 or -1
    values of kx,ky,kz normalized with respect to az'''
    
    assert s==1 or s==-1, 's must be either 1 or -1' 
    
    #Gamma parameters of SWM model in eV
     # gamma parameters from Schneider's thesis
#    g0 = 3.37
#    g1 = 0.363
#    g2 = -0.0243
#    g3 = 0.31
#    g4 = 0.07
#    g5 = 0.05
#    Delta = -0.007
    
    # gamma parameters from Basko's paper
    g0 = 3.15
    g1 = 0.375
    g2 = -0.020
    g3 = 0.315
    g4 = 0.044
    g5 = 0.038
    Delta = -0.008
    
    a0= 2.46e-10    # lattice parameter in single graphite lattice in m
    a= 1.42E-10     # Distance between neighboring carbon atoms in the same layer
    az= 3.35E-10    # Distance between graphene layers
    
    kx,ky,kz= kx*2*pi/az, ky*2*pi/az, kz*2*pi/az #units m^-1
    
    px,py= hbar*kx, hbar*ky  #units eV*s/m
    p= np.sqrt(px**2 +py**2) #units eV*s/m
    phi= np.arctan2(py,px)  
    C=np.cos(kz*az)
    
    v=3/2 *g0*a/hbar #units m/s
    G1=2*g1*C
    G2=2*g2*C**2
    alp3=2*g3/g0*C #adimensional
    alp4=2*g4/g0*C #adimensional
    G5=2*g5*C**2 +Delta
    zeta= alp4 + (G5-G2)/(2*G1) #adimensional
    
    square = np.sqrt((v*p)**4/G1**2 -2*alp3*(v*p)**3/G1 *np.cos(3*phi) +alp3**2*(v*p)**2)

    vx = 4*zeta*v**2 *p*np.cos(phi)/G1 + s*(4*v**4 *p**3/G1**2 *np.cos(phi) -6*alp3*v**3 *p**2/G1 *np.cos(3*phi)*np.cos(phi) + 6*alp3*(v*p)**3/G1 *np.sin(3*phi) *(-1*np.sin(phi)/p)  + 2*alp3**2 *v**2*p*np.cos(phi)  )/2/square
    vy = 4*zeta*v**2 *p*np.sin(phi)/G1 + s*(4*v**4 *p**3/G1**2 *np.sin(phi) -6*alp3*v**3 *p**2/G1 *np.cos(3*phi)*np.sin(phi) + 6*alp3*(v*p)**3/G1 *np.sin(3*phi) *(np.cos(phi)/p) + 2*alp3**2 *v**2*p*np.sin(phi)  )/2/square
    
    #vx,vy=vx/np.sqrt(vx**2+vy**2), vy/np.sqrt(vx**2+vy**2) #normalizing values, just for visualization
    vtot=np.array([vx,vy])

    #mask= np.isnan(vtot) #only take the elements at which p is different to zero (resuting in v=nan)
    mask=(p<1E-15)
    vtot=np.where(mask, 0, vtot) #convert the nan to zero
    
    return vtot #units m/s

def p_t(s,kz,kx0,ky0,B):
    '''Solution of the equation of motion for px, py. Returns the time evolution of px and py. ''' 
    az= 3.35E-10 #lattice parameter in z direction. units: m
    
    #Function to integrate. -e to account for the negative charge of the electron. 
    #divided by e so it has units eV*s/m 
    f = lambda k,t: np.array([(-e*B/e)*vel(k[0],k[1],kz,s)[1], -1*(-e*B/e)*vel(k[0],k[1],kz,s)[0]])
    
    p0=np.array([kx0,ky0]) #initial conditions 
    T=30E-8 #some factor to fit some periods of cyclotron oscillation.
    
    t=np.linspace(0,T,8000) #since the differential equation is given for adimensional p, this t has units m/eV 
    p = odeint(f,p0,t)
    px_f,py_f = p[:,0]*2*pi*hbar/az ,p[:,1]*2*pi*hbar/az #reescaling to momentum units eV*s/m
    t=t*2*pi*hbar/az #getting the time in seconds
    return t, px_f, py_f

def p_t_test(s,kz,kx0,ky0,B):
    '''Solution of the equation of motion for px, py. Returns the time evolution of px and py. ''' 
    az= 3.35E-10 #lattice parameter in z direction. units: m
    
    #Function to integrate. -e to account for the negative charge of the electron. 
    #divided by e so it has units eV*s/m 
    f = lambda k,t: np.array([(-e*B/e)*vel(k[0],k[1],kz,s)[1], -1*(-e*B/e)*vel(k[0],k[1],kz,s)[0]])
    
    p0=np.array([kx0,ky0]) #initial conditions 
    T=30E-8 #some factor to fit some periods of cyclotron oscillation.
    
    t=np.linspace(0,T,8000) #since the differential equation is given for adimensional p, this t has units m/eV 
    p = odeint(f,p0,t)
    px_f,py_f = p[:,0]*2*pi*hbar/az ,p[:,1]*2*pi*hbar/az #reescaling to momentum units eV*s/m
    t=t*2*pi*hbar/az #getting the time in seconds
    return t, px_f, py_f



def Fourier_pn(t, px, py):
    '''Fourier transform of the momentum components'''
    #getting the index at which px goes back to its initial state (period of the signal)
    ind=np.where(np.isclose(px,px[0],atol=1E-11))[0] #Finds indices at which px is close to px[0]    
    #the last result is of the form ind=[0,1,2,2000,2001, 2002, 4000,...], The index 2000 is the position of the periodicity. This line
    ind=ind[np.flatnonzero(np.diff(ind) >1)+1]  #calculates concecutive differences of ind. When the difference is higher than one, we found the period.
      
    if len(ind)==0: #if the orbit is too small, the period is small too
             #just forget about such values
        return np.array([0]), np.array([0]), np.array([0])

    i=ind[0]  #Taking only one period of the function

    
    t=t[:i]
    px=px[:i]
    py=py[:i]
  
    FT_px= np.fft.fft(px)
    FT_py= np.fft.fft(py)
    w= np.fft.fftfreq(len(t), t[1]-t[0]) 
    
#    return w, FT_px, FT_py, t, px, py
    return w, FT_px, FT_py


def Fourier_pn_test(t, px, py):
    '''Fourier transform of the momentum components'''
    #getting the index at which px goes back to its initial state (period of the signal)
    ind=np.where(np.isclose(px,px[0],atol=1E-11))[0] #Finds indices at which px is close to px[0]    
    #the last result is of the form ind=[0,1,2,2000,2001, 2002, 4000,...], The index 2000 is the position of the periodicity. This line
    ind=ind[np.flatnonzero(np.diff(ind) >1)+1]  #calculates concecutive differences of ind. When the difference is higher than one, we found the period.
      
    if len(ind)==0: #if the orbit is too small, the period is small too
             #just forget about such values
        return np.array([0]), np.array([0]), np.array([0])

    i=ind[0]  #Taking only one period of the function

    
    t=t[:i]
    px=px[:i]
    py=py[:i]
  
    FT_px= np.fft.fft(px)
    FT_py= np.fft.fft(py)
    w= np.fft.fftfreq(len(t), t[1]-t[0]) 
    
    return w, FT_px, FT_py, t, px, py

def contour_area(line):
    ''' Calculates the area of closed path (Isoenergy line) generated by the function plt.contour
    INPUT: 
    line= object generated by the collections routine of the plt.contour object
    for instance:  
    cs=plt.contour(kx,ky, E1, 3)    
    line = cs.collections[0]
    OUTPUT:
    enclosed area of the path approximated by Green's theorem'''
    
    area=0 
    for i in range( len(line.get_paths()) ): #This object has as many elements as disconnected paths .
        vs = line.get_paths()[i].vertices #Vertices
        x=vs[:,0]
        y=vs[:,1]
        area= area+ abs(0.5*sum( y[:-1]*np.diff(x) -x[:-1]*np.diff(y))) #area enclosed by a closed path using Green's theorem
    return area

def cyclotron_mw(Energies,kz,s,B):
    assert abs(s)==1, 's must be either 1 or -1'   
    
    dkx, dky = 0.0001,0.0001 
    kx=np.arange(-0.02,0.02,dkx)
    ky=np.arange(-0.02,0.02,dky)
    az= 3.35E-10

    kx,ky=np.meshgrid(kx,ky)
    E1=energy(kx,ky,kz,s)
    
    ''' Isoenergy lines (contours) in kx,ky space (kz=0)'''
    #Energies=np.arange(-0.047,-0.035, 0.00001)
    Iso=plt.contour(kx,ky, E1, levels=Energies )
    
    '''getting the areas''' 
    area=np.zeros(len(Energies))
    for i in range(len(Energies)):
        line = Iso.collections[i]
        area[i]=contour_area(line)
    
    area= area*(2*pi*hbar/az)**2 #area in units of momentum^2  (eV*s/m)**2

    '''With the areas as a function of E, we can now obtain the cyclotron mass 
    and the cyclotron frequency'''
    m_cycl = 1/2/pi * np.diff(area)/np.diff(Energies) *c**2 #mass in eV 
    m_kg = m_cycl*e #mass in kg 
    w_cycl = e*B/m_kg
    return m_cycl, w_cycl


def Momentum_E(E,kz,s):
    '''Gets the values kx such that energy(kx,ky=0,kz,s)=E'''
    assert abs(s)==1, 's must be either 1 or -1'    
    
    f = lambda kx: energy(kx,0,kz,s) - E #function whose roots will be found 
    
    kx1=np.arange(-0.02,0.02,0.0001)
    ''' Finding local maxima and minima of the energy''' 
    pmin= argrelextrema(s*energy(kx1,0,kz,s), np.less)[0] 
    pmax= argrelextrema(s*energy(kx1,0,kz,s), np.greater)[0] 
    kmin= kx1[pmin] #minimum for s=1, maximum for s=-1
    kmax= kx1[pmax]
    
    E0=energy(kmin[0], 0, kz, s)
    
    
    if len(kmin)==2:  #if there are two local minima (maxima for s=-1) in the band  
        Eleg=energy(kmin[1], 0, kz, s)    
        Esaddle=energy(kmax[0], 0, kz, s)
        '''Checking that the input E value is in the energy values of that band.
        otherwise, return an empty array'''
        if s==1 and E<E0:
            return np.array([])
        if s==-1 and E>Eleg:
            return np.array([])
        
        ''' finding the kx values, here denoted as px'''          
        
        if s==1 and Eleg < E and E < Esaddle: 
            px1=optimize.root_scalar(f, bracket=[kmin[0], kmax[0]]).root
            px2=optimize.root_scalar(f, bracket=[kmax[0], kmin[1]]).root
            return np.array([px1,px2])
        elif s==-1 and Esaddle < E and E < E0: 
            px1=optimize.root_scalar(f, bracket=[kmin[0], kmax[0]]).root
            px2=optimize.root_scalar(f, bracket=[kmax[0], kmin[1]]).root
            return np.array([px1,px2])
        elif s==1:
            px1=optimize.root_scalar(f, bracket=[kmin[0], 0.1]).root  
            return np.array([px1])
        elif s==-1:
            px1=optimize.root_scalar(f, bracket=[kmin[1], 0.1]).root  #for s=-1, the energy value in the leg
            return np.array([px1])                                    #is the absolute maxima, so a root is 
                                                                      #always found choosing it as a bound in the bracket
        
    if len(kmin)==1: #If there is only one minima (maxima for s=1) in the band
        '''Checking that the input E value is in the energy values of that band.
        otherwise, return an empty array'''
        if s*E<s*E0:
            return np.array([])
        else:
            px1=optimize.root_scalar(f, bracket=[kmin[0], 0.1]).root 
            return np.array([px1])  
        