import os
import numpy as np
import shutil
import tempfile
import pytest
import json
import pandas
from ultranest.mlfriends import MLFriends, ScalingLayer, AffineLayer, MaxPrincipleGapAffineLayer
from ultranest import NestedSampler, ReactiveNestedSampler, read_file
from ultranest.integrator import warmstart_from_similar_file, _update_region_bootstrap, _get_cumsum_range
import ultranest.mlfriends
from numpy.testing import assert_allclose

def sample_ellipsoid(rng, nsamples, ndim, sigma=0.01, center=0.5):
    """ Sample from a unit sphere with constant density """
    z = rng.normal(size=(nsamples, ndim))
    z /= ((z**2).sum(axis=1)**0.5).reshape((nsamples, 1))
    u = z * rng.uniform(size=(nsamples, 1))**(1./ndim)
    return u * sigma + center

def generate_two_blob_points(rng, d, Nlive1, Nlive2, offset2, sigma):
    """ generate live points from two spheres """
    return np.vstack((sample_ellipsoid(rng, Nlive1, d, sigma=sigma), sample_ellipsoid(rng, Nlive2, d, sigma=sigma) + offset2))

def test_get_cumsum_range1():
    # cumulative probabilities are: array([0.1, 0.3, 0.6, 1. ])
    pi = np.array([0.1, 0.2, 0.3, 0.4])
    dp = 0.2
    ilo, ihi = _get_cumsum_range(pi, dp)
    assert ilo == 1, ilo
    assert ihi == 2, ihi

def test_get_cumsum_range_equal_prob():
    p = np.ones(100) * 1.0
    p = p * 1. / p.sum()
    print(p, p.sum())
    for percentile in 1, 5, 10, 20, 45:
        ilo, ihi = _get_cumsum_range(p, percentile / 100.)
        print(percentile, ilo, ihi, np.cumsum(p))
        print(np.cumsum(p)[ilo], np.cumsum(p)[ihi])
        # due to rounding issues, can slip to a lower index
        assert ilo in (percentile, percentile - 1)
        assert ihi in (100 - percentile - 1, 100 - percentile - 2)
        assert np.cumsum(p)[ilo] >= percentile / 100.
        assert np.cumsum(p)[ihi] <= 1 - percentile / 100.
        assert p[ilo:ihi].sum() <= (1 - percentile / 100.) * 2, (p[ilo:ihi], p[ilo:ihi].sum(), percentile)


def test_get_cumsum_range_random_prob():
    np.random.seed(100)
    for i in range(100):
        size = int(10**np.random.uniform(0, 4))
        p = np.random.uniform(size=size)
        p = p * 1. / p.sum()
        dp = np.random.uniform()
        ilo, ihi = _get_cumsum_range(p, dp)
        print(dp, p, np.cumsum(p), '-->', ilo, ihi, np.cumsum(p)[ilo], np.cumsum(p)[ihi])
        # check that the selected interval contains the desired probability
        assert p[ilo:ihi].sum() <= (1 - dp) * 2, (p[ilo:ihi], p[ilo:ihi].sum(), (1 - dp) * 2)


def test_failing_update_region_bootstrap():
    rng = np.random.RandomState(10)
    u = rng.uniform(size=(200, 4))
    # make linearly dependent, so building a region should fail
    u[:,-1] = u[:,0]
    # boot-strap an affine layer after clustering
    transformLayer = ScalingLayer()
    transformLayer.optimize(u, u)
    region = MLFriends(u, transformLayer)
    with pytest.raises(np.linalg.LinAlgError):
        _update_region_bootstrap(region, nbootstraps=30)


