#!/usr/bin/env python
##################################################
#  Multi-Level control variate Monte Carlo       #
#                                                #
# Jun Nie                                        #
# Last modification: 01-08-2018                  #
##################################################

import sys, os
import numpy as np
import statistics as stats

class MLMFMC:
    """
    Multi-level Multi-fidelity Monte Carlo method    
    TODO: simulation fault:     
    """
    def __init__(self, pp):
                                    # configuration
        #self.pp = pp
                                    # === iniitialization
                                    # test case initialization
        self.case_init(pp)
                                    # MLMC initialization
        
    def case_init(self, pp):
        ''' Initialize test case and solver '''
        self.method = pp.method
        self.model  = pp.model
        self.interface = pp.interface
        self.maxIterations = pp.maxIterations
        
        if self.interface.case == 'heat_equation':
            from model_heat_equation import MLMC_HEAT_EQUATION
            self.case = MLMC_HEAT_EQUATION()
        else:
            print "\n:: ERROR: case %s not implemented!" %self.interface.case
            sys.exit(1)
                                    # directory for results
        self.case_dir = os.getcwd()
        self.data_dir = self.case_dir + '/%s/' %pp.data_dir
        
    def core_run(self):
        ''' core function: switch bwetween different MC methods '''
        num_mf = self.model.num_mf
        num_lev = self.method.num_lev
                                    # === CV-based
        if num_mf > 1:
            num_hf_lev = self.model.model_HF['Nlev']  # number of levels
                                    # ML-CV-MC
            if num_hf_lev > 1:
                Ycorr = True        # Ycorr or Qcorr
                if Ycorr:
                    return self.mlcvmc_Ycorr(self.model.model_LF, self.model.model_HF)
                else:
                    return self.mlcvmc_Qcorr(self.model.model_LF, self.model.model_HF)
                                    # CV-MC
            else:
                return self.cvmc(self.model.model_LF, self.model.model_HF)
                                    # === ML-MC
        elif num_lev > 1:
            Ysum = True             # Ysum or Qsum
            if Ysum:
                return self.mlmc_Ysum(self.model.model_HF)
            else:
                return self.mlmc_Qsum(self.model.model_HF)
                                    # === MC
        else:
            return self.mc(self.model.model_HF)
    
    def mc(self, model):
        ''' main function -- standard Monte Carlo. '''
        print "\nThis is standard Monte Carlo method.\n"
                                    # Initialize
        self.max_iter = 25 if self.maxIterations <= 0 else self.maxIterations
        self.eps_sq_div_2 = self.method.eps**2 / 2.
#        cost = model['COST'][-1]
        self.N = 0
        pilotSamples = self.method.pilot_samples
        delta_N = pilotSamples
        self.sums = np.zeros( (4) )
        
        niter = 0
        while( np.abs(delta_N) and niter<= self.max_iter ):
            numSamples = delta_N
            if numSamples:
                                    # generate new MC parameter sets (random variables)
                rvs = self.draw_samples(numSamples)
                qois = self.run_samples_1fidelity(model, rvs)
                self.N += numSamples
                sums = stats.accumulate_mc_sums(qois)
                self.sums += sums
            #
                                    ## mean and variance estimates
            mu, var = stats.variance_Ysum(self.sums, self.N)
            if niter == 0:
                print "MC estimator with pilot samples: N = %d" %pilotSamples
                print "  -->mean = %.4f, mse = %.4f " %(mu, var/self.N)
                print "  -->Epsilon squared target = %.4f" %self.eps_sq_div_2
                                    # target N
            N_target = var / self.eps_sq_div_2
            delta_N = stats.one_sided_delta(self.N, N_target)
            
            niter += 1
            print "\nMC iteration %d with sample increments %d:" %(niter, delta_N)
        #end while
                                    # === final sts
        self.mu, self.var = stats.variance_Ysum(self.sums, self.N)
        self.mse = self.var / self.N
        print "\nSummary of standard MC estimator:"
        print "  -->Total number of samples: N = ", self.N
        print "  -->Mean = %f, with mse = %f" %(self.mu, self.mse)
        
        return self.mu, self.mse, self.N
    
    def mlmc_init(self, model):
        ''' initialization of mlmc. '''
        self.max_iter = 25 if self.maxIterations <= 0 else self.maxIterations
#        self.convergenceTol = self.method.eps
        self.eps_sq_div_2 = self.method.eps**2 / 2.
        self.num_lev = self.method.num_lev
        assert self.num_lev <= model['Nlev'], \
            'The number of multi-levels should be no larger than that of mesh levels!'
                                    # === cost
        cost = model['COST']        # computational cost of one simulation on each level
        self.unitary_cost = float( cost[-1] )
        self.cost_l = np.zeros(self.num_lev)  # equivalent cost of one sample on each level
        for lev in range(self.num_lev):
            self.cost_l[lev] = cost[lev]
            if lev>0: self.cost_l[lev] += cost[lev-1]
                                    # Initialize for pilot sample
        self.N_l    = np.zeros((self.num_lev),dtype=int)
        self.sum_l  = np.zeros((6, self.num_lev))
        self.mean_l = np.zeros((self.num_lev))
        self.var_l  = np.zeros((self.num_lev))
        self.mean = 0.; self.var = 0.; self.mse = 0.
    
    def mlmc_Ysum(self, model):
        ''' main function -- multi-level Monte Carlo:
        This function performs "geometrical" MLMC on a single model form
        with multiple discretization levels. '''
        print "\nThis is multi-level Monte Carlo method.\n"
                                    # === initialization
        self.mlmc_init(model)
        pilotSamples = self.method.pilot_samples
        delta_N_l = pilotSamples * np.ones((self.num_lev), dtype=int)
        niter = 0; self.estimator_var0 = 0.
                                    # === converge on sample counts per level (N_l)
        while( np.sum(np.abs(delta_N_l)) and niter<= self.max_iter ):
            for lev in range(self.num_lev):
                                    # set the number of current samples from the defined increment
                numSamples = delta_N_l[lev]
                if numSamples:                      # numSamples > 0
                                    # generate new MC parameter sets (random variables)
                    rvs = self.draw_samples(numSamples)
                                    # TODO: output samples if needed
