import numpy
import scipy.interpolate
import importlib
import convolveLN
importlib.reload(convolveLN)

def generateSample(mu,cv):
   return numpy.array([convolveLN.ln(mu,cv) for i in range(10000)])

def getEdges(a):
   #convert array of bin centers a to array of bin edges with one more element than a
   a1=numpy.zeros(len(a)+1)
   a2=numpy.zeros(a1.shape)
   a1[1:]=a
   a2[:-1]=a
   e=0.5*(a1+a2)
   e[0]=2*e[1]-e[2]
   e[-1]=2*e[-2]-e[-3]
   return e

def deltaDistribution(xTarget):
   yTarget=numpy.zeros(len(xTarget))
#fill bin at 0
   n=numpy.argmin(numpy.abs(xTarget))
   yTarget[n]=1
   return yTarget


def scale(xTarget,xOrig,yOrig,scale,sign=1):
   f = scipy.interpolate.CubicSpline(xOrig, yOrig)
   x0=numpy.min(xOrig)
   x1=numpy.max(xOrig)
   yTarget=[f(sign*scale*x) if sign*x>x0 and sign*scale*x<x1 else 0 for x in xTarget]
   s=numpy.sum(yTarget)
   if s==0:
      yTarget=deltaDistribution(xTarget)
   else:
      yTarget/=s
   return yTarget

def cDistScaled(gx,qa,cv,sign=+1,nbins=100,umax=5):
   sigmaS=numpy.sqrt(numpy.log(1+cv*cv))
   muS=numpy.log(qa/numpy.sqrt(1+cv*cv))
   q=1/numpy.sum(qa)
   muS1=muS+numpy.log(q)
#temporary space to convolve
   qx=numpy.linspace(0,umax,nbins)
   qy=convolveLN.convolveLN(qx,sigmaS,muS1)
   gy=scale(gx,qx,qy,q,sign)
   return gy

def shift(xTarget,xOrig,yOrig,shift):
   f = scipy.interpolate.CubicSpline(xOrig, yOrig)
   x0=numpy.min(xOrig)
   x1=numpy.max(xOrig)
   yTarget=[f(x-shift) if x-shift>x0 and x-shift<x1 else 0 for x in xTarget]
   s=numpy.sum(yTarget)
   if s==0:
        #fill bin at 0
      n=numpy.argmin(numpy.abs(xTarget))
      yTarget[n]=1
   else:
      yTarget/=s
   return yTarget

def calculateDistribution(xTarget,muTarget,mu,cv,dydp,nbins=100,umax=5):
   #calculate distribution of random variable y with a functional dependence 
   #on random parameters p, each with a log-normal 
   #distribution with a mean mu(p) and coefficients of variation cv(p) 
   #and dydp giving partial derivatives of y on  p
   #muTarget is the target mean of distribution
   #xTarget is the set of values x where distribution is evaluated at
   #Parameters mu, cv and dydp are given as vectors/lists
   
   qa=mu*dydp
   y0=numpy.max(numpy.abs(qa))
   fx=numpy.linspace(-5*y0,5*y0,201)
   
   sign=qa>=0
   fy1=cDistScaled(fx,qa[sign],cv[sign],nbins=nbins,umax=umax)
   fy2=cDistScaled(fx,-qa[~sign],cv[~sign],sign=-1,nbins=nbins,umax=umax)

   fy=scipy.signal.convolve(fy1,fy2,'same')

   a=numpy.average(fx,weights=fy)
   return shift(xTarget,fx,fy,muTarget-a)