def test_clustering_recursion(plot=False):
    # generate two blobs separated by 2 sigma
    # check that they are *not* separated with AffineLayer+MLFriends
    # check that they are separated with MaxPrincipleGapAffineLayer+MLFriends
    nbootstraps = 30

    sigma = 0.001
    d = 20
    
    Nlive = 100
    Nlive1 = Nlive // 2
    Nlive2 = Nlive // 2
    
    offset2 = 1.0 * sigma
    
    nwithclusters = 0
    nwithclusters2 = 0
    noverclustered = 0
    gapped_nwithclusters = 0
    gapped_nwithclusters2 = 0
    gapped_noverclustered = 0

    for seed in range(25):
        rng = np.random.RandomState(54 + seed)
        u = generate_two_blob_points(rng, d, Nlive1, Nlive2, offset2, sigma)

        # boot-strap an affine layer after clustering
        transformLayer = AffineLayer()
        transformLayer.optimize(u, u)
        region = MLFriends(u, transformLayer)
        _update_region_bootstrap(region, nbootstraps)
        region.create_ellipsoid()
        layer = transformLayer.create_new(u, region.maxradiussq)
        nextregion = MLFriends(u, layer)
        _update_region_bootstrap(nextregion, nbootstraps=30)
        nextLayer = layer.create_new(u, nextregion.maxradiussq)
        nextNextLayer = nextLayer.create_new(u, region.maxradiussq)

        nwithclusters += layer.nclusters
        nwithclusters2 += nextLayer.nclusters
        noverclustered += nextLayer.nclusters > 2
        if plot:
            import matplotlib.pyplot as plt
            plt.figure("AffineLayer", figsize=(20, 20))
            plt.subplot(5, 5, seed + 1)
            plt.title('%d -> %d -> %d' % (transformLayer.nclusters, nextLayer.nclusters, nextNextLayer.nclusters))
            plt.scatter(u[:,0], u[:,1], label='points')
            # Plot the principal vectors
            plt.quiver(0.5, 0.5, transformLayer.invT[0, 0], transformLayer.invT[0, 1], angles='xy', scale_units='xy', scale=1, color='r', label='First Principal Vector')
            plt.quiver(0.5, 0.5, transformLayer.invT[1, 0], transformLayer.invT[1, 1], angles='xy', scale_units='xy', scale=1, color='b', label='Second Principal Vector')
            ylo, yhi = plt.ylim()
            ymax = max(0.5 - ylo, yhi - 0.5)
            xlo, xhi = plt.xlim()
            xmax = max(0.5 - xlo, xhi - 0.5)
            xymax = 1.5 * max(ymax, xmax)
            plt.xlim(0.5 - xymax, 0.5 + xymax)
            plt.ylim(0.5 - xymax, 0.5 + xymax)

        # boot-strap an MaxPrincipleGapAffineLayer layer after clustering
        transformLayer = MaxPrincipleGapAffineLayer()
        transformLayer.optimize(u, u)
        region = MLFriends(u, transformLayer)
        _update_region_bootstrap(region, nbootstraps)
        region.create_ellipsoid()
        layer = transformLayer.create_new(u, region.maxradiussq)
        nextregion = MLFriends(u, layer)
        _update_region_bootstrap(nextregion, nbootstraps=30)
        nextLayer = layer.create_new(u, nextregion.maxradiussq)
        nextNextLayer = nextLayer.create_new(u, region.maxradiussq)

        gapped_nwithclusters += layer.nclusters
        gapped_nwithclusters2 += nextLayer.nclusters
        gapped_noverclustered += nextLayer.nclusters > 2
        if plot:
            plt.figure("MaxPrincipleGapAffineLayer", figsize=(20, 20))
            plt.subplot(5, 5, seed + 1)
            plt.title('%d -> %d -> %d' % (transformLayer.nclusters, nextLayer.nclusters, nextNextLayer.nclusters))
            plt.scatter(u[:,0], u[:,1], label='points')
            # Plot the principal vectors
            plt.quiver(0.5, 0.5, transformLayer.invT[0, 0], transformLayer.invT[0, 1], angles='xy', scale_units='xy', scale=1, color='r', label='First Principal Vector')
            plt.quiver(0.5, 0.5, transformLayer.invT[1, 0], transformLayer.invT[1, 1], angles='xy', scale_units='xy', scale=1, color='b', label='Second Principal Vector')
            ylo, yhi = plt.ylim()
            ymax = max(0.5 - ylo, yhi - 0.5)
            xlo, xhi = plt.xlim()
            xmax = max(0.5 - xlo, xhi - 0.5)
            xymax = 1.5 * max(ymax, xmax)
            plt.xlim(0.5 - xymax, 0.5 + xymax)
            plt.ylim(0.5 - xymax, 0.5 + xymax)

    if plot:
        plt.savefig('layercov_MaxPrincipleGapAffineLayer.pdf')
        plt.close()
        plt.savefig('layercov_AffineLayer.pdf')
        plt.close()

    print("clustering statistics: (%d runs)" % (seed+1))
    print("  number of clusters iteration 1, iteration 2, number of overclusterings")
    print("AffineLayer:")
    print("  ", nwithclusters, nwithclusters2, noverclustered)
    print("MaxPrincipleGapAffineLayer:")
    print("  ", gapped_nwithclusters, gapped_nwithclusters2, gapped_noverclustered)
    # with the affine layer we only see one cluster, because they are
    # so close together and the covariance spans them
    assert nwithclusters in (25, 26, 27)
    assert nwithclusters2 in (25, 26, 27)
    # MaxPrincipleGapAffineLayer builds a more local covariance
    # so the subsequent iteration splits the cluster
    assert gapped_nwithclusters in (25, 26, 27)
    assert gapped_nwithclusters2 in (49, 50, 51, 52)
    assert noverclustered in (0, 1, 2)
    assert gapped_noverclustered in (0, 1, 2)