#                    exportSampleSets = False
#                    if exportSampleSets:
#                        self.export_all_samples()
                                    # compute allResponses using hierarchical model
                    qois = self.run_samples_2level(model,lev,numSamples,rvs)
                                    # accumulate new samples that calculate sums
                    sums = stats.accumulate_ml_Ysums(qois)
                                    # update raw evaluation counts and sum_l
                    self.N_l[lev] += numSamples
                    self.sum_l[:,lev] += sums[:]
                # end of numSamples
            # end of loop over num_lev
            
                                    # === total MSE = eps^2
        # is equally apportioned (eps^2 / 2) among discretization MSE and
        # estimator variance (\Sum var_Y_l / N_l).  Since we do not know the
        # discretization error, we compute an initial estimator variance and
        # then seek to reduce it by a relative_factor <= 1.'''
                                    # mean and variance estimates
            self.mean_l, self.var_l = stats.variance_Ysum(self.sum_l, self.N_l.astype(float))
                                    # MSE reference is MC applied to HF
            if (niter == 0): # eps^2 / 2 = var * relative factor
                print "MLMC estimator with pilot samples:", delta_N_l
                self.estimator_var0 = np.sum( self.var_l / self.N_l)
                mean0 = self.sum_l[4,-1] / self.N_l[-1]
                var0  = (self.sum_l[5,-1]/self.N_l[-1] - mean0**2) * self.N_l[-1]/(self.N_l[-1]-1.)
                mcMSEIter0 = var0 / self.N_l[-1]
                mse_ratio  = self.estimator_var0 / mcMSEIter0
#                #self.eps_sq_div_2   = self.estimator_var0 * self.convergenceTol
#                self.eps_sq_div_2   = mcMSEIter0 * self.convergenceTol
                print "  -->Mean square error reduced from %.4f (MC) to %.4f (ML); factor = %.2f" %(mcMSEIter0, self.estimator_var0, mse_ratio)
                print "  -->MC estimator for pilot samples: mean = %.4f, mse = %.4f " %(mean0, mcMSEIter0)
                print "  -->Epsilon squared target = %.4f" %self.eps_sq_div_2
                                    # optimal number of samples
            N_target  = stats.optimal_num_sample(self.var_l, self.cost_l, self.eps_sq_div_2)
            delta_N_l = np.maximum(0, N_target - self.N_l)
                                    # update niter 
            niter += 1
            print "\nMLMC iteration %d with sample increments:\n" %niter, delta_N_l
        # end while
                                    # === aggregate expected value of estimators for Y, Y^2, Y^3, Y^4. 
        # Final expected value is sum of expected values from telescopic sum.
        # Note: raw moments have no bias correction (no additional variance from an estimated center).        
        self.sts = stats.ml_moments(self.sum_l, self.N_l)
                                    # === compute the equivalent number of HF evaluations 
        self.equivHFEvals = np.sum( self.N_l * self.cost_l)
        self.equivHFEvals /= self.unitary_cost    # normalize into equivalent HF evals
        numMCSamples = self.sts[3] * self.N_l[-1] / self.eps_sq_div_2
        cost_saving  = float(numMCSamples) / self.equivHFEvals
        
        print "\nSummary of MLMC estimator:"
        print "  -->Total number of samples: N_l = ", self.N_l
        print "  -->The equivalent number of evaluations on finest level %d" %self.equivHFEvals
        print "  -->Compared to MC estimator which needs Nsamples = %d to achieve the same mse, \
                computational cost saving is %.2f" % (numMCSamples, cost_saving)
        print "  -->Mean = %f, with mse = %f; mu_mc = %f, mse_mc = %f" \
                %(self.sts[0], self.sts[1], self.sts[2], self.sts[3])
        
        return self.sts[0], self.sts[1], self.equivHFEvals, self.N_l
    
    def mlmc_Qsum(self, model):
        ''' main function -- multi-level Monte Carlo:
        This function performs "geometrical" MLMC on a single model form
        with multiple discretization levels.'''
        pass
    
    def cvmc_init(self, model_lf, model_hf):
        ''' initialize the CVMC estimator. '''
        self.max_iter = 25 if self.maxIterations <= 0 else self.maxIterations
