from numpy import *
from functools import reduce
from genx.gui_logging import iprint

# "Ordinary" implementaion of Parrats recursion formula
# theta-vector, lamda- can be a vector,n-1Dvector, d-1Dvector, sigma-1Dvector
def Refl(theta, lamda, n, d, sigma, return_int=True):
    d = d[1:-1]
    sigma = sigma[:-1]
    # Length of k-vector in vaccum
    k = 2*math.pi/lamda
    # Calculates the wavevector in each layer
    Qj = 2*n[-1]*k*sqrt(n[:,newaxis]**2/n[-1]**2-cos(theta*math.pi/180)**2)
    # Fresnel reflectivity for the interfaces
    rp=(Qj[1:]-Qj[:-1])/(Qj[1:]+Qj[:-1])*exp(-Qj[1:]*Qj[:-1]/2*sigma[:,newaxis]**2)
    # Ignoring the top and bottom layer for the calc.
    p = exp(1.0j*d[:,newaxis]*Qj[1:-1])
    # Setting up a matrix for the reduce function. Reduce only takes one array
    # as argument
    rpp=array(list(map(lambda x,y:[x,y],rp[1:],p)))
    # Paratt's recursion formula
    def formula(rtot,rint):
        return (rint[0]+rtot*rint[1])/(1+rtot*rint[0]*rint[1])
    # Implement the recursion formula
    r = reduce(formula,rpp,rp[0])

    if return_int:
        return abs(r)**2
    else:
        return r

# "Ordinary" implementaion of Parrats recursion formula
# Q-vector,n-1Dvector, d-1Dvector, sigma-1Dvector
def ReflQ(Q,lamda,n,d,sigma, return_int=True):
    # Length of k-vector in vaccum
    d=d[1:-1]
    sigma=sigma[:-1]
    Q0=4*pi/lamda
    Q = Q.astype(complex128)
    # Calculates the wavevector in each layer
    Qj=sqrt((n[:,newaxis]**2 - n[-1]**2)*Q0**2 + n[-1]**2*Q**2)
    # Fresnel reflectivity for the interfaces
    rp=(Qj[1:]-Qj[:-1])/(Qj[1:]+Qj[:-1])*exp(-Qj[1:]*Qj[:-1]/2*sigma[:,newaxis]**2)
    #print rp.shape #For debugging
    #print d.shape
    #print Qj[1:-1].shape
    p=exp(1.0j*d[:,newaxis]*Qj[1:-1]) # Ignoring the top and bottom layer for the calc.
    #print p.shape #For debugging
    # Setting up a matrix for the reduce function. Reduce only takes one array
    # as argument
    rpp=array(list(map(lambda x,y:[x,y],rp[1:],p)))
    #print rpp.shape
    # Paratt's recursion formula
    def formula(rtot,rint):
        return (rint[0]+rtot*rint[1])/(1+rtot*rint[0]*rint[1])
    # Implement the recursion formula
    r=reduce(formula,rpp,rp[0])
    #print r.shape
    # return the reflectivity 
    if return_int:
        return abs(r)**2
    else:
        return r

# Parrats recursion formula for varying n given by n_func-function
def Refl_nvary(theta,lamda,n_func,d,sigma, return_int=True):
    d=d[1:-1]
    sigma=sigma[:-1]
    # Length of k-vector in vaccum
    k=2*math.pi/lamda
    # Calculates the wavevector in each layer
    #print n_func
    ss=transpose((sin(theta[:,newaxis]*pi/180.0)/lamda)*ones(len(n_func)))
    #print ss.shape
    #print theta.shape
    #print len(n_func)
    n=array(list(map(lambda f,val:f(val),n_func,ss)))
    #print n
    Qj=2*k*n[-1]*sqrt(n**2/n[-1]**2-cos(theta*math.pi/180)**2)
    # Fresnel reflectivity for the interfaces
    rp=(Qj[1:]-Qj[:-1])/(Qj[1:]+Qj[:-1])*exp(-Qj[1:]*Qj[:-1]/2*sigma[:,newaxis]**2)
    #print rp.shape #For debugging
    #print d.shape
    #print Qj[1:-1].shape
    p=exp(1.0j*d[:,newaxis]*Qj[1:-1]) # Ignoring the top and bottom layer for the calc.
    #print p.shape #For debugging
    # Setting up a matrix for the reduce function. Reduce only takes one array
    # as argument
    rpp=array(list(map(lambda x,y:[x,y],rp[1:],p)))
    #print rpp.shape
    # Paratt's recursion formula
    def formula(rtot,rint):
        return (rint[0]+rtot*rint[1])/(1+rtot*rint[0]*rint[1])
    # Implement the recursion formula
    r=reduce(formula,rpp,rp[0])
    if return_int:
        return abs(r)**2
    else:
        return r