def test_run():
    def loglike(y):
        z = np.log10(y)
        a = np.array([-0.5 * sum([((xi - 0.83456 + i*0.1)/0.5)**2 for i, xi in enumerate(x)]) for x in z])
        b = np.array([-0.5 * sum([((xi - 0.43456 - i*0.1)/0.5)**2 for i, xi in enumerate(x)]) for x in z])
        loglike.ncalls += len(a)
        return np.logaddexp(a, b)
    loglike.ncalls = 0

    def transform(x):
        return 10**(10. * x - 5.)

    paramnames = ['Hinz', 'Kunz']

    sampler = NestedSampler(paramnames, loglike, transform=transform, num_live_points=400, vectorized=True)
    r = sampler.run(log_interval=50)

    ncalls = loglike.ncalls
    if sampler.mpi_size > 1:
        ncalls = sampler.comm.gather(ncalls, root=0)
        if sampler.mpi_rank == 0:
            print("ncalls on the different MPI ranks:", ncalls)
        ncalls = sum(sampler.comm.bcast(ncalls, root=0))
    assert abs(r['ncall'] - ncalls) <= 2 * sampler.mpi_size, (r['ncall'], ncalls)
    open('nestedsampling_results.txt', 'a').write("%.3f\n" % r['logz'])
    sampler.plot()


def test_dlogz_reactive_run_SLOW():
    def loglike(y):
        return -0.5 * np.sum(((y - 0.5)/0.001)**2, axis=1)

    paramnames = ['Hinz', 'Kunz']

    sampler = ReactiveNestedSampler(paramnames, loglike, vectorized=True)
    print("running for ess")
    firstresults = sampler.run(min_num_live_points=50, cluster_num_live_points=0, max_num_improvement_loops=3, min_ess=10000, viz_callback=None)
    print()
    print({k:v for k, v in firstresults.items() if 'logzerr' in k})
    print()
    assert firstresults['logzerr'] > 0.1 * 2
    print("running again for logz")
    for niter, results in enumerate(sampler.run_iter(min_num_live_points=1, cluster_num_live_points=0, max_num_improvement_loops=10, dlogz=0.1, viz_callback=None, region_class=ultranest.mlfriends.RobustEllipsoidRegion)):
        print("logzerr in iteration %d" % niter, results['logzerr'])
    print()
    print({k:v for k, v in results.items() if 'logzerr' in k})
    assert results['logzerr'] < 0.1 * 3