#        self.eps = self.method.eps
        self.eps_sq_div_2 = self.method.eps**2 / 2.
                                                    # === cost per sample on each level
        self.lf_cost = model_lf['COST'][-1]
        self.hf_cost = model_hf['COST'][-1]
        self.cost_ratio = float(self.hf_cost) / self.lf_cost
                                    # === accumulators for computing means, variances,
                                    # and covariances across fidelity levels. 
        self.var_H = 0.
        self.sum_L_refined = np.zeros((4))
        self.sum_LL = np.zeros((4))
        self.sum_LH = np.zeros((4))
        self.sum_HH = np.zeros((4))
        self.eval_ratio = 0; self.rho2_LH = 0; self.mse_ratio  = 0.
                                    # === number of samples
        self.N_lf = 0
        self.N_hf = 0
    
    def cvmc(self, model_lf, model_hf):
        ''' main function -- multi-level control-variate Monte Carlo:
        This function performs control variate MC across two combinations of 
        model form and discretization level.'''
        print "\nThis is control-variate Monte Carlo method."
                                    # === initialization
        self.cvmc_init(model_lf, model_hf)
        pilotSamples = self.method.pilot_samples
        niter = 0
                                    # === Compute Pilot Samples
        numSamples = np.int( pilotSamples )
        print "\nCVMC pilot sample: --> LF = %d, HF = %d" %(numSamples, numSamples)
                                    # generate new MC parameter sets and compute allResponses
        rvs = self.draw_samples(numSamples)
        qois = self.run_samples_2fidelity(model_lf, model_hf, rvs)
        self.N_lf += numSamples; self.N_hf += numSamples
        sums = stats.accumulate_cv_sums(qois)
        #self.sum_L_refined += sums[0]
        self.sum_LL += sums[0]; self.sum_LH += sums[1]; self.sum_HH += sums[2]
        self.sum_L_refined += sums[0]
                                    # compute CV parameters: the LF/HF evaluation ratio
        var_H, rho2_LH, _ = stats.compute_control(self.sum_LL, self.sum_LH, self.sum_HH, self.N_hf)
                                    # Compute the LF/HF evaluation ratio
        if (rho2_LH < 1.):
            eval_ratio = np.sqrt( self.cost_ratio * rho2_LH / (1. - rho2_LH))   # r+1
                                    # compute the ratio of MC and CVMC mean squared errors (controls convergence)
        Lamda  = 1. - rho2_LH * (1. - 1. / eval_ratio)    # Ng 2014
        mc_mse = var_H / self.N_hf; cvmc_mse = mc_mse * Lamda
        mcMSEIter0 = mc_mse
        mse_ratio  = cvmc_mse / mcMSEIter0
        print "  Mean square error reduced from %f (MC) to %f (CV); factor = %f" %(mc_mse, cvmc_mse, Lamda)
#        print "  Average MSE reduction factor since pilot MC = %f targeting convergence tol = %f \n" \
#                %(mse_ratio,self.convergenceTol)
        
                                    # === Compute HF + LF increment targeting specified MSE reduction
                                    # bypass refinement if maxIterations == 0 or convergenceTol already
                                    # satisfied by pilot sample
        if ( self.maxIterations >= 1  and cvmc_mse >  self.eps_sq_div_2 ):
            # Assuming rho_LH, evaluation_ratio and var_H to be relatively invariant,
            # we seek a relative reduction in MSE using the convergence tol spec:
            #   convTol = CV_mse / MC^0_mse = mse_ratio * N0 / N
            #   delta_N = mse_ratio*N0/convTol - N0 = (mse_ratio/convTol - 1) * N0
#            incr = (mse_ratio / self.convergenceTol - 1.) * numSamples
            incr = var_H * Lamda / self.eps_sq_div_2  - self.N_hf
            numSamples = np.int( np.floor(incr + .5) )      # round
            if numSamples:
                                                    # generate new MC parameter sets and compute allResponses
                rvs = self.draw_samples(numSamples)
                qois = self.run_samples_2fidelity(model_lf, model_hf, rvs)
                niter += 1
                print "\nIteration %d with numSamples = %d" %(niter,numSamples)
                self.N_lf += numSamples; self.N_hf += numSamples
                
                sums = stats.accumulate_cv_sums(qois)
                self.sum_LL += sums[0]; self.sum_LH += sums[1]; self.sum_HH += sums[2]
                self.sum_L_refined += sums[0]
                                    # update ratios:
                var_H, rho2_LH, _ = stats.compute_control(self.sum_LL, self.sum_LH, self.sum_HH, self.N_hf)
                if (rho2_LH < 1.):
                    eval_ratio = np.sqrt( self.cost_ratio * rho2_LH / (1. - rho2_LH))
                Lamda  = 1. - rho2_LH * (1. - 1. / eval_ratio)
                mc_mse = var_H / self.N_hf; cvmc_mse = mc_mse * Lamda
                mse_ratio = cvmc_mse / mcMSEIter0
                print "  Mean square error reduced from %f (MC) to %f (CV); factor = %f" %(mc_mse, cvmc_mse, Lamda)
#                print "  Average MSE reduction factor since pilot MC = %f targeting convergence tol = %f" \
#                        %(mse_ratio,self.convergenceTol)
                #mse_ratio  = stats.MSE_ratio(eval_ratio, var_H, rho2_LH, niter, self.N_hf)
            # end
        # end
        
                                    # === Compute LF increment based on the evaluation ratio
        delta_lf = stats.lf_increment(eval_ratio, self.N_lf, self.N_hf)
        if delta_lf:
            rvs = self.draw_samples(delta_lf)
            qois = self.run_samples_1fidelity(model_lf, rvs)
            sums = stats.accumulate_cv_sums_lf(qois)
            self.sum_L_refined += sums
            self.N_lf += delta_lf
                                    # === Compute/apply CV parameter to estimate uncentered raw moments