def ReflProfiles(theta,lamda,n,d,sigma,profile, return_int=True):
    d=d[1:-1]
    sigma=sigma[:-1]
    # Length of k-vector in vaccum
    k=2*math.pi/lamda
    # Calculates the wavevector in each layer
    Qj=2*n[-1]*k*sqrt(n[:,newaxis]**2/n[-1]**2 - cos(theta*math.pi/180)**2)
    # Function to map caclulate the roughness values
    def w(QiQj,sigma,profile):
        # erf profile
        if profile==0:
            return exp(-QiQj*sigma**2/2)
        # exponential profile
        elif profile==1:
            return 1/(1+sqrt(QiQj)*sigma**2/2)
        # linear profile
        elif profile==2:
            return sin(sqrt(3)*sigma*sqrt(QiQj))/sqrt(3)/sigma/sqrt(QiQj)
        else:
            # If any wrong number are typed in return the erf profile
            return exp(-QiQj*sigma**2/2)

    interface=array(list(map(w,Qj[:-1]*Qj[1:],sigma,profile)))
    # Fresnel reflectivity for the interfaces
    rp=(Qj[1:]-Qj[:-1])/(Qj[1:]+Qj[:-1])*interface
    #print rp.shape #For debugging
    #print d.shape
    #print Qj[1:-1].shape
    p=exp(1.0j*d[:,newaxis]*Qj[1:-1]) # Ignoring the top and bottom layer for the calc.
    #print p.shape #For debugging
    # Setting up a matrix for the reduce function. Reduce only takes one array
    # as argument
    rpp=array(list(map(lambda x,y:[x,y],rp[1:],p)))
    #print rpp.shape
    # Paratt's recursion formula
    def formula(rtot,rint):
        return (rint[0]+rtot*rint[1])/(1+rtot*rint[0]*rint[1])
    # Implement the recursion formula
    r=reduce(formula,rpp,rp[0])
    #print r.shape
    # return the reflectivity 
    if return_int:
        return abs(r)**2
    else:
        return r

# paratts algorithm for n as function of lamda or theta
def Refl_nvary2(theta,lamda,n_vector,d,sigma, return_int=True):
    d=d[1:-1]
    sigma=sigma[:-1]
    # Length of k-vector in vaccum
    k=2*math.pi/lamda
    # Calculates the wavevector in each layer
    #print n_func
    #ss=transpose((sin(theta[:,newaxis]*pi/180.0)/lamda)*ones(len(n_func)))
    #print ss.shape
    #print theta.shape
    #print len(n_func)
    #n=array(map(lambda f,val:f(val),n_func,ss))
    n=n_vector
    #print n
    Qj=2*n[-1]*k*sqrt(n**2/n[-1]**2-cos(theta*math.pi/180)**2)
    #print sigma.shape, Qj.shape
    # Fresnel reflectivity for the interfaces
    rp=(Qj[1:]-Qj[:-1])/(Qj[1:]+Qj[:-1])*exp(-Qj[1:]*Qj[:-1]/2*sigma[:,newaxis]**2)
    #print rp.shape #For debugging
    #print d.shape
    #print Qj[1:-1].shape
    p=exp(1.0j*d[:,newaxis]*Qj[1:-1]) # Ignoring the top and bottom layer for the calc.
    #print p.shape #For debugging
    # Setting up a matrix for the reduce function. Reduce only takes one array
    # as argument
    rpp=array(list(map(lambda x,y:[x,y],rp[1:],p)))
    #print rpp.shape
    # Paratt's recursion formula
    def formula(rtot,rint):
        return (rint[0]+rtot*rint[1])/(1+rtot*rint[0]*rint[1])
    # Implement the recursion formula
    r=reduce(formula,rpp,rp[0])
    if return_int:
        return abs(r)**2
    else:
        return r


def reflq_kin(q, lamda, n, d, sigma, correct_q=True, return_int=True):
    """Calculates the reflectivity in the kinematical approximation"""
    d = d[:-1]
    d[0] = 0
    z = d.sum() - d.cumsum()
    sigma = sigma[:-1]
    q0 = 4*pi/lamda
    # Kinematical reflectivity for the interfaces
    iprint(n.shape, len(n.shape))
    if len(n.shape) == 1:
        if correct_q:
            # The internal wave vector calacuted with the thickness averaged refractive index.
            n_mean = (n[1:-1] * d[1:] / d.sum()).sum()
            q_corr = sqrt((n_mean ** 2 - n[-1] ** 2) * q0 ** 2 + (n[-1] * q) ** 2)
        else:
            q_corr = q
        rp = (n[:-1] - n[1:])[:, newaxis]*exp(-(q_corr*sigma[:, newaxis])**2/2)
    else:
        if correct_q:
            # The internal wave vector calacuted with the thickness averaged refractive index.
            n_mean = (n[1:-1] * d[1:][:, newaxis] / d.sum()).sum(axis=0)
            q_corr = sqrt((n_mean ** 2 - n[-1] ** 2) * q0 ** 2 + (n[-1] * q) ** 2)
        else:
            q_corr = q
        rp = (n[:-1] - n[1:]) * exp(-(q_corr * sigma[:, newaxis]) ** 2 / 2)
    p = exp(1.0j*z[:, newaxis]*q_corr)

    r = (rp*p).sum(axis=0)*q0**2/q_corr**2/2.

    if return_int:
        return abs(r)**2
    else:
        return r