def test_reactive_run():
    np.random.seed(1)
    evals = set()

    def loglike(z):
        #print(loglike.ncalls, z[0,0])
        [evals.add(str(x[0])) for x in z]
        a = np.array([-0.5 * sum([((xi - 0.83456 + i*0.1)/0.5)**2 for i, xi in enumerate(x)]) for x in z])
        b = np.array([-0.5 * sum([((xi - 0.43456 - i*0.1)/0.5)**2 for i, xi in enumerate(x)]) for x in z])
        loglike.ncalls += len(a)
        return np.logaddexp(a, b)
    loglike.ncalls = 0

    def transform(x):
        return 10. * x - 5.

    paramnames = ['Hinz', 'Kunz']

    sampler = ReactiveNestedSampler(paramnames, loglike, transform=transform,
        draw_multiple=False, vectorized=True)
    r = sampler.run(log_interval=50, min_num_live_points=400)

    # test that the number of likelihood calls is correct

    ncalls = loglike.ncalls
    nunique = len(evals)
    if sampler.mpi_size > 1:
        ncalls = sampler.comm.gather(ncalls, root=0)
        if sampler.mpi_rank == 0:
            print("ncalls on the different MPI ranks:", ncalls)
        ncalls = sum(sampler.comm.bcast(ncalls, root=0))

        allevals = sampler.comm.gather(evals, root=0)
        if sampler.mpi_rank == 0:
            print("evals on the different MPI ranks:", [len(e) for e in allevals])
            allevals = len(set.union(*allevals))
        else:
            allevals = None
        nunique = sampler.comm.bcast(allevals, root=0)

    if sampler.mpi_rank == 0:
        print('ncalls:', ncalls, 'nunique:', nunique)

    assert abs(r['ncall'] - ncalls) <= 2 * sampler.mpi_size, (r['ncall'], ncalls)
    assert ncalls == nunique, (ncalls, nunique)
    if sampler.mpi_rank == 0:
        open('nestedsampling_reactive_results.txt', 'a').write("%.3f\n" % r['logz'])

    print(r)
    assert r['niter'] > 100
    assert -10 < r['logz'] < 10
    assert 0.01 < r['logzerr'] < 0.5
    assert 1 < r['ess'] < 10000


    sampler.plot()


def test_plateau_SLOW():
    def loglike(y):
        a = -0.5 * ((y/0.1)**2).sum()
        if a < -1:
            return -1e100
        return a

    def transform(x):
        return x * 2 - 1

    paramnames = ['Hinz', 'Kunz']

    sampler = ReactiveNestedSampler(paramnames, loglike, transform=transform)
    print(sampler.run(min_num_live_points=400)['logz'])
    print(sampler.run_sequence['nlive'][:-400])
    assert sampler.run_sequence['nlive'][-400] == 400, sampler.run_sequence['nlive'][:-400]

def test_flat():
    def loglike(y):
        return 0
    paramnames = ['Hinz', 'Kunz']

    sampler = ReactiveNestedSampler(paramnames, loglike)
    print(sampler.run(min_num_live_points=400)['logz'])
    print(sampler.run_sequence['nlive'][:-400])


def test_reactive_run_extraparams():
    np.random.seed(1)

    def loglike(z):
        return -0.5 * z[-1].sum()
    loglike.ncalls = 0

    def transform(x):
        z = 10. * x - 5.
        return np.append(z, np.abs(z).sum())

    paramnames = ['Hinz', 'Kunz']

    sampler = ReactiveNestedSampler(paramnames, loglike, transform=transform,
        derived_param_names=['ctr_distance'])
    r = sampler.run()
    assert r['samples'].shape[1] == 3
    sampler.plot()

def test_return_summary():
    sigma = np.array([0.1, 0.01])
    centers = np.array([0.5, 0.75])
    paramnames = ['a', 'b']
    ndim = len(paramnames)

    def loglike(theta):
        like = -0.5 * (((theta - centers)/sigma)**2) - 0.5 * np.log(2 * np.pi * sigma**2) * ndim
        return like.sum()

    def transform(x):
        return x

    sampler = ReactiveNestedSampler(paramnames, loglike, transform=transform)
    r = sampler.run()

    print(r)
    assert r['paramnames'] == paramnames
    assert r['niter'] > 100
    assert -10 < r['logz'] < 10
    assert 0.01 < r['logzerr'] < 0.5
    assert 1 < r['ess'] < 10000
    assert 0.4 < r['posterior']['mean'][0] < 0.6
    assert 0.74 < r['posterior']['mean'][1] < 0.76
    assert 0.4 < r['posterior']['median'][0] < 0.6
    assert 0.74 < r['posterior']['median'][1] < 0.76
    assert 0.05 < r['posterior']['stdev'][0] < 0.2
    assert 0.005 < r['posterior']['stdev'][1] < 0.02

    assert 0.35 < r['posterior']['errlo'][0] < 0.45
    assert 0.72 < r['posterior']['errlo'][1] < 0.75
    assert 0.55 < r['posterior']['errup'][0] < 0.65
    assert 0.75 < r['posterior']['errup'][1] < 0.78

    N, ndim2 = r['samples'].shape
    assert ndim2 == ndim
    assert N > 10
    N, ndim2 = r['weighted_samples']['points'].shape
    assert ndim2 == ndim
    assert N > 10

    assert r['weighted_samples']['logw'].shape == (N,)
    assert r['weighted_samples']['weights'].shape == (N,)
    assert r['weighted_samples']['bootstrapped_weights'].shape[0] == N
    assert r['weighted_samples']['logl'].shape == (N,)