#        moments = stats.cv_raw_moments(self.sum_LL,self.sum_LH, self.sum_HH, 
#                                       self.sum_L_refined, self.N_hf, self.N_lf, rho2_LH)
        self.sts = stats.cv_moments(self.sum_LL,self.sum_LH,self.sum_HH,
                                    self.sum_L_refined,self.N_hf,self.N_lf)
                                    # === compute the equivalent number of HF evaluations
        equivHFEvals = self.N_hf + float(self.N_lf) / self.cost_ratio
        numMCSamples = self.sts[3] / self.eps_sq_div_2
        cost_saving = numMCSamples / equivHFEvals
        print "\nSummary of CVMC estimator:"
        print "  -->Total number of samples: N_hf = %d, N_lf=%d" %(self.N_hf, self.N_lf)
        print "  -->The equivalent number of HF evaluations %d" %equivHFEvals
        print "  -->Compared to MC estimator which needs Nsamples %d to achieve the same mse, \
                computational cost saving is %.2f" % (numMCSamples, cost_saving)
        print "  -->Mean = %f, with mse = %f; mu_LF = %f, E_LF = %f" \
                %(self.sts[0], self.sts[1], self.sts[4], self.sts[5])
        
        return [self.sts[0], self.sts[1], equivHFEvals, self.N_hf, self.N_lf]
    
    def mlcvmc_init(self, model_lf, model_hf):
        ''' initialize the MLCVMC estimator: '''
        self.max_iter = 25 if self.maxIterations <= 0 else self.maxIterations
#        self.eps = self.method.eps
        self.eps_sq_div_2 = self.method.eps**2 / 2.
                
        self.num_hf_lev = model_hf['Nlev']                       # layers of multi-level
        self.num_cv_lev = np.minimum(self.num_hf_lev, model_lf['Nlev'])   # layers of control variate
                                    # === cost per sample on each level
        cost_hf = model_hf['COST']      # computational cost of one simulation on each level
        cost_lf = model_lf['COST']
        self.lf_lev_cost = np.zeros(self.num_cv_lev)
        self.hf_lev_cost = np.zeros(self.num_hf_lev)
        self.unitary_cost = float( cost_hf[-1] ) # unitary cost at the finest resolution level
                                    # cost of one sample(l+l-1)
        for lev in range(self.num_cv_lev):
            self.lf_lev_cost[lev] = cost_lf[lev]
            if lev: self.lf_lev_cost[lev] += cost_lf[lev-1]
        for lev in range(self.num_hf_lev):
            self.hf_lev_cost[lev] = cost_lf[lev]
            if lev: self.hf_lev_cost[lev] += cost_hf[lev-1]
        self.cost_ratio = self.hf_lev_cost / self.lf_lev_cost
                                    # === accumulators for computing means, variances,
                                    # and covariances across fidelity levels. 
        self.var_H = np.zeros((self.num_cv_lev))
        
        self.sum_L_refined = np.zeros((6, self.num_cv_lev))
        self.sum_LL = np.zeros((6, self.num_cv_lev))
        self.sum_LH = np.zeros((6, self.num_cv_lev))
        self.sum_HH = np.zeros((6, self.num_hf_lev))
        
        self.eval_ratios = np.zeros((self.num_cv_lev))
        self.Lambda  = np.zeros((self.num_cv_lev))
        self.rho2_LH = np.zeros((self.num_cv_lev))
                                    # === number of samples
        self.N_lf = np.zeros((self.num_cv_lev),dtype=int)
        self.N_hf = np.zeros((self.num_hf_lev),dtype=int)
        self.mu_mlcvmc = 0.; self.mse_mlcvmc = 0.
    
    def mlcvmc_Ycorr(self, model_lf, model_hf):
        ''' main function -- multi-level control-variate Monte Carlo:
        This function performs "geometrical" MLMC across discretization
        levels for the high fidelity model form where CVMC si employed
        across two model forms to exploit correlation in the discrepancies
        at each level (Y_l).'''
        print "\n This is multi-level control-variate Monte Carlo method.\n"
                                    # === initialization
        self.mlcvmc_init(model_lf, model_hf)
        pilotSamples = self.method.pilot_samples
        #delta_N_lf = np.int(pilotSamples) * np.ones((self.num_cv_lev), dtype=int)
        delta_N_hf = np.int(pilotSamples) * np.ones((self.num_hf_lev), dtype=int)
        niter = 0; self.estimator_var0 = 0.; self.estimator_mean0 = 0.
        finalCVRefinement = True
                                    # === converge on sample counts per level (N_l)
        while( np.sum(np.abs(delta_N_hf)) and niter<= self.max_iter ):
            sum_sqrt_var_cost = 0.
                                    # === loop over multi-levels
            for lev in range(self.num_hf_lev):
                                    # set the number of current samples from the defined increment
                numSamples = delta_N_hf[lev]
                                    # aggregate variances across QoI for estimating N_l (justification:
                                    # for independent QoI, sum of QoI variances = variance of QoI sum)
                if numSamples:                      # numSamples > 0
                                    # generate new MC parameter sets (random variables)
                    rvs = self.draw_samples(numSamples)
                                    # compute HF using hierarchical model
                    qois_hf = self.run_samples_2level(model_hf, lev, numSamples, rvs)
                    
                                    # === control variate betwen LF and HF for this discretization level:
                                    # if unequal number of levels, loop over all HF levels for MLMC and
                                    # apply CVMC when LF levels are available.  LF levels are assigned as
                                    # control variates to the leading set of HF levels, since these will
                                    # tend to have larger variance.
                    if lev < self.num_cv_lev:
                                    # compute allResp w/ LF model form reusing allVars from MLMC step
                        qois_lf = self.run_samples_2level(model_lf, lev, numSamples, rvs)
                                    # accumulate new samples that calculate sums
                        sums = stats.accumulate_mlcv_Ysums(qois_lf, qois_hf)
                        self.sum_LL[:, lev] += sums[0]
                        self.sum_LH[:, lev] += sums[1]
                        self.sum_HH[:, lev] += sums[2]
                        self.sum_L_refined[:,lev] += sums[0]
                                    # update evaluation counts and sum_l
                        self.N_lf[lev] += numSamples; self.N_hf[lev] += numSamples
                                    # compute the average evaluation ratio and Lambda factor
                        self.var_H[lev], self.rho2_LH[lev], self.eval_ratios[lev] = stats.eval_ratio(
                                self.sum_LL[:,lev], self.sum_LH[:,lev], self.sum_HH[:,lev], self.cost_ratio[lev], self.N_hf[lev])
                        self.Lambda[lev] = 1.- self.rho2_LH[lev] * (self.eval_ratios[lev] - 1.) / self.eval_ratios[lev]
                        
                                    # === no LF model for this level; 
                                    # accumulate only multilevel sums
                    else:
                        sums = stats.accumulate_ml_Ysums(qois_hf)
                        self.sum_HH[:, lev] += sums
                        self.N_hf[lev] += numSamples
                                    # Y variances for this level
                        _, self.var_H[lev] = stats.variance_Ysum(self.sum_HH[:, lev], self.N_hf[lev])
                # end of numSamples
                                    # accumulate sum of sqrt's of estimator var * cost used in N_target
                if lev < self.num_cv_lev:
                    sum_sqrt_var_cost += np.sqrt( self.var_H[lev] * self.hf_lev_cost[lev] / \
                                                 (1.-self.rho2_LH[lev]) ) * self.Lambda[lev]
                else:
                    sum_sqrt_var_cost += np.sqrt( self.var_H[lev] * self.hf_lev_cost[lev] )
            # end of loop over num_lev
            
                                        # === total MSE = eps^2
            # compute epsilon target based on relative tolerance: total MSE = eps^2
            # which is equally apportioned (eps^2 / 2) among discretization MSE and
            # estimator variance (\Sum var_Y_l / N_l).  Since we do not know the
            # discretization error, we compute an initial estimator variance and
            # then seek to reduce it by a relative_factor <= 1.
            if (niter == 0):    # eps^2 / 2 = var * relative factor
                mean_l, _ = stats.variance_Ysum(self.sum_HH, self.N_hf)
                self.estimator_mean0 = np.sum(mean_l)
                self.estimator_var0  = np.sum( self.var_H / self.N_hf )