def calculateDistribution1(fy,muTarget,mu,cv,dydp):
   #calculate distribution of random variable y with a functional dependence 
   #on random parameters p, each with a log-normal 
   #distribution with a mean mu(p) and coefficients of variation cv(p) 
   #and dydp giving partial derivatives of y on  p
   #muTarget is the target mean of y distribution
   #fy is the set of values y where distribution is evaluated at
   #Parameters mu, cv and dydp are given as vectors/lists
   
   print(f'muTarget={muTarget}')
   print('mu={}'.format(mu))
   #parameters qa should be around 1
   qa=mu*dydp
   print('qa={}'.format(qa))
   sign=qa>=0
   Ax=numpy.sum(qa[sign])
   Ay=numpy.sum(-qa[~sign])
   print('A={},{}'.format(Ax,Ay))
   if Ax>0:
      qa[sign]/=Ax
   if Ay>0:
      qa[~sign]/=Ay
   cvPrime=cv
   print('cvPrime={}'.format(cvPrime))

   #convolve qa multiplied (1,cv) distributions 
   sigmaS=numpy.sqrt(numpy.log(1+cvPrime*cvPrime))
   print('sigmaS={}'.format(sigmaS))
   muS=numpy.zeros(qa.shape)
   muS[sign]=numpy.log(qa[sign]/numpy.sqrt(1+cvPrime[sign]*cvPrime[sign]))
   muS[~sign]=numpy.log(-qa[~sign]/numpy.sqrt(1+cvPrime[~sign]*cvPrime[~sign]))
   print('muS={}'.format(muS))

   #do the convolution
   nbins=100
   umax=5
   fw=numpy.linspace(0,umax,nbins)
   wx=numpy.array([])
   wy=numpy.array([])
   
   if Ax>0:
      wx=convolveLN.convolveLN(fw,sigmaS[sign],muS[sign])
   
   if Ay>0:
      wy=convolveLN.convolveLN(fw,sigmaS[~sign],muS[~sign])

   #for correlate, binning of wx and wy should be equal, but one has scale A1, the other A2
   #how do you deal with that?
   if Ax>0:
      fwx=Ax*fw
   if Ay>0:
      fwy=Ay*fw

   fw=fwx
   if not Ay>Ax:
      if wy.size>0:
      #re-sample wy
         wy = scipy.interpolate.CubicSpline(fwy, wy)(fw)

   if Ay>Ax:
      fw=fwy
      #resample wx
      if wx.size>0:
         wx = scipy.interpolate.CubicSpline(fwx, wx)(fw)
   
   #output is twice the length of w and symmetric around 0
   fw1=numpy.concatenate((numpy.flip(-fw[1:]),fw))
   w=numpy.zeros(fw1.size)

   if wx.size==0:
      w[:len(wy)]=numpy.flip(wy)
   if wy.size==0:
      w[-len(wx):]=wx
   if wx.size>0 and wy.size>0:
      w=scipy.signal.correlate(wx,wy)
      #w is a convolution of input LN distributions with a mean of A


   #now shift it to the target meat muTarget
   B=muTarget-(Ax-Ay)
   print('B={}'.format(B))

   #interpolate on w and fill y
   f = scipy.interpolate.CubicSpline(fw1, w)
     
   y=numpy.zeros(fy.shape)
   
   for i in range(len(y)):
      try:
         y[i]=f(fy[i]-B)
      except ValueError:
         pass
   
   return y/numpy.sum(y)

def generateDistribution(fx,muTarget,mu,cv,dydp, dbg=0):
   #sum randomly generated variables to approximate distribution of 
   #variable y with a functional dependence y(p) on random parameters p, 
   #each with a log-normal (LN) distribution 
   #with a mean mu(p) and coefficient of variation cv(p)
   #target mean of y is muTarget
   #distribution is evaluated at points fy, 
   #dydp are derivatives of y with respect to parameters p
   
   qa=mu*dydp
   A=numpy.sum(qa)
   #qa*=A
   print('qa={}'.format(qa))
   cvPrime=cv
   samples=[a*generateSample(1,c) for a,c in zip(qa,cvPrime)]
   mean=[numpy.mean(s) for s in samples]
   std=[numpy.std(s) for s in samples]
   convSample=numpy.zeros(len(samples[0]))
   msg='[{}] mean={:.2g} std={:.2g} cv={:.2f}'
   for i in range(len(cv)):
      convSample+=samples[i]
      if dbg:
         print(msg.format(i,mean[i],std[i],std[i]/mean[i]))
   convSample+=muTarget-A   
   convM=numpy.mean(convSample)
   convS=numpy.std(convSample)
   if dbg:
      code="x"
      print(msg.format(code,convM,convS,convS/convM))
   h,be=numpy.histogram(convSample,bins=getEdges(fx))
   #h should be len(fx)
   return h/numpy.sum(h)