@pytest.mark.parametrize("dlogz", [2.0, 0.5, 0.1])
def test_run_resume(dlogz):
    sigma = 0.01
    ndim = 1

    def loglike(theta):
        like = -0.5 * (((theta - 0.5)/sigma)**2).sum(axis=1) - 0.5 * np.log(2 * np.pi * sigma**2) * ndim
        return like

    def transform(x):
        return x

    paramnames = ['a']
    def myadd(row):
        assert False, (row, 'should not need to add more points in resume')

    last_results = None
    #for dlogz in 0.5, 0.1, 0.01:
    np.random.seed(int(dlogz*100))
    folder = tempfile.mkdtemp()
    try:
        for i in range(2):
            sampler = ReactiveNestedSampler(paramnames, loglike, transform=transform,
                log_dir=folder, resume=True, vectorized=True)
            r = sampler.run(log_interval=50, dlogz=dlogz, min_num_live_points=400)
            sampler.print_results()
            sampler.pointstore.close()
            if i == 1:
                sampler.pointstore.add = myadd
            del r['weighted_samples']
            del r['samples']
            if last_results is not None:
                print("ran with dlogz:", dlogz)
                print("first run gave:", last_results)
                print("second run gave:", r)
                assert last_results['logzerr'] < 1.0
                assert r['logzerr'] < 1.0
                assert np.isclose(last_results['logz'], r['logz'], atol=0.5)
            last_results = r
    finally:
        shutil.rmtree(folder, ignore_errors=True)