#                eps_sq_div_2 = self.estimator_var0 * self.convergenceTol
#                eps_sq_div_2 = self.convergenceTol **2 / 2.
                print 'MLMC estimator with pilot samples: N_hf = N_lf = ', self.N_hf
                print '  -->mean = %f, mse = %f' %(self.estimator_mean0, self.estimator_var0)
                print '  -->Epsilon squared target = %f\n' %self.eps_sq_div_2
            
                                    # === All CV lf_increment() calls now follow all ML level evals:
            for lev in range(self.num_cv_lev):
                if delta_N_hf[lev]:
                                    # now execute additional LF sample increment, if needed
                    delta_lf = stats.lf_increment(self.eval_ratios[lev],self.N_lf[lev],self.N_hf[lev])
                    print ", for lev = %d" %lev
                    if delta_lf:
                                    # generate new MC parameter sets
                        rvs = self.draw_samples(delta_lf)
                        qois = self.run_samples_2level(model_lf, lev, delta_lf, rvs)
                        sums = stats.accumulate_ml_Ysums(qois)
                        self.sum_L_refined[:,lev] += sums
                        self.N_lf[lev] += delta_lf
                    #end delta_lf
                #end
            #end
                                    # === update targets based on variance estimates
                                    # optimal number of samples
            fact = sum_sqrt_var_cost / self.eps_sq_div_2
            for lev in range(self.num_hf_lev):
                if lev < self.num_cv_lev:
                    N_target = fact * np.sqrt(self.var_H[lev] / self.hf_lev_cost[lev] * (1. - self.rho2_LH[lev]))
                else:
                    N_target = fact * np.sqrt(self.var_H[lev] / self.hf_lev_cost[lev])
                delta_N_hf[lev] = stats.one_sided_delta(self.N_hf[lev], N_target)
                                    # update niter            
            niter += 1
            print "\nMLCVMC iteration %d with sample increments:\n" %niter, delta_N_hf
        # end of while. Iteration complete.
                                    # ===roll up raw moments from combining final
                                    #    CVMC and MLMC estimators.
        self.sts = stats.mlcv_moments(self.sum_LL,self.sum_LH,self.sum_HH,
                                      self.sum_L_refined,self.N_hf,self.N_lf)

                                    #  === compute the equivalent number of HF evaluations         
