# -*- coding: utf-8 -*-
"""
Created on Mon Nov 25 09:44:15 2019

@author: Jorge Navarro
"""

import numpy as np
from scipy import optimize
from scipy.signal import resample
from scipy.interpolate import interp1d
import os
from scipy.constants import e, h, hbar, pi, m_e, Boltzmann as kb, femto
from scipy.special import polygamma as pgamma


def readFile(folder,fileName):
    '''Reads a measurement file, returning the data
    Input:
    folder: Location of the folder to read
    fileName: The name of the file.
    Output:
    Data: all numerical data of the measurement. First column: Magnetic field, second and third columns: Absorption meaasurements'''
    os.chdir(folder) #Folder where the files are stored
    f=open(fileName, 'r')
    f1=f.readlines() #Reads files line per line
    #Just to check if things go right
    #print(f1[2])
    #print(f1[16].split()[0])
    
    '''Find the first row where there are numerical values, and save the position i'''
    i=0
    for x in f1:
        try: #Checks if first element of the row is a float (converting the string to a float)
            if isinstance(float(x.split()[0]), float): #....But an error arises if it is a string
                break
        except ValueError: #The mentioned error is a ValueError, so we tell the computer it is ok to have it, just pass it
            i=i+1
        except IndexError: #Also an error if there are no characters in the line, an index error
            i=i+1
    
    Titles = f1[i-1].split('\t') #Get the labels for each column
    #print(Titles)
    Data = np.loadtxt(open(fileName,'r'), skiprows=i)
    f.close()
    return Data, Titles



def Bextrema(dcond,B):
    '''Gets the local maximum and minimum points of dcond in a B interval.
    Useful to get the weak localization points.
    model with w*tau << 1 (w:light frequency, tau: scattering time).
    INPUT:
    dcond: Derivative of the absorption (derivative of conductivity)
    B: Magnetic field
    OUTPUT:
    Positions of B for maximum and minimum dcond around B=0, and their index position'''
    
    ind=np.argmin(abs(B)) # Finds the zero in magnetic field
            
    # Find local maxima-minima in dcond. This gives the mobilities as
    ddcond=np.gradient(dcond)
    Deltai=100
    iextrema =np.argwhere(abs(ddcond[ind-Deltai:ind+Deltai])<1E-7) #Finds extrema of dcond around B=0
    
    #We keep only two terms in iextrema, as it might detect several extremal points close to each other
    for j in range(len(iextrema)):
        if j==0:
            imin=iextrema[0]
        elif (iextrema[j]-iextrema[j-1])> 4:
            imax=iextrema[j]
            
    return [B[ind-Deltai+imin], B[ind-Deltai+imax], ind-Deltai+imin, ind-Deltai+imax]