@pytest.mark.parametrize("storage_backend", ['hdf5', 'tsv', 'csv'])
def test_reactive_run_resume_eggbox(storage_backend):
    def loglike(z):
        chi = (np.cos(z / 2.)).prod(axis=1)
        loglike.ncalls += len(z)
        return (2. + chi)**5
    loglike.ncalls = 0

    def transform(x):
        return x * 10 * np.pi

    paramnames = ['a', 'b']
    ndim = len(paramnames)

    #last_results = None
    folder = tempfile.mkdtemp()
    np.random.seed(1)
    try:
        for i in range(2):
            print()
            print("====== Running Eggbox problem [%d] =====" % (i+1))
            print()
            sampler = ReactiveNestedSampler(paramnames,
                loglike, transform=transform,
                log_dir=folder, resume=True, vectorized=True, draw_multiple=False,
                storage_backend=storage_backend)
            initial_ncalls = int(sampler.ncall)
            num_live_points = 100
            loglike.ncalls = 0
            r = sampler.run(max_iters=200 + i*200,
                max_num_improvement_loops=0,
                min_num_live_points=num_live_points,
                cluster_num_live_points=0)
            sampler.print_results()
            if storage_backend == 'hdf5':
                print("pointstore:", sampler.pointstore.fileobj['points'].shape)
            sampler.pointstore.close()
            print(loglike.ncalls, r['ncall'], initial_ncalls)

            ncalls = loglike.ncalls
            if sampler.mpi_size > 1:
                ncalls = sampler.comm.gather(ncalls, root=0)
                if sampler.mpi_rank == 0:
                    print("ncalls on the different MPI ranks:", ncalls)
                ncalls = sum(sampler.comm.bcast(ncalls, root=0))
            ncalls = ncalls + initial_ncalls
            assert abs(r['ncall'] - ncalls) <= 2 * sampler.mpi_size, (i, r['ncall'], ncalls, r['ncall'] - ncalls)
            assert paramnames == r['paramnames'], 'paramnames should be in results'

            results2 = json.load(open(folder + '/info/results.json'))
            print('CSV content:')
            print(open(folder + '/info/post_summary.csv').read())
            post_summary = pandas.read_csv(folder + '/info/post_summary.csv')
            print(post_summary, post_summary.columns)
            for k, v in r.items():
                if k in results2:
                    print("checking results[%s] ..." % k)
                    assert results2[k] == r[k], (k, results2[k], r[k])
            
            assert r['paramnames'] == paramnames
            samples = np.loadtxt(folder + '/chains/equal_weighted_post.txt', skiprows=1)
            data = np.loadtxt(folder + '/chains/weighted_post.txt', skiprows=1)
            data_u = np.loadtxt(folder + '/chains/weighted_post_untransformed.txt', skiprows=1)
            assert (data[:,:2] == data_u[:,:2]).all()
            
            assert_allclose(samples.mean(axis=0), r['posterior']['mean'])
            assert_allclose(np.median(samples, axis=0), r['posterior']['median'])
            assert_allclose(np.std(samples, axis=0), r['posterior']['stdev'])
            for k, v in r.items():
                if k == 'posterior':
                    for k1, v1 in v.items():
                        if k1 == 'information_gain_bits':
                            continue
                        for param, value in zip(paramnames, v[k1]):
                            print("checking %s of parameter '%s':" % (k1, param), value)
                            assert np.isclose(post_summary[param + '_' + k1].values, value), (param, k1, post_summary[param + '_' + k1].values, value)
                elif k == 'samples':
                    assert_allclose(samples, r['samples'])
                elif k == 'paramnames':
                    assert v == paramnames
                elif k == 'weighted_samples':
                    print(k, v.keys())
                    assert_allclose(data[:,0], v['weights'])
                    assert_allclose(data[:,1], v['logl'])
                    assert_allclose(data[:,2:], v['points'])
                    assert_allclose(data_u[:,2:], v['upoints'])
                elif k == 'maximum_likelihood':
                    print(k, v.keys())
                    assert_allclose(data[-1,1], v['logl'])
                    assert_allclose(data[-1,2:], v['point'])
                    assert_allclose(data_u[-1,2:], v['point_untransformed'])
                    
                elif k.startswith('logzerr') or '_bs' in k or 'Herr' in k:
                    print("   skipping", k, np.shape(v))
                    #assert_allclose(r[k], v, atol=0.5)
                elif k == 'insertion_order_MWW_test':
                    print('insertion_order_MWW_test:', r[k], v)
                    assert r[k] == v, (r[k], v)
                else:
                    print("  ", k, np.shape(v))
                    assert_allclose(r[k], v)

            logw = r['weighted_samples']['logw']
            v = r['weighted_samples']['points']
            L = r['weighted_samples']['logl']

            assert results2['niter'] == len(r['samples'])

        # the results are not exactly the same, because the sampling adds
        #ncalls = loglike.ncalls
        #sampler = ReactiveNestedSampler(paramnames,
        #    loglike, transform=transform,
        #    log_dir=folder, resume=True, vectorized=True, num_test_samples=0)
        #print("pointstore:", sampler.pointstore.fileobj['points'].shape)
        #assert ncalls == loglike.ncalls, (ncalls, loglike.ncalls)
        if storage_backend == 'hdf5':
            sequence, results = read_file(folder, ndim, random=False, num_bootstraps=0)

            print("sampler results: ********************")
            print({k:v for k, v in r.items() if np.asarray(v).size < 20 and k != 'weighted_samples'})
            print("reader results: ********************")
            print({k:v for k, v in results.items() if np.asarray(v).size < 20 and k != 'weighted_samples'})
            for k, v in results.items():
                if k == 'posterior' or k == 'samples':
                    pass
                elif k == 'weighted_samples' or k == 'maximum_likelihood':
                    for k2, v2 in results[k].items():
                        if k2 == 'bootstrapped_weights': continue
                        print("  ", k, "::", k2, np.shape(v2))
                        assert_allclose(r[k][k2], v2)
                elif k.startswith('logzerr') or '_bs' in k or 'Herr' in k:
                    print("   skipping", k, np.shape(v))
                    #assert_allclose(r[k], v, atol=0.5)
                elif k == 'insertion_order_MWW_test':
                    print('insertion_order_MWW_test:', r[k], v)
                    assert r[k] == v, (r[k], v)
                else:
                    print("  ", k, np.shape(v))
                    assert_allclose(r[k], v)

            logw = r['weighted_samples']['logw']
            v = r['weighted_samples']['points']
            L = r['weighted_samples']['logl']

            assert sequence['logz'][-1] - r['logz'] < 0.5, (results['logz'][-1], r['logz'])
            assert sequence['logzerr'][-1] <= r['logzerr_single'], (results['logzerr'][-1], r['logzerr'])
            #assert_allclose(sequence['logz_final'], r['logz_single'], atol=0.3)
            #assert_allclose(sequence['logzerr_final'], r['logzerr_single'], atol=0.1)
            assert r['niter'] <= sequence['niter'] <= r['niter'], (sequence['niter'], r['niter'])
            assert results['niter'] == len(sequence['logz']) == len(sequence['logzerr']) == len(sequence['logvol']) == len(sequence['logwt'])
            assert results['niter'] == len(results['samples'])
            data = np.loadtxt(folder + '/chains/weighted_post.txt', skiprows=1)
            assert_allclose(data[:,0], results['weighted_samples']['weights'])
            assert_allclose(data[:,1], results['weighted_samples']['logl'])
            assert_allclose(v, results['weighted_samples']['points'])
            assert_allclose(logw, results['weighted_samples']['logw'])
            assert_allclose(L, results['weighted_samples']['logl'])

            assert_allclose(L, sequence['logl'])
            #assert_allclose(logw + L, sequence['logwt'])
            assert sequence['logvol'].shape == logw.shape == (len(L),), (sequence['logvol'].shape, logw.shape)
            assert sequence['logwt'].shape == logw.shape == (len(L),), (sequence['logwt'].shape, logw.shape)
            #assert_allclose(logw, sequence['logvols'])
            #assert results['samples_untransformed'].shape == v.shape == (len(L), ndim), (results['samples_untransformed'].shape, v.shape)

    finally:
        shutil.rmtree(folder, ignore_errors=True)