#        equivHFEvals = self.N_hf[0] * hf_cost[0] + self.N_lf[0] * lf_cost[0]; # 1st lev
                                    # subsequent levels incur 2 model costs
        equivHFEvals = 0
        equivHFEvals2 = 0           # equivHFEvals with out cost from LF
        for lev in range(self.num_hf_lev):
            equivHFEvals += self.N_hf[lev] * self.hf_lev_cost[lev]
            equivHFEvals2 += self.N_hf[lev] * self.hf_lev_cost[lev]
                                    # subsequent levels incur 2 model costs
        for lev in range(self.num_cv_lev):
            equivHFEvals += self.N_lf[lev] * self.lf_lev_cost[lev]
                                    # normalize into equivalent HF evals
        equivHFEvals /= self.unitary_cost
        equivHFEvals2 /= self.unitary_cost
        numMCSamples = self.sts[5] / self.eps_sq_div_2
        cost_saving = numMCSamples / equivHFEvals
        
        print "\nSummary of MLCVMC estimator:"
        print "  -->Total number of samples:"
        print '     --> N_hf = ', self.N_hf
        print '     --> N_lf = ', self.N_lf
#        print "  -->The equivalent number of HF evaluations %d" %equivHFEvals
#        print "  -->Compared to MC estimator which needs Nsamples = %d to achieve the same mse, \
#                computational cost saving is %.2f" % (numMCSamples, cost_saving)
        print "  MLCV estimator: Mean = %f, with mse = %f;" %(self.sts[0], self.sts[1])
        print " iMLMC estimator: Mean = %f, with mse = %f"  %(self.sts[2], self.sts[3])

        return [self.sts[0], self.sts[1], equivHFEvals, self.sts[2], self.sts[3], equivHFEvals2, self.N_hf, self.N_lf]
    
    def mlcvmc_Qcorr(self, model_lf, model_hf):
        ''' main function: mlmfmc
        This function performs "geometrical" MLMC across discretization
        levels for the high fidelity model form where CVMC is employed
        across two model forms.  It generalizes the Y_l correlation case
        to separately target correlations for each QoI level embedded
        within the level discrepancies.'''
        print "\n This is multi-level control-variate Monte Carlo method .\n"
                                    # === initialization
        self.mlcvmc_init(model_lf, model_hf)
        pilotSamples = self.method.pilot_samples
        #delta_N_lf = np.int(pilotSamples) * np.ones((self.num_cv_lev), dtype=int)
        delta_N_hf = np.int(pilotSamples) * np.ones((self.num_hf_lev), dtype=int)
        niter = 0; self.estimator_var0 = 0.; self.estimator_mean0 = 0.
        finalCVRefinement = True
                                    # === converge on sample counts per level (N_l)
        while( np.sum(np.abs(delta_N_hf)) and niter<= self.max_iter ):
            sum_sqrt_var_cost = 0.
                                    # === loop over multi-levels
            for lev in range(self.num_hf_lev):
                                    # set the number of current samples from the defined increment
                numSamples = delta_N_hf[lev]
                                    # aggregate variances across QoI for estimating N_l (justification:
                                    # for independent QoI, sum of QoI variances = variance of QoI sum)
                if numSamples:                      # numSamples > 0
                                    # generate new MC parameter sets (random variables)
                    rvs = self.draw_samples(numSamples)
                                    # compute HF using hierarchical model
                    qois_hf = self.run_samples_2level(model_hf, lev, numSamples, rvs)
                                    #=== control variate betwen LF and HF for this discretization level:
                                    # if unequal number of levels, loop over all HF levels for MLMC and
                                    # apply CVMC when LF levels are available.  LF levels are assigned as
                                    # control variates to the leading set of HF levels, since these will
                                    # tend to have larger variance.
                    if lev < self.num_cv_lev:
                                    # compute allResp w/ LF model form reusing allVars from MLMC step
                        qois_lf = self.run_samples_2level(model_lf, lev, numSamples, rvs)
                                    # accumulate new samples that calculate sums
                        sums = stats.accumulate_mlcv_Ysums(qois_lf, qois_hf)
                        self.sum_LL[:, lev] += sums[0]
                        self.sum_LH[:, lev] += sums[1]
                        self.sum_HH[:, lev] += sums[2]
                        self.sum_L_refined[:,lev] += sums[0]
                                    # update evaluation counts and sum_l
                        self.N_lf[lev] += numSamples; self.N_hf[lev] += numSamples
                                    # compute the average evaluation ratio and Lambda factor
                        self.var_H[lev], self.rho2_LH[lev], self.eval_ratios[lev] = stats.eval_ratio(
                                self.sum_LL[:,lev], self.sum_LH[:,lev], self.sum_HH[:,lev], self.cost_ratio[lev], self.N_hf[lev])
                        self.Lambda[lev] = 1.- self.rho2_LH[lev] * (self.eval_ratios[lev] - 1.) / self.eval_ratios[lev]
                                    # === no LF model for this level; 
                                    # accumulate only multilevel sums
                    else:   #no LF model for this level; accumulate only multilevel sums
                        sums = stats.accumulate_ml_Ysums(qois_hf)
                        self.sum_HH[:, lev] += sums
                        self.N_hf[lev] += numSamples
                                    # Y variances for this level
                        _, self.var_H[lev] = stats.variance_Ysum(self.sum_HH[:, lev], self.N_hf[lev])
                # end of numSamples
                                    # accumulate sum of sqrt's of estimator var * cost used in N_target
                if lev < self.num_cv_lev:
                    sum_sqrt_var_cost += np.sqrt( self.var_H[lev] * self.hf_lev_cost[lev] / \
                                                 (1.-self.rho2_LH[lev]) ) * self.Lambda[lev]
                else:
                    sum_sqrt_var_cost += np.sqrt( self.var_H[lev] * self.hf_lev_cost[lev] )
            # end of loop over num_lev
            
                                        # === total MSE = eps^2
            # compute epsilon target based on relative tolerance: total MSE = eps^2
            # which is equally apportioned (eps^2 / 2) among discretization MSE and
            # estimator variance (\Sum var_Y_l / N_l).  Since we do not know the
            # discretization error, we compute an initial estimator variance and
            # then seek to reduce it by a relative_factor <= 1.
            if (niter == 0):    # eps^2 / 2 = var * relative factor
                mean_l, _ = stats.variance_Ysum(self.sum_HH, self.N_hf)
                self.estimator_mean0 = np.sum(mean_l)
                self.estimator_var0  = np.sum( self.var_H / self.N_hf )
