import numpy
import os
import scipy.interpolate
import filionQuadrature

def ln(mu,cv):
   #generate log-normal distributed variable with mean mu and coefficient of variation cv
    m=numpy.log(mu/numpy.sqrt(1+cv*cv))
    s=numpy.sqrt(numpy.log(1+cv*cv))
    return ln_s(m,s)

def ln_s(mu_s,sigma_s):
   #generate r.v. with log-normal distribution
   #mu_s and sigma_s are parameters of the normal distribution in the exponent
   #see ln() for targeted mean and coefficient of variation
    z=numpy.random.normal()
    return numpy.exp(mu_s+sigma_s*z)

def readCoef(sigma):
   #read coefficients of gamma distribution set that approximates LN distro
   #with mu_s=0 and sigma_s=sigma, valid for 0.05<=sigma<=3
   #from 
   #Furman E, Hackmann D, Kuznetsov A. 
   #On log-normal convolutions: An analytical–numerical method with 
   #applications to economic capital determination.
   #Insurance: Mathematics and Economics, 90 (2020) 120-134.
   #return fa and fb which are alpha and beta for n=20 approximations 
    
    fdir=os.path.join(os.path.expanduser('~'),'software','src','PBPK',
      'lnnCoeficients20','cBase.txt')
    with open(fdir,'r') as f:
        lines=f.read().splitlines()
    fm=[q.split(' ') for q in lines]
    fm=numpy.array([[float(q) for q in x] for x in fm])
    fn=(fm.shape[1]-1)//2
    #print(fn)
    fa=fm[:,1:fn+1]
    fb=fm[:,fn+1:2*fn+1]
    fx=fm[:,0]
    falpha=[]
    fbeta=[]
    for i in range(fn):
        f2a = scipy.interpolate.interp1d(fx, fa[:,i], kind='cubic')
        try:
         falpha.append(f2a(sigma))
        except ValueError:
         #interpolation failed
         if sigma<fx[0]:
            falpha.append(fa[0,i])
            print('Using alpha[{}]={} at sigma={}'.format(fx[0],fa[0,i],sigma))

        f2b = scipy.interpolate.interp1d(fx, fb[:,i], kind='cubic')
        try:
         fbeta.append(f2b(sigma))
        except ValueError:
         if sigma<fx[0]:
            fbeta.append(fb[0,i])
            print('Using beta[{}]={} at sigma={}'.format(fx[0],fb[0,i],sigma))

    return numpy.array(falpha),numpy.array(fbeta)

def getLTransform(z,fa,fb, mu=0):
   #evaluate Laplace transform of a gamma-set convolution with 
   #coefficients alpha,fa and beta,fb
   #at point z with an optional parameter mu equal to log-normal parameter mu_s, 
   #see labeling at ln() and ln_s() functions
    s=1
    fn=len(fa)
    for i in range(fn):
        s*=numpy.power((1+z*numpy.exp(mu)/fb[i]),-fa[i])
    return s

def getLTransformGrid(sigma,fz, mu=0):
   #evaluate the Laplace transform of a gamma-set convolution 
   #that approximates a log-normal distribution with parameters
   #sigma_s=sigma, mu_s=mu at an array of complex numbers fz
    print('Reading for sigma={}'.format(sigma))
    fa,fb=readCoef(sigma)
    fn=len(fa)
    return numpy.array([getLTransform(z,fa,fb,mu) for z in fz])

def getComplexConjugatedLTransformAtMinusComplexConjugatedZGrid(sigma,fz, mu=0):
   #evaluate the Laplace transform of a gamma-set convolution 
   #that approximates a log-normal distribution with parameters
   #sigma_s=sigma, mu_s=mu at an array of complex numbers fz
    print('Reading for sigma={}'.format(sigma))
    fa,fb=readCoef(sigma)
    fn=len(fa)
    zPrime=-numpy.conj(fz)
    return numpy.conj(numpy.array([getLTransform(z,fa,fb,mu) for z in zPrime]))



def getLTransformPart(z, fa, fb, i):
   #get a Laplace transform of a single gamma distribution, specified by index
   #into a set of alpha and beta arrays fa and fb, respectively
    return numpy.power((1+z/fb[i]),-fa[i])

def getLTransformGridPart(sigma,fz,i):
   #evaluate the Laplace transform of a single gamma distribution appearing at
   #index i of alpha/beta pairs that approximate log-normal distribution 
   #for sigma_s=sigma  
    fa,fb=readCoef(sigma)
    fn=len(fa)
    print('Using alpha={} beta={}'.format(fa[i],fb[i]))
    return numpy.array([getLTransformPart(z,fa,fb,i) for z in fz])

def addExpA(fz,fq,x):
    #apply the exp(c*x)*exp(Re(a)*u*x term, see paper for reasoning
    return fq*numpy.exp(numpy.real(fz)*x)


def getZArray(n=2000):
    #internal computation parameters
    #number of grid points 1M suggested in paper, good fit found for few thousand
    #n=2000
    #integral range, 250 to 500 suggested in paper
    u0=0
    u1=300
    #constant c (to the right of all poles, saw artefacts for c of 2 or above)
    c=1
    #constant a (inclined integral path, seems to work best for Re(a) close to 0,
   #-0.1>a>-0.5 was the range tested in paper, using upper limit
    fa=complex(-0.1,1)
    
    #internal structures
    qh=numpy.linspace(u0,u1,2*n+1)
    h=(u1-u0)/2/n
    #curve segmentation
    qz=c+fa/numpy.absolute(fa)*qh
    return qz,fa,u0,h

def inverseL(fx,qz,qPhi,fa,u0,h,n):
   #Filion quadratures by parts of complex numbers
   fI=numpy.zeros(len(fx),dtype=complex)
   for i in range(len(fx)):
      x=fx[i]
      qPhi1=addExpA(qz,qPhi,x)
      fRc=filionQuadrature.filionQuadrature(numpy.real(qPhi1),x,u0,h,n,'cos')
      fIc=filionQuadrature.filionQuadrature(numpy.imag(qPhi1),x,u0,h,n,'cos')
      fRs=filionQuadrature.filionQuadrature(numpy.real(qPhi1),x,u0,h,n,'sin')
      fIs=filionQuadrature.filionQuadrature(numpy.imag(qPhi1),x,u0,h,n,'sin')
      fI[i]=complex(fRc-fIs,fIc+fRs)
   fI1=fa*fI
   return numpy.imag(fI1)

def convolveLN(fx,sigma,mu,sign=numpy.array([])):
   #convolve a set of LN with parameters sigma and mu as numpy arrays and
   #return the pdf of convolution evaluated at the set of points fx
    
   #zArray
   #number of grid points 1M suggested in paper, good fit found for few thousand
   n=2000
   qz,fa,u0,h=getZArray(n)
   #function values at grid points
   qPhi=numpy.ones(qz.size,dtype=complex)
   for i in range(len(sigma)):
      m=mu[i]
      s=sigma[i]
      try:
         if not sign[i]:
            qPhi*=getComplexConjugatedLTransformAtMinusComplexConjugatedZGrid(s,qz,m)
            continue
      except IndexError:
         pass
      
      qPhi*=getLTransformGrid(sigma[i],qz,mu[i])
        
   return inverseL(fx,qz,qPhi,fa,u0,h,n)