import numpy as np
import dynesty
import dynesty.utils as dyutil
from utils import get_rstate
"""
Run a series of basic tests changing various things like
maxcall options and potentially other things
"""

nlive = 100

size = 10  # box size


def loglike(x):
    return -0.5 * np.sum(x**2)


def loglike_inf(x):
    r2 = np.sum(x**2)
    if r2 > size * size or np.abs(x[0]) < 0.1:
        return -np.inf
    return -0.5 * r2


def prior_transform(x):
    return (2 * x - 1) * size


def test_maxcall():
    # hard test of dynamic sampler with high dlogz_init and small number
    # of live points
    ndim = 2
    rstate = get_rstate()
    sampler = dynesty.NestedSampler(loglike,
                                    prior_transform,
                                    ndim,
                                    nlive=nlive,
                                    rstate=rstate)
    sampler.run_nested(dlogz=1, maxcall=1000)

    sampler = dynesty.DynamicNestedSampler(loglike,
                                           prior_transform,
                                           ndim,
                                           nlive=nlive)
    sampler.run_nested(dlogz_init=1, maxcall=1000)


def test_inf():
    # hard test of dynamic sampler with high dlogz_init and small number
    # of live points
    ndim = 2
    rstate = get_rstate()
    sampler = dynesty.NestedSampler(loglike_inf,
                                    prior_transform,
                                    ndim,
                                    nlive=nlive,
                                    rstate=rstate)
    sampler.run_nested()

    sampler = dynesty.DynamicNestedSampler(loglike_inf,
                                           prior_transform,
                                           ndim,
                                           nlive=nlive,
                                           rstate=rstate)
    sampler.run_nested(dlogz_init=1)


def test_unravel():
    # hard test of dynamic sampler with high dlogz_init and small number
    # of live points
    ndim = 2
    rstate = get_rstate()
    sampler = dynesty.NestedSampler(loglike,
                                    prior_transform,
                                    ndim,
                                    nlive=nlive,
                                    rstate=rstate)
    sampler.run_nested()
    dyutil.unravel_run(sampler.results)

    sampler = dynesty.DynamicNestedSampler(loglike,
                                           prior_transform,
                                           ndim,
                                           nlive=nlive,
                                           rstate=rstate)
    sampler.run_nested(dlogz_init=1, maxcall=1000)
    dyutil.unravel_run(sampler.results)
    logps = sampler.results.logl
    dyutil.reweight_run(sampler.results, logps / 4.)


def test_livepoints():
    # hard test of dynamic sampler with high dlogz_init and small number
    # of live points
    ndim = 2
    rstate = get_rstate()
    live_u = rstate.uniform(size=(ndim, 2))
    live_v = np.array([prior_transform(_) for _ in live_u])
    live_logl = np.array([loglike(_) for _ in live_v])
    live_points = [live_u, live_v, live_logl]
    sampler = dynesty.NestedSampler(loglike,
                                    prior_transform,
                                    ndim,
                                    nlive=nlive,
                                    live_points=live_points,
                                    rstate=rstate)
    sampler.run_nested()
    dyutil.unravel_run(sampler.results)

@@ -12,16 +12,30 @@
 y as np%0A
+import pytest%0A
 import d
@@ -41,16 +41,16 @@
 dynesty%0A
-
 import d
@@ -72,16 +72,17 @@
  dyutil%0A
+%0A
 from uti
@@ -507,37 +507,32 @@
 maxcall():%0A    #
- hard
  test of dynamic
@@ -528,82 +528,29 @@
  of 
-dynamic sampler with high dlogz_init and small number%0A    # of live points
+maxcall functionality
 %0A   
@@ -1168,95 +1168,38 @@
   # 
-hard t
+T
 est of 
-dynamic sampler with high dlogz_init and small number%0A    # of live points
+logl that returns -inf
 %0A   
@@ -1852,95 +1852,24 @@
   # 
-hard test of dynamic sampler with high dlogz_init and small number%0A    # of live points
+test unravel_run
 %0A   
@@ -2697,95 +2697,64 @@
   # 
-hard test of dynamic sampler with high dlogz_init and small number%0A    # of
+Test the providing of initial
  live
- 
+-
 points
+ to the sampler
 %0A   
@@ -3312,32 +3312,32 @@
 er.run_nested()%0A
-
     dyutil.unrav
@@ -3340,28 +3340,630 @@
 nravel_run(sampler.results)%0A
+%0A%0Adef test_kl():%0A    # Test the providing of initial live-points to the sampler%0A    ndim = 2%0A    rstate = get_rstate()%0A    sampler1 = dynesty.NestedSampler(loglike,%0A                                     prior_transform,%0A                                     ndim,%0A                                     nlive=nlive,%0A                                     rstate=rstate)%0A    sampler1.run_nested()%0A    res1 = sampler1.results%0A    res2 = dyutil.resample_run(res1)%0A    kl = dyutil.kl_divergence(res2, res1)%0A    assert (kl%5B-1%5D %3E 0)%0A    with pytest.raises(Exception):%0A        kl = dyutil.kl_divergence(res1, res2)%0A

add a test of kl_divergence code