#                eps_sq_div_2 = self.estimator_var0 * self.convergenceTol
#                eps_sq_div_2 = self.convergenceTol **2 / 2.
                print 'MLMC estimator with pilot samples: N_hf = N_lf = ', self.N_hf
                print '  -->mean = %f, mse = %f' %(self.estimator_mean0, self.estimator_var0)
                print '  -->Epsilon squared target = %f\n' %self.eps_sq_div_2
            
                                    # === All CV lf_increment() calls now follow all ML level evals:
            for lev in range(self.num_cv_lev):
                if delta_N_hf[lev]:
                                    # now execute additional LF sample increment, if needed
                    delta_lf = stats.lf_increment(self.eval_ratios[lev],self.N_lf[lev],self.N_hf[lev])
                    print ", for lev = %d" %lev
                    if delta_lf:
                                    # generate new MC parameter sets
                        rvs = self.draw_samples(delta_lf)
                        qois = self.run_samples_2level(model_lf, lev, delta_lf, rvs)
                        sums = stats.accumulate_ml_Ysums(qois)
                        self.sum_L_refined[:,lev] += sums
                        self.N_lf[lev] += delta_lf
                    #end delta_lf
                #end
            #end
                                    # === update targets based on variance estimates
                                    # optimal number of samples
            fact = sum_sqrt_var_cost / self.eps_sq_div_2
            for lev in range(self.num_hf_lev):
                if lev < self.num_cv_lev:
                    N_target = fact * np.sqrt(self.var_H[lev] / self.hf_lev_cost[lev] * (1. - self.rho2_LH[lev]))
                else:
                    N_target = fact * np.sqrt(self.var_H[lev] / self.hf_lev_cost[lev])
                delta_N_hf[lev] = stats.one_sided_delta(self.N_hf[lev], N_target)
                                    # update niter            
            niter += 1
            print "\nMLCVMC iteration %d with sample increments:\n" %niter, delta_N_hf
        # end of while. Iteration complete.
                                    # ===roll up raw moments from combining final
                                    #    CVMC and MLMC estimators.
        self.sts = stats.mlcv_moments(self.sum_LL,self.sum_LH,self.sum_HH,
                                      self.sum_L_refined,self.N_hf,self.N_lf)

                                    #  === compute the equivalent number of HF evaluations         
#        equivHFEvals = self.N_hf[0] * hf_cost[0] + self.N_lf[0] * lf_cost[0]; # 1st lev
                                    # subsequent levels incur 2 model costs
        equivHFEvals = 0
        equivHFEvals2 = 0           # equivHFEvals with out cost from LF
        for lev in range(self.num_hf_lev):
            equivHFEvals += self.N_hf[lev] * self.hf_lev_cost[lev]
            equivHFEvals2 += self.N_hf[lev] * self.hf_lev_cost[lev]
                                    # subsequent levels incur 2 model costs
        for lev in range(self.num_cv_lev):
            equivHFEvals += self.N_lf[lev] * self.lf_lev_cost[lev]
                                    # normalize into equivalent HF evals
        equivHFEvals /= self.unitary_cost
        equivHFEvals2 /= self.unitary_cost
        numMCSamples = self.sts[5][-1] / self.eps_sq_div_2
        cost_saving = numMCSamples / equivHFEvals
        
        print "\nSummary of MLCVMC estimator:"
        print "  -->Total number of samples:"
        print '     --> N_hf = ', self.N_hf
        print '     --> N_lf = ', self.N_lf
        print "  -->The equivalent number of HF evaluations %d" %equivHFEvals
        print "  -->Compared to MC estimator which needs Nsamples = %d to achieve the same mse, \
                computational cost saving is %.2f" % (numMCSamples, cost_saving)
        print "  -->Mean = %f, with mse = %f; mu_ml = %f, var_ml = %f" \
                %(self.sts[0], self.sts[1], self.sts[2], self.sts[3])

        return [self.sts[0], self.sts[1], equivHFEvals, self.sts[2], self.sts[3], equivHFEvals2, self.N_hf, self.N_lf]
    
    def run_sample(self, nmodes, nx, rv):
        ''' single simulation with given random sample '''
        qoi = self.case.run_case(nmodes,nx,rv)
        return qoi
    
    def run_samples(self, nmodes, nx, rvs):
        ''' running multiple simulation
        TODO: Add parallel'''
        qois = [self.run_sample(nmodes,nx,rv) for rv in rvs.transpose()]