def reflq_pseudo_kin(q, lamda, n, d, sigma, return_int=True):
    """Calculates the reflectivity in a pseudo kinematical approximation.
    The mean refractive index of the film is simulated with the single reflection approximation and the deviation from
    the mean is simulated with the kinematical approximation.
    """
    d = d[:-1]
    d[0] = 0
    z = d.sum() - d.cumsum()
    sigma = sigma[:-1]
    q0 = 4*pi/lamda
    # Q = Q.astype(complex128)
    # The internal wave vector calacuted with the thickness averaged refractive index.
    n_mean = (n[1:-1]*d[1:]/d.sum()).sum()
    q_corr = sqrt((n_mean**2 - n[-1]**2)*q0**2 + n[-1]**2*q**2)
    q_sub = sqrt((n[0]**2 - n[-1]**2)*q0**2 + n[-1]**2*q**2)
    q_amb = n[-1]*q
    # Top interface
    rp_top = (q_corr - q_amb) / (q_corr + q_amb) * exp(-q_corr*q/2 * sigma[-1]**2)
    rp_sub = (q_sub - q_corr) / (q_sub + q_corr) * exp(-q_sub*q_corr/2 * sigma[0]**2)
    #rp_top = -(n[-1] - n_mean)*exp(-(q_corr*sigma[-1])**2/2)*q0**2/q_corr**2/2.
    #rp_sub = -(n_mean - n[0])*exp(-(q_corr*sigma[0])**2/2)*q0**2/q_corr**2/2.
    # Kinematical reflectivity for the interfaces
    n_diff = n - n_mean
    n_diff[0] = 0
    n_diff[-1] = 0
    rp = (n_diff[:-1] - n_diff[1:])[:, newaxis]*exp(-(q_corr*sigma[:, newaxis])**2/2)
    p = exp(1.0j*z[:, newaxis]*q_corr)

    r_kin = (rp*p).sum(axis=0)*q0**2/q_corr**2/2.
    r_sra = rp_top + rp_sub*exp(1.0j*d.sum()*q_corr)
    r = r_kin + r_sra

    if return_int:
        return abs(r)**2
    else:
        return r


def reflq_sra(q, lamda, n, d, sigma, return_int=True):
    """Single reflection approximation calculation of the reflectivity"""
    # Length of k-vector in vaccum
    d = d[1:-1]
    sigma=sigma[:-1]
    q0 = 4*pi/lamda
    # Calculates the wavevector in each layer
    if len(n.shape) == 1:
        qj = sqrt((n[:, newaxis]**2 - n[-1]**2)*q0**2 + (n[-1]*q)**2)
    else:
        qj = sqrt((n ** 2 - n[-1] ** 2) * q0 ** 2 + (n[-1] * q) ** 2)
    # Fresnel reflectivity for the interfaces
    rp = (qj[:-1] - qj[1:])/(qj[1:] + qj[:-1])*exp(-qj[1:]*qj[:-1]/2*sigma[:, newaxis]**2)
    # The wave does not transverse the ambient and substrate - ignoring them
    # Also, the wave travels from top -> bottom, the array has the first element as the substrate
    # - need to reverse the order.
    phaseterm = (d[:, newaxis]*qj[1:-1])[::-1].cumsum(axis=0)[::-1]
    p = exp(1.0J*phaseterm)
    # Adding the first interface (top -> last in array) since p is not calculated for that layer (p = 1)
    r = rp[-1] + (rp[:-1]*p).sum(axis=0)
    if return_int:
        return abs(r)**2
    else:
        return r

from . import USE_NUMBA
if USE_NUMBA:
    # try to use numba to speed up the calculation intensive functions:
    try:
        from .paratt_numba import Refl, ReflQ, Refl_nvary2
    except Exception as e:
        iprint('Could not use numba, no speed up from JIT compiler:\n'+str(e))

if __name__=='__main__':
    theta=arange(0,5,0.01)+1e-13
    #c=paratt.Refl(theta,1.54,array([1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1]),array([80,20,80,20,80,20]),0)
    #c1=Refl(theta,1.54,array([1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1]),array([80,20,80,20,80,20]),4)
    c2=Refl(theta,1.54,array([1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1-2.24e-5+2.89e-6j,1-7.57e-6+1.73e-7j,1]),array([2,80,20,80,20,80,20,2]),array([0,0,0,0,0,0,0,0]))
    #print c.shape
    #gplt.plot(theta,log10(c1),theta,log10(c2))
    #gplt.plot(theta,log10(c2))