def test_reactive_run_warmstart_gauss():
    center = 0

    def loglike(z):
        chi2 = (((z - center)/0.001)**2).sum(axis=1)
        loglike.ncalls += len(z)
        return -0.5 * chi2
    loglike.ncalls = 0

    def transform(x):
        return x * 20000 - 10000

    paramnames = ['a']

    folder = tempfile.mkdtemp()
    np.random.seed(1)
    first_ncalls = None
    resume_ncalls = None
    try:
        for i, resume in enumerate(['overwrite', 'resume', 'resume-similar']):
            print()
            print("====== Running Gauss problem [%d] =====" % (i+1))
            print()
            center = (i+1) * 1e-4
            try:
                sampler = ReactiveNestedSampler(paramnames,
                    loglike, transform=transform,
                    log_dir=folder, resume=resume, vectorized=True, draw_multiple=False,
                    warmstart_max_tau=0.5)
            except Exception as e:
                # we expect an error for resuming with a changed likelihood
                if resume != 'resume':
                    raise e
                else:
                    assert 'loglikelihood function changed' in str(e), e
                    print("Exception as expected:", e)
                    continue
            initial_ncalls = int(sampler.ncall)
            if i == 0:
                assert initial_ncalls == 0
            num_live_points = 100
            loglike.ncalls = 0
            r = sampler.run(
                max_num_improvement_loops=0,
                min_num_live_points=num_live_points,
                cluster_num_live_points=0, viz_callback=None, frac_remain=
                0.5)
            sampler.print_results()
            print("pointstore:", sampler.pointstore.fileobj['points'].shape)
            sampler.pointstore.close()
            print(loglike.ncalls, r['ncall'], initial_ncalls)

            ncalls = loglike.ncalls
            if sampler.mpi_size > 1:
                ncalls = sampler.comm.gather(ncalls, root=0)
                if sampler.mpi_rank == 0:
                    print("ncalls on the different MPI ranks:", ncalls)
                ncalls = sum(sampler.comm.bcast(ncalls, root=0))
            ncalls = ncalls + initial_ncalls
            if i == 0:
                first_ncalls = ncalls
            if i == 2:
                resume_ncalls = loglike.ncalls
            assert abs(r['ncall'] - ncalls) <= 2 * sampler.mpi_size, (i, r['ncall'], ncalls, r['ncall'] - ncalls)
            assert paramnames == r['paramnames'], 'paramnames should be in results'

    finally:
        shutil.rmtree(folder, ignore_errors=True)
    
    # make sure warm start is much faster
    assert resume_ncalls < first_ncalls - 800, (resume_ncalls, first_ncalls)