#        qois = np.zeros(rvs.shape[-1])
#        for it in range(rvs.shape[-1]):
#            qois[it] = self.run_sample(nmodes,nx,rvs[:,it])
        return np.array(qois)
    
    def run_samples_2level(self, model, lev, numSamples, rvs):
        ''' running multiple simulations on two consecutive levels'''
        nmodes = model['nmodes']
                                    # l
        nx = model['MESHSIZE'][lev]
        qf = self.run_samples(nmodes, nx, rvs)
                                    # l-1
        if lev:
            nx = model['MESHSIZE'][lev-1]
            qc = self.run_samples(nmodes, nx, rvs)
        else:                       # l=0
            qc = np.zeros(numSamples)
        
        return np.array([qc, qf])
    
    def run_samples_1fidelity(self, model_lf, rvs):
        ''' running multiple simulations with model'''
        nmodes_lf = model_lf['nmodes']
        nx_lf = model_lf['MESHSIZE'][-1]
        q_lf = self.run_samples(nmodes_lf, nx_lf, rvs)
        
        return np.array([q_lf])
    
    def run_samples_2fidelity(self, model_lf, model_hf, rvs):
        ''' running multiple simulations with two fidelity models'''
        nmodes_lf = model_lf['nmodes']
        nmodes_hf = model_hf['nmodes']
        nx_lf = model_lf['MESHSIZE'][-1]
        nx_hf = model_hf['MESHSIZE'][-1]
        q_lf = self.run_samples(nmodes_lf, nx_lf, rvs)
        q_hf = self.run_samples(nmodes_hf, nx_hf, rvs)
        
        return np.array([q_lf, q_hf])
    
    def run_samples_2level2fidelity(self, model_lf, model_hf, lev, numSamples, rvs):
        ''' running multiple simulations on two consecutive levels with two fidelity models'''
        nmodes_lf = model_lf['nmodes']
        nmodes_hf = model_hf['nmodes']
        nx_lf = model_lf['MESHSIZE'][lev]
        nx_hf = model_hf['MESHSIZE'][lev]
        qf_lf = self.run_samples(nmodes_lf,nx_lf,rvs)
        qf_hf = self.run_samples(nmodes_hf,nx_hf,rvs)
        if lev:
            nx_lf = model_lf['MESHSIZE'][lev-1]
            nx_hf = model_hf['MESHSIZE'][lev-1]
            qc_lf = self.run_samples(nmodes_lf,nx_lf,rvs)
            qc_hf = self.run_samples(nmodes_hf,nx_hf,rvs)
            
        return np.array( [[qc_lf, qf_lf], [qc_hf, qf_hf] ])
        
    def draw_samples(self, numSamples):
        ''' draw samples with size of nm'''
        return self.case.generate_sample(numSamples)
        
    def Qcorr(self):
        ''' correlation based on Q, QoI correlation'''
        pass
    
    def Ycorr(self):
        ''' correlation based on Y_l=Q_l-Q_{l-1}, discrepancy correlation'''
        pass
    
    def export_all_samples(self, rvs, prefix):
        ''' save random samples '''
        print "\n:: Writing random samples (%s)..." %prefix,
        ns = rvs.shape[-1]
        filename = os.getcwd() + '/TestCase/heat_equation/samples_%s' %prefix
        
        fh = open(filename, "a+")
        
        for i in range(ns):
            fh.write("% 8.4e    % 8.4e    % 8.4e    % 8.4e    % 8.4e    % 8.4e    % 8.4e \n"
                      %(rvs[0, i], rvs[1, i], rvs[2, i], rvs[3, i],
                        rvs[4, i], rvs[5, i], rvs[6, i] ) )

        print " done."
        fh.close()
    
class pp_mlcvmc:
    """ pp_file that define mlcvmc"""
    data_dir = 'data_01'
    maxIterations = 15
    
    class method:
        uq_method = 'mlcvmc'            # 'mlmc', 'mc', 'cvmc'
        pilot_samples = 10
        eps = 0.5
        num_lev = 4
        
    class model:
        num_mf   = 2                     # multi-fidelity models
        model_HF = {'id': 'HF',
                    'Nlev': 4,
                    'nmodes': 21,
                    'MESHSIZE': np.array([30,60,100,200]),
                    'COST': 21.*np.array([30,60,100,200]),
                    'NM': 0
                    }
        
        model_LF = {'id': 'LF',
                    'Nlev': 4,
                    'nmodes': 3,
                    'MESHSIZE': np.array([5,15,30,60]),
                    'COST': 3.*np.array([5,15,30,60]),
                    'NM': np.array([838449,   9008 ,   761,    154]),
                    
                    #'NM': np.array([22161,159,20,20]),
                    'r': np.array([50,5,3,3])
                    }
    
    class variables:
        pass
    
    class interface:
        case = 'heat_equation'
    
    class responses:
        pass

if __name__ == '__main__':
    MLMFMC(pp_mlcvmc).core_run()