'''
SEG2020 is a script that runs all the computations for the different results 
presented at the SEG2020 conference (Houston, Texas).

***Improving BEL1D accuracy for geophysical imaging of the subsurface***

It runs at first the numerical benchmark for a dataset that is created directly.
    - Creating the dataset 
    - Running BEL1D (initialization + first iteration)
    - Presenting graphs of the results
    - Applying IPR
    - Presenting graphs of the improved result
    - Comparing to McMC results from DREAM (results provided in github)

Then, it runs the same for the Mont Rigi dataset and presents simply the obtained profile.

Author: 
Hadrien MICHEL
ULiège, UGent, F.R.S.-FNRS
hadrien[dot]michel[at]uliege[dot]be
(c) August 2020
'''
def round_to_5(x,n=1): 
    import math as mt
    # Modified from: https://stackoverflow.com/questions/3410976/how-to-round-a-number-to-significant-figures-in-python
    tmp = [(round(a, -int(mt.floor(mt.log10(abs(a)))) + (n-1)) if a != 0.0 else 0.0) for a in x]
    return tmp

if __name__=='__main__':
    import numpy as np
    from pyBEL1D import BEL1D
    from pyBEL1D.utilities import Tools
    from matplotlib import pyplot
    from pathos import multiprocessing as mp
    from pathos import pools as pp
    import time
    ################################################################################
    ###                                                                          ###
    ###                          Numerical Benchmarking                          ###
    ###                                                                          ###
    ################################################################################
    Benchmark = True
    if Benchmark:
        # 1) Creating the dataset:
        from pygimli.physics import sNMR
        KernelBench = "Data/sNMR/Tx50Rx50.mrsk" # A kernel file generated by MRSMatlab (Mueller-Petke et al., 2012)
        nbLayer = 3 # 2 layers and a half-space
        TimingBench = np.arange(0.005,0.5,0.001) # The time vector for the model
        # KFile = sNMR.MRS()
        # KFile.loadKernel(KernelBench)
        # ModellingMethod = sNMR.MRS1dBlockQTModelling(nlay=nbLayer,K=KFile.K,zvec=KFile.z,t=TimingBench)
        ModelBench = np.asarray([25, 25, 0.05, 0.25, 0.1, 0.1, 0.2, 0.05])
        Noise = 10 # nV
        # DatasetBench = ModellingMethod.response(ModelBench)
        # DatasetBench += np.random.normal(loc=0, scale=Noise*1e-9,size=DatasetBench.shape) # Adding noise to the dataset
        # np.save('sNMR_Bench_Dataset',DatasetBench)# To use in the McMC algorithm
        DatasetBench = np.load('sNMR_Bench_Dataset.npy') # Load the dataset that was already created and run in McMC
        # 2) Initializing BEL1D:
        priorBench = np.array([[0.0, 50.0, 0.0, 0.15, 0.0, 0.5], [0.0, 50.0, 0.15, 0.50, 0.0, 0.5], [0.0, 0.0, 0.0, 0.15, 0.0, 0.5]])
        # Initialsizing the parameters:
        start = time.time()
        ModelParam = BEL1D.MODELSET.SNMR(prior=priorBench, Kernel=KernelBench, Timing=TimingBench)
        # 3) Running pre-BEL operations:
        nbSampled = 10000
        PreBEL_Bench = BEL1D.PREBEL(ModelParam,nbModels=nbSampled)
        PreBEL_Bench.run()
        # 4) Sampling 10000 models from the posterior:
        PostBEL_Bench = BEL1D.POSTBEL(PreBEL_Bench)
        PostBEL_Bench.run(Dataset=DatasetBench,nbSamples=nbSampled,NoiseModel=Noise)
        end = time.time()

        PostBEL_Bench.KDE.ShowKDE(Xvals=PostBEL_Bench.CCA.transform(PostBEL_Bench.PCA['Data'].transform(np.reshape(DatasetBench,(1,-1)))))
        PostBEL_Bench.ShowDataset(RMSE=True)
        CurrentGraph = pyplot.gcf()
        CurrentGraph = CurrentGraph.get_axes()[0]
        CurrentGraph.plot(TimingBench,DatasetBench[:len(TimingBench)],'k',linestyle='None',marker='o',markerfacecolor='None')
        # Graph for the CCA space parameters loads
        _, ax = pyplot.subplots()
        B = PostBEL_Bench.CCA.y_loadings_
        B = np.divide(np.abs(B).T,np.repeat(np.reshape(np.sum(np.abs(B),axis=0),(1,B.shape[0])),B.shape[0],axis=0).T)
        ind =  np.asarray(range(B.shape[0]))+1
        ax.bar(x=ind,height=B[0],label=r'${}$'.format(PostBEL_Bench.MODPARAM.paramNames["NamesSU"][0]))
        for i in range(B.shape[0]+1)[1:-1]:
            ax.bar(x=ind,height=B[i],bottom=np.reshape(np.sum(B[0:i],axis=0),(B.shape[0],)),label=r'${}$'.format(PostBEL_Bench.MODPARAM.paramNames["NamesSU"][i]))
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width, box.height*0.8])
        ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.4), ncol=3)
        ax.set_ylabel('Relative contribution')
        ax.set_xlabel('CCA dimension')
        pyplot.show(block=False)

        PostBEL_Bench.ShowPostModels(TrueModel=ModelBench,RMSE=True)
        PostBEL_Bench.ShowPostCorr(TrueModel=ModelBench)
        pyplot.show(block=False)
        
        MODELS_1stITER = PostBEL_Bench.SAMPLES
        MODELS_1stITER_DATA = PostBEL_Bench.SAMPLESDATA
        PRE_MODS = PreBEL_Bench.MODELS
        PRE_DATA = PreBEL_Bench.FORWARD

        Postbel = PostBEL_Bench
        Prebel = PreBEL_Bench
        # 5) Applying IPR
        nbIter = 100 # maximum number of iterations
        tolerance = 5e-3 # Tolerance on the normalized difference between the distributions
        nbParam = int(priorBench.size/2 - 1)
        means = np.zeros((nbIter,nbParam))
        stds = np.zeros((nbIter,nbParam))
        timings = np.zeros((nbIter,))
        MODELS_ITER = np.zeros((nbIter,nbSampled,nbParam)) # Store the models that have been computed
        diverge = True
        distancePrevious = 1e10
        MixingUpper = 0
        MixingLower = 1
        for idxIter in range(nbIter):
            PostbelLast = Postbel
            PrebelLast = Prebel
            if idxIter == 0: # Initialization: already done (see 2 and 3)
                # PostbelTest.KDE.ShowKDE(Xvals=PostbelTest.CCA.transform(PostbelTest.PCA['Data'].transform(np.reshape(Dataset,(1,-1)))))
                means[idxIter,:], stds[idxIter,:] = Postbel.GetStats()
                timings[idxIter] = end-start

                ModLastIter = Prebel.MODELS
            else:
                ModLastIter = Postbel.SAMPLES
                # Here, we will use the POSTBEL2PREBEL function that adds the POSTBEL from previous iteration to the prior (Iterative prior resampling)
                # However, the computations are longer with a lot of models, thus you can opt-in for the "simplified" option which randomely select up to 10 times the numbers of models
                MixingUpper += 1
                MixingLower += 1
                Mixing = MixingUpper/MixingLower
                Prebel = BEL1D.PREBEL.POSTBEL2PREBEL(PREBEL=Prebel,POSTBEL=Postbel,Dataset=DatasetBench,NoiseModel=Noise,Simplified=True,nbMax=nbSampled,MixingRatio=Mixing)
                # Since when iterating, the dataset is known, we are not computing the full relationship but only the posterior distributions directly to gain computation timing
                print(idxIter+1)
                Postbel = BEL1D.POSTBEL(Prebel)
                Postbel.run(DatasetBench,nbSamples=nbSampled,NoiseModel=None)
                means[idxIter,:], stds[idxIter,:] = Postbel.GetStats()
                end = time.time()
                timings[idxIter] = end-start
            # The distance is computed on the normalized distributions. Therefore, the tolerance is relative.
            diverge, distance = Tools.ConvergeTest(SamplesA=ModLastIter,SamplesB=Postbel.SAMPLES, tol=tolerance)
            print('Wasserstein distance: {}'.format(distance))
            if not(diverge) or (abs((distancePrevious-distance)/distancePrevious)*100<1):# If the distance between the distributions is not changing, we converged as well
                # Convergence acheived if:
                # 1) Distance below threshold
                # 2) Distance does not vary significantly (less than 2.5%)
                print('Model has converged at iter {}!'.format(idxIter+1))
                MODELS_ITER[idxIter,:,:] = Postbel.SAMPLES
                break
            distancePrevious = distance
            MODELS_ITER[idxIter,:,:] = Postbel.SAMPLES
            start = time.time()
        timings = timings[:idxIter+1]
        means = means[:idxIter+1,:]
        stds = stds[:idxIter+1,:]
        MODELS_ITER = MODELS_ITER[:idxIter+1,:,:]
        np.save('ModelsIteration',MODELS_ITER)

        # 6) Graphs for the results:
        Postbel.ShowDataset(RMSE=True)
        CurrentGraph = pyplot.gcf()
        CurrentGraph = CurrentGraph.get_axes()[0]
        CurrentGraph.plot(TimingBench,DatasetBench[:len(TimingBench)],'k',linestyle='None',marker='o',markerfacecolor='None')
        #Postbel.ShowPostCorr(TrueModel=ModelBench,OtherMethod=PRE_MODS)
        Postbel.ShowPostModels(TrueModel=ModelBench,RMSE=True)
        # Graph for the CCA space parameters loads
        _, ax = pyplot.subplots()
        B = Postbel.CCA.y_loadings_
        B = np.divide(np.abs(B).T,np.repeat(np.reshape(np.sum(np.abs(B),axis=0),(1,B.shape[0])),B.shape[0],axis=0).T)
        ind =  np.asarray(range(B.shape[0]))+1
        ax.bar(x=ind,height=B[0],label=r'${}$'.format(Postbel.MODPARAM.paramNames["NamesSU"][0]))
        for i in range(B.shape[0]+1)[1:-1]:
            ax.bar(x=ind,height=B[i],bottom=np.reshape(np.sum(B[0:i],axis=0),(B.shape[0],)),label=r'${}$'.format(Postbel.MODPARAM.paramNames["NamesSU"][i]))
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width, box.height*0.8])
        ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.4), ncol=3)
        ax.set_ylabel('Relative contribution')
        ax.set_xlabel('CCA dimension')
        pyplot.show(block=False)

        # Add KDE graph at last iteration:
        PrebelGraph = BEL1D.PREBEL.POSTBEL2PREBEL(PREBEL=PrebelLast,POSTBEL=PostbelLast, NoiseModel=Noise,RemoveOutlier=True,Simplified=True,nbMax=nbSampled,MixingRatio=Mixing)
        PostbelGraph = BEL1D.POSTBEL(PREBEL=PrebelGraph)
        print('Postbel for graphs initialized')
        PostbelGraph.run(Dataset=DatasetBench, nbSamples=10000)
        print('Printing KDE Graphs')
        PostbelGraph.KDE.ShowKDE(Xvals=PostbelGraph.CCA.transform(PostbelGraph.PCA['Data'].transform(np.reshape(DatasetBench,(1,-1)))))

        print('Total computation time: {} seconds'.format(np.sum(timings)))

        # Comparison with McMC:
        OtherMethod = np.load('Data/sNMR/SEG2020_Bench.npy')
        fig = pyplot.figure(figsize=[10,10])# Creates the figure space
        axs = fig.subplots(nbParam, nbParam)
        for i in range(nbParam):
            for j in range(nbParam):
                if i == j: # Diagonal
                    if i != nbParam-1:
                        axs[i,j].get_shared_x_axes().join(axs[i,j],axs[-1,j])# Set the xaxis limit
                    axs[i,j].hist(np.squeeze(MODELS_ITER[0,:,j]),color='gray',density=True,alpha=0.5)
                    axs[i,j].hist(np.squeeze(MODELS_ITER[3,:,j]),color='b',density=True,alpha=0.5)
                    axs[i,j].hist(np.squeeze(MODELS_ITER[6,:,j]),color='m',density=True,alpha=0.5)
                    axs[i,j].hist(OtherMethod[:,j],color='y',density=True,alpha=0.5)
                    axs[i,j].hist(np.squeeze(MODELS_ITER[-1,:,j]),color='g',density=True,alpha=0.5)
                    axs[i,j].plot([ModelBench[i],ModelBench[i]],np.asarray(axs[i,j].get_ylim()),'r')
                    if nbParam > 8:
                        axs[i,j].set_xticks([])
                        axs[i,j].set_yticks([])
                elif i > j: # Below the diagonal -> Scatter plot
                    if i != nbParam-1:
                        axs[i,j].get_shared_x_axes().join(axs[i,j],axs[-1,j])# Set the xaxis limit
                    if j != nbParam-1:
                        if i != nbParam-1:
                            axs[i,j].get_shared_y_axes().join(axs[i,j],axs[i,-1])# Set the yaxis limit
                        else:
                            axs[i,j].get_shared_y_axes().join(axs[i,j],axs[i,-2])# Set the yaxis limit
                    axs[i,j].plot(np.squeeze(MODELS_ITER[0,:,j]),np.squeeze(MODELS_ITER[0,:,i]),color='gray',linestyle='None',marker='.')
                    axs[i,j].plot(np.squeeze(MODELS_ITER[3,:,j]),np.squeeze(MODELS_ITER[3,:,i]),'.b')
                    axs[i,j].plot(np.squeeze(MODELS_ITER[6,:,j]),np.squeeze(MODELS_ITER[6,:,i]),'.m')
                    axs[i,j].plot(np.squeeze(MODELS_ITER[-1,:,j]),np.squeeze(MODELS_ITER[-1,:,i]),'.g')
                    axs[i,j].plot(ModelBench[j],ModelBench[i],'.r')
                    if nbParam > 8:
                        axs[i,j].set_xticks([])
                        axs[i,j].set_yticks([])
                elif OtherMethod is not None:
                    if i != nbParam-1:
                        axs[i,j].get_shared_x_axes().join(axs[i,j],axs[-1,j])# Set the xaxis limit
                    if j != nbParam-1:
                        if i != 0:
                            axs[i,j].get_shared_y_axes().join(axs[i,j],axs[i,-1])# Set the yaxis limit
                        else:
                            axs[i,j].get_shared_y_axes().join(axs[i,j],axs[i,-2])# Set the yaxis limit
                    axs[i,j].plot(np.squeeze(MODELS_ITER[0,:,j]),np.squeeze(MODELS_ITER[0,:,i]),color='gray',linestyle='None',marker='.')
                    axs[i,j].plot(OtherMethod[:,j],OtherMethod[:,i],'.y')
                    axs[i,j].plot(ModelBench[j],ModelBench[i],'.r')
                    if nbParam > 8:
                        axs[i,j].set_xticks([])
                        axs[i,j].set_yticks([])
                else:
                    axs[i,j].set_visible(False)
                if j == 0: # First column of the graph
                    if ((i==0)and(j==0)) or not(i==j):
                        axs[i,j].set_ylabel(r'${}$'.format(Postbel.MODPARAM.paramNames["NamesSU"][i]))
                if i == nbParam-1: # Last line of the graph
                    axs[i,j].set_xlabel(r'${}$'.format(Postbel.MODPARAM.paramNames["NamesSU"][j]))
                if j == nbParam-1:
                    if not(i==j):
                        axs[i,j].yaxis.set_label_position("right")
                        axs[i,j].yaxis.tick_right()
                        axs[i,j].set_ylabel(r'${}$'.format(Postbel.MODPARAM.paramNames["NamesSU"][i]))
                if i == 0:
                    axs[i,j].xaxis.set_label_position("top")
                    axs[i,j].xaxis.tick_top()
                    axs[i,j].set_xlabel(r'${}$'.format(Postbel.MODPARAM.paramNames["NamesSU"][j]))
        fig.suptitle("Posterior model space visualtization")
        from matplotlib.lines import Line2D
        custom_lines = [Line2D([0], [0], color='gray', lw=4),
                        Line2D([0], [0], color='y', lw=4),
                        Line2D([0], [0], color='b', lw=4),
                        Line2D([0], [0], color='m', lw=4),
                        Line2D([0], [0], color='g', lw=4),
                        Line2D([0], [0], color='r', lw=4)]
        fig.legend(custom_lines, ['Prior', 'DREAM', '3rd Iteration', '6th Iteration', 'Last Iteration', 'Benchmark'],loc='lower center',ncol=6)
        for ax in axs.flat:
            ax.label_outer()
        pyplot.show(block=False)

        # Graph with the models:
        DREAM = OtherMethod # Already loaded
        Prior = PRE_MODS
        Iter5 = np.squeeze(MODELS_ITER[3,:,:])
        Iter10 = np.squeeze(MODELS_ITER[6,:,:])
        IterLast = np.squeeze(MODELS_ITER[-1,:,:])
        pltidx = [1,4,2,5,8,3,6,9]
        fig = pyplot.figure(figsize=[10,10])
        for idx in range(8):
            ax = pyplot.subplot(3,3,pltidx[idx])
            pyplot.hist(Prior[:,idx].ravel(),bins=50,density=True,alpha=0.5,label='Prior')
            pyplot.hist(Iter5[:,idx].ravel(),bins=50,density=True,alpha=0.5,label='3rd iteration')
            pyplot.hist(Iter10[:,idx].ravel(),bins=50,density=True,alpha=0.5,label='6th iteration')
            pyplot.hist(IterLast[:,idx].ravel(),bins=50,density=True,alpha=0.5,label='Last iteration')
            pyplot.hist(DREAM[:,idx].ravel(),bins=50,density=True,alpha=0.5,label='DREAM')
            ax.plot([ModelBench[idx],ModelBench[idx]],np.asarray(ax.get_ylim()),label='Benchmark')
            pyplot.plot()
            if pltidx[idx]==1:
                ax.set_ylabel('Layer 1')
            elif pltidx[idx]==4:
                ax.set_ylabel('Layer 2')
            elif pltidx[idx]==8:
                ax.set_xlabel('Water content [/]')
            elif pltidx[idx]==9:
                ax.set_xlabel('Relaxation time [sec]')
        handles, labels = ax.get_legend_handles_labels()
        ax = pyplot.subplot(3,3,7)# Not used but labels needed
        ax.set_xlabel('Thickness [m]')
        ax.set_ylabel('Half-space')
        ax.spines['bottom'].set_color('None')
        ax.spines['top'].set_color('None') 
        ax.spines['right'].set_color('None')
        ax.spines['left'].set_color('None')
        ax.xaxis.set_ticks([])
        ax.yaxis.set_ticks([])
        #pyplot.axis('off')
        CurrentGraph = pyplot.gcf()
        CurrentGraph.legend(handles, labels,loc='lower center', ncol=6)

        # Graph for the i-th pusle moment:
        import matplotlib
        from matplotlib import colors
        from scipy import stats
        
        RMS = np.sqrt(np.square(np.subtract(DatasetBench,Postbel.SAMPLESDATA)).mean(axis=-1))
        quantiles = np.divide([stats.percentileofscore(RMS,a,'strict') for a in RMS],100)
        sortIndex = np.argsort(RMS)
        sortIndex = np.flip(sortIndex)
        
        fig = pyplot.figure()
        ax = fig.add_subplot(1, 1, 1)
        
        colormap = matplotlib.cm.get_cmap('jet')
        for j in sortIndex:
            ax.plot(Postbel.MODPARAM.forwardFun["Axis"],np.squeeze(Postbel.SAMPLESDATA[j,-5*len(Postbel.MODPARAM.forwardFun["Axis"]):-4*len(Postbel.MODPARAM.forwardFun["Axis"])]),color=colormap(quantiles[j]))
        ax.plot(TimingBench,DatasetBench[-5*len(Postbel.MODPARAM.forwardFun["Axis"]):-4*len(Postbel.MODPARAM.forwardFun["Axis"])],'k',linestyle='None',marker='o',markerfacecolor='None')
        ax.set_xlabel(r'${}$'.format(Postbel.MODPARAM.paramNames["DataAxis"]),fontsize=14)
        ax.set_ylabel(r'${}$'.format(Postbel.MODPARAM.paramNames["DataName"]),fontsize=14)
                
        fig.subplots_adjust(bottom=0.30)
        ax_colorbar = fig.add_axes([0.10, 0.15, 0.80, 0.05])
        nb_inter = 1000
        color_for_scale = colormap(np.linspace(0,1,nb_inter,endpoint=True))
        cmap_scale = colors.ListedColormap(color_for_scale)
        scale = [stats.scoreatpercentile(RMS,a,limit=(np.min(RMS),np.max(RMS)),interpolation_method='lower') for a in np.linspace(0,100,nb_inter,endpoint=True)]
        norm = colors.BoundaryNorm(scale,len(color_for_scale))
        data = np.atleast_2d(np.linspace(np.min(RMS),np.max(RMS),nb_inter,endpoint=True))
        ax_colorbar.imshow(data, aspect='auto',cmap=cmap_scale,norm=norm)
        ax_colorbar.set_xlabel('Root Mean Square Error {}'.format(Postbel.MODPARAM.paramNames["DataUnits"]),fontsize=12)
        ax_colorbar.yaxis.set_visible(False)
        nbTicks = 5
        ax_colorbar.set_xticks(ticks=np.linspace(0,nb_inter,nbTicks,endpoint=True))
        ax_colorbar.set_xticklabels(labels=round_to_5([stats.scoreatpercentile(RMS,a,limit=(np.min(RMS),np.max(RMS)),interpolation_method='lower') for a in np.linspace(0,100,nbTicks,endpoint=True)],n=5),rotation=15,ha='center')

        pyplot.show()

    ################################################################################
    ###                                                                          ###
    ###                           Case study: Mont Rigi                          ###
    ###                                                                          ###
    ################################################################################
    MtRigi = False
    if MtRigi:
        # Load the field data
        from pygimli.physics import sNMR
        Dataset = "Data/sNMR/SEG2020_MtRigi.mrsd"
        Kernel = "Data/sNMR/SEG2020_MtRigi.mrsk"
        ModelParam = sNMR.MRS()
        sNMR.MRS.loadKernel(ModelParam,Kernel)
        sNMR.MRS.loadMRSI(ModelParam,Dataset)
        FieldData = np.ravel(ModelParam.dcube)
        TimingField = ModelParam.t
        Noise = 18 #nV
        # Initialize BEL1D:
        nbSampled = 5000
        priorMtRigi = np.asarray([[0.0, 7.5, 0.30, 0.80, 0.0, 0.200], [0, 0, 0.0, 0.15, 0.100, 0.400]])
        start = time.time()
        MODEL_MtRigi = BEL1D.MODELSET().SNMR(prior=priorMtRigi,Kernel=Kernel, Timing=TimingField)
        PREBEL_MtRigi = BEL1D.PREBEL(MODEL_MtRigi,nbModels=nbSampled)
        PREBEL_MtRigi.run()
        POSTBEL_MtRigi = BEL1D.POSTBEL(PREBEL_MtRigi)
        POSTBEL_MtRigi.run(FieldData,nbSamples=nbSampled,NoiseModel=Noise)
        end = time.time()
        POSTBEL_MtRigi.ShowDataset(RMSE=True,Prior=True)
        CurrentGraph = pyplot.gcf()
        CurrentGraph = CurrentGraph.get_axes()[0]
        CurrentGraph.plot(TimingField,FieldData[:len(TimingField)],'k',linestyle='None',marker='o',markerfacecolor='None')
        pyplot.show(block=False)
        POSTBEL_MtRigi.ShowPostModels(RMSE=True)
        # Iterations:
        Postbel = POSTBEL_MtRigi
        Prebel = PREBEL_MtRigi
        nbIter = 100 # maximum number of iterations
        tolerance = 5e-3 # Tolerance on the normalized difference between the distributions
        nbParam = int(priorMtRigi.size/2 - 1)
        means = np.zeros((nbIter,nbParam))
        stds = np.zeros((nbIter,nbParam))
        timings = np.zeros((nbIter,))
        MODELS_ITER = np.zeros((nbIter,nbSampled,nbParam)) # Store the models that have been computed
        diverge = True
        distancePrevious = 1e10
        MixingUpper = 0
        MixingLower = 1
        for idxIter in range(nbIter):
            PostbelLast = Postbel
            PrebelLast = Prebel
            if idxIter == 0: # Initialization: already done (see 2 and 3)
                # PostbelTest.KDE.ShowKDE(Xvals=PostbelTest.CCA.transform(PostbelTest.PCA['Data'].transform(np.reshape(Dataset,(1,-1)))))
                means[idxIter,:], stds[idxIter,:] = Postbel.GetStats()
                timings[idxIter] = end-start

                ModLastIter = Prebel.MODELS
            else:
                ModLastIter = Postbel.SAMPLES
                # Here, we will use the POSTBEL2PREBEL function that adds the POSTBEL from previous iteration to the prior (Iterative prior resampling)
                # However, the computations are longer with a lot of models, thus you can opt-in for the "simplified" option which randomely select up to 10 times the numbers of models
                MixingUpper += 1
                MixingLower += 1
                Mixing = MixingUpper/MixingLower
                Prebel = BEL1D.PREBEL.POSTBEL2PREBEL(PREBEL=Prebel,POSTBEL=Postbel,Dataset=FieldData,NoiseModel=Noise,Simplified=True,nbMax=nbSampled,MixingRatio=Mixing)
                # Since when iterating, the dataset is known, we are not computing the full relationship but only the posterior distributions directly to gain computation timing
                print(idxIter+1)
                Postbel = BEL1D.POSTBEL(Prebel)
                Postbel.run(FieldData,nbSamples=nbSampled,NoiseModel=Noise)
                means[idxIter,:], stds[idxIter,:] = Postbel.GetStats()
                end = time.time()
                timings[idxIter] = end-start
            # The distance is computed on the normalized distributions. Therefore, the tolerance is relative.
            diverge, distance = Tools.ConvergeTest(SamplesA=ModLastIter,SamplesB=Postbel.SAMPLES, tol=tolerance)
            print('Wasserstein distance: {}'.format(distance))
            if not(diverge) or (abs((distancePrevious-distance)/distancePrevious)*100<1):# If the distance between the distributions is not changing, we converged as well
                # Convergence acheived if:
                # 1) Distance below threshold
                # 2) Distance does not vary significantly (less than 2.5%)
                print('Model has converged at iter {}!'.format(idxIter+1))
                MODELS_ITER[idxIter,:,:] = Postbel.SAMPLES
                break
            distancePrevious = distance
            MODELS_ITER[idxIter,:,:] = Postbel.SAMPLES
            start = time.time()
        timings = timings[:idxIter+1]
        means = means[:idxIter+1,:]
        stds = stds[:idxIter+1,:]
        MODELS_ITER = MODELS_ITER[:idxIter+1,:,:]
        np.save('ModelsIteration',MODELS_ITER)

        Postbel.ShowPostModels(RMSE=True)
        Postbel.ShowDataset(RMSE=True,Prior=True)
        CurrentGraph = pyplot.gcf()
        CurrentGraph = CurrentGraph.get_axes()[0]
        CurrentGraph.plot(TimingField,FieldData[:len(TimingField)],'k',linestyle='None',marker='o',markerfacecolor='None')
        pyplot.show(block=False)
        # i = int(np.ceil(idxIter/2))
        # Postbel.ShowPostCorr(OtherMethod=np.squeeze(MODELS_ITER[i,:,:]))

        print('CPU time: {} seconds'.format(np.sum(timings)))

        pyplot.show()