def test_run_compat():
    from ultranest.solvecompat import pymultinest_solve_compat as solve

    ndim = 2
    sigma = 0.01
    centers = 0.5
    paramnames = ['a', 'b']

    def loglike(theta):
        like = -0.5 * (((theta - centers)/sigma)**2).sum() - 0.5 * np.log(2 * np.pi * sigma**2) * ndim
        return like

    def transform(x):
        params = x.copy()
        params[0] = 10 * x[0] - 5.
        params[1] = 10**(x[1] - 1)
        return params

    result = solve(LogLikelihood=loglike, Prior=transform,
        n_dims=ndim, outputfiles_basename=None,
        verbose=True, resume='resume', importance_nested_sampling=False)

    print()
    print('evidence: %(logZ).1f +- %(logZerr).1f' % result)
    print()
    print('parameter values:')
    for name, col in zip(paramnames, result['samples'].transpose()):
        print('%15s : %.3f +- %.3f' % (name, col.mean(), col.std()))


def test_run_warmstart_gauss_SLOW():
    center = None
    stdev = 0.001

    def loglike(z):
        chi2 = (((z - center) / stdev)**2).sum(axis=1)
        loglike.ncalls += len(z)
        return -0.5 * chi2
    loglike.ncalls = 0

    def transform(x):
        return x * 20000 - 10000

    paramnames = ['a']

    folder = tempfile.mkdtemp()
    np.random.seed(1)
    ncalls = []
    try:
        for i, resume in enumerate(['overwrite', 'resume-hot', 'resume-hot', 'resume-hot']):
            print()
            print("====== Running Gauss problem [%d] =====" % (i+1))
            print()
            center = [0, 0, stdev, 1][i]
            print("center:", center, "folder:", folder)
            if i == 0:
                sampler = ReactiveNestedSampler(paramnames,
                    loglike, transform=transform,
                    log_dir=folder, resume=resume, vectorized=True)
            else:
                aux_param_names, aux_loglike, aux_transform, vectorized = warmstart_from_similar_file(
                    os.path.join(folder, 'chains', 'weighted_post_untransformed.txt'),
                    paramnames, loglike=loglike, transform=transform, vectorized=True,
                )
                sampler = ReactiveNestedSampler(aux_param_names,
                    aux_loglike, transform=aux_transform, vectorized=True)

            sampler.run(viz_callback=None)
            sampler.print_results()
            print("expected posterior:", center, '+-', stdev)
            print(sampler.results.keys())
            print(sampler.results['posterior'].keys())
            print(sampler.results['posterior']['mean'], sampler.results['posterior']['stdev'])
            print(sampler.results['weighted_samples']['upoints'], sampler.results['weighted_samples']['weights'])
            assert center - stdev < sampler.results['posterior']['mean'][0] < center + stdev, (center, sampler.results['posterior'])
            assert stdev * 0.8 < sampler.results['posterior']['stdev'][0] < stdev * 1.2, (center, sampler.results['posterior'])
            ncalls.append(sampler.ncall)
    finally:
        shutil.rmtree(folder, ignore_errors=True)
    print(ncalls)
    
    # make sure hot start is much faster
    assert ncalls[1] < ncalls[0] - 800, (ncalls)
    assert ncalls[2] < ncalls[0] - 800, (ncalls)

if __name__ == '__main__':
    #test_run_compat()
    #test_run_resume(dlogz=0.5)
    #test_reactive_run_resume(dlogz=0.5, min_ess=1000)
    #test_reactive_run()
    #test_run()
    #test_reactive_run_warmstart_gauss()
    #test_reactive_run_extraparams()
    #test_reactive_run_resume_eggbox('hdf5')
    #test_dlogz_reactive_run()
    #test_plateau()
    test_clustering_recursion(plot=True)
