from pylab import *
from scipy import *
from scipy import random

def my_integrant2(x):
    """ For testing, we are integration the function
       1/(1-cos(x)*cos(y)*cos(z))/pi^3
       in the interval [0,pi]**3
    """
    #nbatch,ndim = shape(x)
    return 1.0/(1.0-cos(x[:,0])*cos(x[:,1])*cos(x[:,2]))/pi**3

class Cumulants:
  def __init__(self):
    self.sum=0.0
    self.sqsum=0.0
    self.weightsum=0.0
    self.avgsum=0.0
    self.avg2sum=0.0
    self.avg = 0.0
    self.err = 0.0
    self.chisq = 0.0
    
def SimpleMC(integrant, ndim, unit, maxeval, cum):
    nbatch=1000             # function will be evaluated in bacthes of 1000 evaluations at one time (for efficiency and storage issues)
    neval=0
    for nsamples in range(maxeval,0,-nbatch):  # loop over all_nsample evaluations in batches of nbatch
        n = min(nbatch,nsamples)  # How many evaluations in this pass?
        xr = unit*random.random((n,ndim)) # generates 2-d array of random numbers in the interval [0,1)
        wfun = integrant(xr)  # n function evaluations required in single call
        neval += n  # We just added so many fuction evaluations
        cum.sum += sum(wfun)      # sum_i f_i*w_i = <fw>
        cum.sqsum += sum(wfun*wfun)    # sum_i (f_i*w_i)^2 = <fw^2>/all_nsamples
    
    cum.avg = cum.sum/neval
    w0 = sqrt(cum.sqsum/neval)
    cum.err = sqrt((w0-cum.avg)*(w0+cum.avg)/neval) # sqrt(sqsum-sum**2)
    # Note that variance of the MC sampling is Var(monte-f)^2 = (<f^2>-<f>^2)/N == 1/sigma_i^2
    cum.avg *= unit**ndim
    cum.err *= unit**ndim

class Grid:
    """Contains the grid points g_n(x) with x=[0...1], and g=[0...1]
       for Vegas integration. There are n-dim g_n functions.
       Constraints : g(0)=0 and g(1)=1.
    """
    def __init__(self, ndim, nbins):
        self.g = zeros((ndim,nbins+1))
        self.ndim=ndim
        self.nbins=nbins
        # At the beginning we set g(x)=x
        # The grid-points are x_0 = 1/N, x_1 = 2/N, ... x_{N-1}=1.0. Note that g(0)=0.
        for ibin in range(nbins):
            w = (ibin + 1.0)/nbins  
            self.g[:,ibin] = w

    def RefineGrid(self,fxbin, SharpEdges=False):
        for idim in range(self.ndim):
            fxb = fxbin[idim,:]
            #**** smooth the f^2 value stored for each bin ****
            prev = fxb[0]
            cur = fxb[1]
            fxb[0] = 0.5*(prev + cur) # the first and the last point are different, as they can only be averaged with one neighbor
            norm = fxb[0]
            for ibin in range(1,self.nbins-1):
                s = prev + cur
                prev = cur
                cur = fxb[ibin+1]
                fxb[ibin] = (s + cur)/3.
                norm += fxb[ibin]
            fxb[nbins-1] = 0.5*(prev+cur)
            norm += fxb[nbins-1]
            if( norm == 0 ): return # can not refine the grid if the function is zero.
            fxb *= 1.0/norm         # we normalize the function.
            # Note that normalization is such that the sum is 1.
            #**** compute the importance function for each bin  ****
            imp = zeros(nbins)
            avgperbin = 0.0
            for ibin in range(nbins):
                impfun = 0.0
                if  fxb[ibin]> 0 : 
                    r = fxb[ibin]   # this is the normalized function in this bin
                    impfun = ((r - 1)/log(r))**(1.5) # this just damps the <fx> a bit. Instead of taking impfun = r,
                                    # we use g(r), where g(r)=((r-1)/log(r))^1.5 so that changes are not too rapid
                imp[ibin] = impfun
            
            avgperbin = sum(imp)/len(imp)
            #**** redefine the size of each bin  ****
            newgrid = zeros(self.nbins)
            cur=0.0
            newcur=0.0
            thisbin = 0.0
            ibin = -1
            # we are trying to determine
            #   Int[ f(g) dg, {g, g_{i-1},g_i}] == I/N_g
            #   where I == avgperbin
            for newbin in range(self.nbins-1):
                while (thisbin < avgperbin) :
                    ibin+=1
                    thisbin += imp[ibin]
                    prev = cur
                    cur = self.g[idim,ibin]
                # Explanation is in order : 
                #   prev    -- g^{old}_{l-1}
                #   cur     -- g^{old}_l
                #   thisbin -- Sm = f_{l-k}+.... +f_{l-2}+f_{l-1}+f_l
                #   we know that  Sm is just a bit more than we need, i.e., I/N_g, hence we need to compute how much more
                #   using linear interpolation :
                #   g^{new} = g_l - (g_l-g_{l-1}) * (f_{l-k}+....+f_{l-2}+f_{l-1}+f_l - I/N_g)/f_l
                #    clearly
                #         if I/N_g == f_{l-k}+....+f_{l-2}+f_{l-1}+f_l
                #            we will get g^{new} = g_l
                #     and if I/N_g == f_{l-k}+....+f_{l-2}+f_{l-1}
                #            we will get g^{new} = g_{l-1}
                #     and if I/N_g  is between the two possibilities, we will get linear interpolation between
                #     g_{l-1} and g_l
                #     
                thisbin -= avgperbin   # thisbin <- (f_{l-k}+....+f_{l-2}+f_{l-1}+f_l - I/N_g)
                delta = (cur - prev)*thisbin # delta <-  (g_l-g_{l-1})*(f_{l-k}+....+f_{l-2}+f_{l-1}+f_l - I/N_g)
                if (SharpEdges):
                    newgrid[newbin] = cur - delta/imp[ibin]
                else:
                    bin_m1 = ibin-1
                    if bin_m1<0: bin_m1=0
                    newcur = max( newcur, cur-2*delta/(imp[ibin]+imp[bin_m1]) )
                    newgrid[newbin] = newcur
            newgrid[nbins-1]=1.0
            self.g[idim,:self.nbins]= newgrid[:self.nbins]

            