def Bminima(dcond,B,Bwindow=0.1):
    '''Gets the local minimum points of dcond in a B interval (defined 
    as Bwindow=0.1 Tesla, can be changed).
    INPUT:
    dcond: Derivative of the absorption (derivative of conductivity)
    B: Magnetic field
    OUTPUT:
    Positions of B for minimum dcond around B=0, and their index position'''
    
    Deltai= np.flatnonzero((abs(B-B[0]))>Bwindow)[0] #length of the array window
    index=np.array([])
    for j in range(len(dcond)//Deltai): # Divide the dcond function in chunks, in each of them we will find the minima
        if j!= len(dcond)//Deltai -1:
            imin=np.argmin(dcond[j*Deltai:(j+1)*Deltai-1])
            if dcond[j*Deltai +imin] < dcond[j*Deltai] and  dcond[j*Deltai +imin] < dcond[(j+1)*Deltai-1]:
                index=np.append(index,[j*Deltai +imin])
        else:
            imin=np.argmin(dcond[j*Deltai:])
            if dcond[j*Deltai +imin] < dcond[j*Deltai] and  dcond[j*Deltai +imin] < dcond[-1]:
                index=np.append(index,[j*Deltai +imin])
    index=index.astype(int)

    return B[index], index


def Drude(Bmax,dcond_max,off,B):
    '''Drude conductivity
    Input:
    Bmax: magnetic field at a conductance maximum, it is a negative number
    dcond_max: (local) maximum value of dcond by the Drude formula
    off: offset, value of dcond at B=0
    B: Magnetic field values (array)
    output:
    array: derivative of the conductivity from the Drude model'''
    
    mu=1/np.sqrt(3)/Bmax
    A=(dcond_max-off)*16*np.sqrt(3)/9/mu
    return A*mu**2. *B/(1+(mu*B)**2.)**2 +off



def gaussBaseline(B,dcond):
    '''Returns the baseline by fitting the data to the derivative of a gaussian plus 
    a linear term plus the offset'''
    
    def gaussLin(x,A1,A2,A3,A4):
        return A1*x*np.exp(-A2*(x+2.10)**2) + A3*x +A4
    
    #Initial guess for the parameters
    ind=np.argmin(abs(B))
    A2=1/16 #width of the gausian. A2=1/2sqrt(sigma), and a guess by looking at the figures is sigma +- 16Tesla
    #A4=B[ind]
    
    #Fit the baseline
    params, params_covariance= optimize.curve_fit(gaussLin,B,dcond,p0=[1,A2,-1,0])
    print ('Parameter covariance', params_covariance)
    print ('Parameters', params)
    return gaussLin(B,params[0],params[1],params[2],params[3])



def smooth(x,window_len=11,window='hanning'):
    """smooth the data using a window with requested size. Taken from Dominik's code
    
    This method is based on the convolution of a scaled window with the signal.
    The signal is prepared by introducing reflected copies of the signal 
    (with the window size) in both ends so that transient parts are minimized
    in the begining and end part of the output signal.
    
                    input:
        x: the input signal 
        window_len: the dimension of the smoothing window; should be an odd integer
        window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
            flat window will produce a moving average smoothing.

    output:
        the smoothed signal
        
    example:

    t=linspace(-2,2,0.1)
    x=sin(t)+randn(len(t))*0.1
    y=smooth(x)
    
    see also: 
    
    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
    scipy.signal.lfilter

    TODO: the window parameter could be the window itself if an array instead of a string
    NOTE: length(output) != length(input), to correct this: return y[(window_len/2-1):-(window_len/2)] instead of just y.
    NOTE (Jorge): Now gives the same length as the input, but make sure that the window length is an odd number.
    """

    if x.ndim != 1:
        raise ValueError( "smooth only accepts 1 dimension arrays." )

    if x.size < window_len:
        raise ValueError( "Input vector needs to be bigger than window size." )


    if window_len<3:
        return x


    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
        raise ValueError( "Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'" )


    s=np.r_[x[window_len-1:0:-1],x,x[-2:-window_len-1:-1]]
    #print(len(s))
    if window == 'flat': #moving average
        w=np.ones(window_len,'d')
    else:
        w=eval('np.'+window+'(window_len)')

    y=np.convolve(w/w.sum(),s,mode='valid')
    #return y
    return y[(window_len//2-1):-(window_len//2 +1)]



def fourier_SdH(dcond,B):
    '''Performs the Fourier transform of dcond(1/B). 
    Input:
    dcond= derivative of the conductivity.
    B: Magnetic field
    Output:
    B_freq: frequencies of the Fourier transform, in Tesla units
    FT_dcond: Fourier transform of dcond. Arbitrary units.
    X, Y: Interpolated and resampled dcond and 1/B, respectively'''
    
    #Interpolating dcond and 1/B such that 1/B is evenly spaced
    f=interp1d(1/B[::-1], dcond[::-1]) #f is a function
    x= np.linspace(1/B[-1],1/B[0], len(B)) #inverse of the magnetic field
    y= f(x)
    
    #Resampling the data 
    #Y,X=resample(y,len(y)//6,x)  
    Y,X=y,x
    
    #Fourier transform of dcond(1/B). Arbitrary units.
    FT_dcond= np.fft.fft(Y) 
    B_freq= np.fft.fftfreq(X.shape[-1], X[2]-X[1]) #Frequencies of the Fourier transform. In Tesla units
    
    return B_freq, FT_dcond, X, Y

def K_shift(dcond,B,phi_dim=100):
    phi=np.linspace(-2*np.pi, 2*np.pi,phi_dim)
    phi=np.vstack(phi) #column vector
    Bfreq,FT,dummyX,dummtY=fourier_SdH(dcond,B)
    return np.cos(phi)*np.real(FT), Bfreq, np.hstack(phi)

def Rdamping(T,C0,m,B):
    alpha= 2*np.pi**2*kb*m_e/hbar/e
    x=alpha*m*T/B
    return (C0*x/np.sinh(x)).flatten()

def Rdamping2(T,m):
    ''' Use this function to fit the temperature damping '''
    m=m*m_e
    T0=3
    #B=8.15662013
    B=8.12359414 # Warning: Insert this value by hand
    A= 2*np.pi**2*kb*T0/(hbar*e*B)
    return (T/T0 *np.sinh(A*m)/np.sinh(A*m*T/T0)).flatten()

def Dingle(Ampl,B,m,T):
    m=m*m_e
    #T=3
    #B=8.15662013
    #B=8.12359414
    x= 2*np.pi**2*kb*T*m/(hbar*e*B)
    return 1/B, np.log(Ampl*np.sinh(x)/x)

'''Shuibnikov-de Haas oscillations modelled by the Lifshitz-Kosevich formula'''
def LifshitzKosevich(B,BF,berry,m,T,TD,Ef,S,g=2.001):
    def RT(p):
        alpha=2*pi**2 *kb*m_e/hbar/e #units T/K
        x= alpha*(m/m_e)*(T/B)
        return x/np.sinh(p*x)
    
    def RD(p):
        alpha=2*pi**2 *kb*m_e/hbar/e #units T/K
        return np.exp( -alpha*p*m*TD/m_e/B )
    
    def RS(p):
        return np.cos(pi/2 *p*g*m/m_e)*np.ones(len(B))
    
    p1=np.array([np.arange(1,10)])
    p1= p1.transpose() #column vector
    
#    print('RT', RT(p1))
#    print('RD', RD(p1))
#    print('RS', RS(p1))
    
    s= (1/np.sqrt(p1))*RT(p1)*RD(p1)*RS(p1)*np.cos(2*pi*p1*(BF/B+berry))
    #print ('s', s)
    
    S= np.sqrt(B)*np.sum(s, axis=0) #Sum over columns
    print('S', S)
    return S
    
    
    
def WeakLocalization(B,N,Blr,Bi,Bp):
    '''derivative of conductance as a funciton of B, from
    Weak localization function from Falko and Altshuler theory. Check https://doi.org/10.1016/j.ssc.2007.03.049
    INPUT:
    B: (1d vector) magnetic field
    N: (scalar) normalization factor
    Blr: (scalar) characteristic field for long range scattering
    Bi: (scalar) characteristic field for intervalley scattering
    Bp: (scalar) characteristic field for dephasing
    OUTPUT:
    (1d vector) contribution to d\sigma/dB by weak localization.
    '''
    def F(z):
        return np.where(z==0, 0, 1/z - pgamma(1,0.5+ 1/z)/z**2 ) #pgamma(1,x) is the trigamma function
          
    #return e**2/(np.pi*h) * ( (1/Bp)*F(abs(B)/Bp) - 1/(Bp+2*Bi)*F(abs(B)/(Bp+2*Bi)) - 2/(Bp+Bi+Blr)*F(abs(B)/(Bp+Bi+Blr))  ) 
    return N* np.sign(B)*( (1/Bp)*F(abs(B)/Bp) - 1/(Bp+2*Bi)*F(abs(B)/(Bp+2*Bi)) - 2/(Bp+Bi+Blr)*F(abs(B)/(Bp+Bi+Blr))  ) 

def WL2(B,N,Blr,Bi,Bp):
    '''magnetoesistivity, from
    Weak localization function from Falko and Altshuler theory. Check https://doi.org/10.1016/j.ssc.2007.03.049
    INPUT:
    B: (1d vector) magnetic field
    N: (scalar) normalization factor
    Blr: (scalar) characteristic field for long range scattering
    Bi: (scalar) characteristic field for intervalley scattering
    Bp: (scalar) characteristic field for dephasing
    OUTPUT:
    (1d vector) contribution to conductance by weak localization.
    '''
    def F(z):
        return np.where(z==0, 0, 1/z - pgamma(1,0.5+1/z)/z**2 )
          
    #return e**2/(np.pi*h) * ( (1/Bp)*F(abs(B)/Bp) - 1/(Bp+2*Bi)*F(abs(B)/(Bp+2*Bi)) - 2/(Bp+Bi+Blr)*F(abs(B)/(Bp+Bi+Blr))  ) 
    return N* ( F(abs(2*B)/Bp) - F(abs(2*B)/(Bp+2*Bi)) - F(abs(2*B)/(Bp+Bi+Blr))  ) 