def Vegas(integrant, unit, maxeval, nstart, nincrease, grid, cum):
    ndim, nbins = grid.ndim,grid.nbins  # dimension of the integral, size of the grid for binning in each direction
    unit_dim = unit**ndim   # converts from unit cube integration to generalized cube with unit length
    nbatch=1000             # function will be evaluated in bacthes of 1000 evaluations at one time (for efficiency and storage issues)
    neval=0
    print """Vegas parameters:
       ndim = """+str(ndim)+"""
       unit = """+str(unit)+"""
       maxeval = """+str(maxeval)+"""
       nstart = """+str(nstart)+"""
       nincrease = """+str(nincrease)+"""
       nbins = """+str(nbins)+"""
       nbaths = """+str(nbatch)+"\n"

    bins = zeros((nbatch,ndim),dtype=int) # in which sampled bin does this point fall?
    wgh = zeros(nbatch)                   # weights for each random point in the batch
    
    all_nsamples = nstart
    for iter in range(10000):
       fxbin = zeros((ndim,nbins))    # after each iteration we reset the average function being binned
       for nsamples in range(all_nsamples,0,-nbatch):  # loop over all_nsample evaluations in batches of nbatch
           n = min(nbatch,nsamples)  # How many evaluations in this pass?
           # We are integrating f(g_1(x),g_2(y),g_3(z))*dg_1/dx*dg_2/dy*dg_3/dz dx*dy*dz
           # This is represented as  1/all_nsamples \sum_{x_i,y_i,z_i} f(g_1(x_i),g_2(y_i),g_3(z_i))*dg_1/dx*dg_2/dy*dg_3/dz
           #  where dg_1/dx = diff*NBINS
           xr = random.random((n,ndim)) # generates 2-d array of random numbers in the interval [0,1)
           for i in range(n):
               weight = 1.0/all_nsamples
               for dim in range(ndim):
                   # We want to evaluate the function f at point g(x), i.e, f(g_1(x),g_2(y),...)
                   # Here we transform the points x,y,z -> g_1(x), g_2(y), g_3(z)
                   # We hence want to evaluate g(x) ~ g(x[i]), where x is the random number and g is the grid function
                   # The discretized g(t) is defined on the grid :
                   #       t[-1]=0, t[0]=1/N, t[1]=2/N, t[2]=3/N ... t[N-1]=1.
                   # We know that g(0)=0 and g(1)=1, so that g[-1]=0.0 and g[N-1]=1.0
                   # To interpolate g at x, we first compute  i=int(x*N) and then we use linear interpolation
                   # g(x) = g[i-1] + (g[i]-g[i-1])*(x*N-i)  ;  if i>0
                   # g(x) =   0    + (g[0]-0)*(x*N-0)       ;  if i=0
                   #
                   pos = xr[i,dim]*nbins               # which grid would it fit ?
                   ipos = int(pos)                     # the grid position is ipos
                   diff = grid.g[dim,ipos] - grid.g[dim,ipos-1]  # g[i]-g[-1]
                   xr[i,dim] = (grid.g[dim,ipos-1] + (pos-ipos)*diff)*unit # g(xr) ~ ( g[i-1]+(g[i]-g[i-1])*(x*N-i) )*[units]
                   bins[i,dim]=ipos                    # remember in which bin this random number falls.
                   weight *= diff*nbins                #  weight for this dimension is dg/dx = (g[i]-g[i-1])*N
               wgh[i] = weight # total weight is  (df/dx)*(df/dy)*(df/dx).../N_{samples}
           # Here we evaluate function f on all randomly generated x points above
           fx = integrant(xr)  # n function evaluations required in single call
           neval += n  # We just added so many fuction evaluations
       
           # Now we compute the integral as weighted average, namely, f(g(x))*dg/dx
           wfun = wgh * fx           # weight * function ~ f_i*w_i
           cum.sum += sum(wfun)      # sum_i f_i*w_i = <fw>
           wfun *= wfun
           cum.sqsum += sum(wfun)    # sum_i (f_i*w_i)^2 = <fw^2>/all_nsamples
           
           for dim in range(ndim):
               # Here we make a better approximation for the function, which we are integrating.
               for i in range(n):
                   fxbin[dim, bins[i,dim] ] += wfun[i] # just bin the function f. We saved the bin position before.
       w0 = sqrt(cum.sqsum*all_nsamples)  # w0 = sqrt(<fw^2>)
       w1 = (w0 + cum.sum)*(w0 - cum.sum) # w1 = (w0^2 - <fw>^2) = (<fw^2>-<fw>^2)
       w = (all_nsamples-1)/w1            # w ~ 1/sigma_i^2 = (N-1)/(<fw^2>-<fw>^2)
       # Note that variance of the MC sampling is Var(monte-f) = (<f^2>-<f>^2)/N == 1/sigma_i^2
       cum.weightsum += w          # weightsum ~ \sum_i 1/sigma_i^2
       cum.avgsum += w*cum.sum     # avgsum    ~ \sum_i <fw>_i / sigma_i^2
       cum.avg2sum += w*cum.sum**2 # avg2cum   ~ \sum_i <fw>_i^2/sigma_i^2
       
       cum.avg = cum.avgsum/cum.weightsum     # I_best = (\sum_i <fw>_i/sigma_i^2 )/(\sum_i 1/sigma_i^2)
       cum.err = sqrt(1/cum.weightsum)        # err ~ sqrt(best sigma^2) = sqrt(1/(\sum_i 1/sigma_i^2))
       # chi2 ~ 1/M *\sum_i (<fw>_i^2-2<fw>_i*I_best+I_best^2)/sigma_i^2= 1/M * \sum_i (<fw>_i-I_best)^2/sigma_i^2
       if iter>0:
           cum.chisq = (cum.avg2sum - 2*cum.avgsum*cum.avg + cum.weightsum*cum.avg**2)/iter
            
       print "Iteration {:3d}: I= {:10.8f} +- {:10.8f}  chisq= {:10.8f} number of evaluations = {:7d} ".format(iter+1, cum.avg*unit_dim, cum.err*unit_dim, cum.chisq, neval)
       
       grid.RefineGrid(fxbin)
       
       cum.sum=0  # reset the sums
       cum.sqsum=0
       all_nsamples += nincrease  # number of MC steps increased a bit when we have better grid
       if (neval>=maxeval): break
    
    cum.avg *= unit**ndim
    cum.err *= unit**ndim
    
if __name__ == '__main__':
    unit=pi
    ndim=3
    maxeval=200000
    nstart =10000
    nincrease=5000
    nbins=128
    
    #random.seed(0)  # so that we get every time the same result
    exact = 1.3932
    
    cum = Cumulants()
    grid = Grid(ndim,nbins)

    Vegas(my_integrant2, pi, maxeval, nstart, nincrease, grid, cum)
    maxeval = 500000
    nstart = 500000
    Vegas(my_integrant2, pi, maxeval, nstart, nincrease, grid, cum)
    
    #SimpleMC(my_integrant2, ndim, pi, maxeval, cum)
    
    print cum.avg, '+-', cum.err, 'exact=', exact
    
    plot(arange(1.,nbins+1.)/nbins, grid.g[0,:nbins], '-')
    show()