import numpy as np
import pytest
from numpy import linalg
from utils import get_rstate, get_printing
import matplotlib

matplotlib.use('Agg')
from matplotlib import pyplot as plt  # noqa
import dynesty  # noqa
from dynesty import plotting as dyplot  # noqa
from dynesty import utils as dyfunc  # noqa
"""
Run a series of basic tests to check whether anything huge is broken.

"""

nlive = 500
printing = get_printing()


class Gaussian:
    def __init__(self, ndim=3, corr=.95):
        self.ndim = ndim
        self.mean = np.linspace(-1, 1, self.ndim)
        self.cov = np.identity(self.ndim)  # set covariance to identity matrix
        self.cov[self.cov ==
                 0] = corr  # set off-diagonal terms (strongly correlated)
        self.cov_inv = linalg.inv(self.cov)  # precision matrix
        self.lnorm = -0.5 * (np.log(2 * np.pi) * self.ndim +
                             np.log(linalg.det(self.cov)))
        self.prior_win = 10  # +/- 10 on both sides
        self.logz_truth = self.ndim * (-np.log(2 * self.prior_win))

    # 3-D correlated multivariate normal log-likelihood
    def loglikelihood(self, x):
        """Multivariate normal log-likelihood."""
        return -0.5 * np.dot(
            (x - self.mean), np.dot(self.cov_inv,
                                    (x - self.mean))) + self.lnorm

    # prior transform
    def prior_transform(self, u):
        """Flat prior between -10. and 10."""
        return self.prior_win * (2. * u - 1.)


@pytest.mark.parametrize("dynamic,periodic,ndim", [(False, False, 3),
                                                   (True, False, 3),
                                                   (True, True, 3),
                                                   (True, False, 1)])
def test_gaussian(dynamic, periodic, ndim):
    rstate = get_rstate()
    g = Gaussian(ndim=ndim)
    if periodic:
        periodic = [0]
    else:
        periodic = None
    if dynamic:
        sampler = dynesty.DynamicNestedSampler(g.loglikelihood,
                                               g.prior_transform,
                                               g.ndim,
                                               nlive=nlive,
                                               rstate=rstate,
                                               periodic=periodic)
    else:
        sampler = dynesty.NestedSampler(g.loglikelihood,
                                        g.prior_transform,
                                        g.ndim,
                                        nlive=nlive,
                                        rstate=rstate,
                                        periodic=periodic)
    sampler.run_nested(print_progress=printing)
    results = sampler.results
    # check plots
    dyplot.runplot(results)
    dyplot.runplot(results, fig=(plt.gcf(), plt.gcf().axes))
    plt.close()
    dyplot.traceplot(results)
    dyplot.traceplot(results,
                     fig=(plt.gcf(), plt.gcf().axes),
                     show_titles=True)
    plt.close()
    dyplot.cornerplot(results, show_titles=True, truths=[-.1, 0, .1])
    plt.close()
    if ndim != 1:
        # cornerbound
        dyplot.cornerbound(results,
                           it=500,
                           prior_transform=g.prior_transform,
                           show_live=True,
                           span=[(-10, 10), (-10, 10)])
        dyplot.cornerbound(results,
                           it=500,
                           show_live=True,
                           span=[(-10, 10), (-10, 10)],
                           fig=(plt.gcf(), plt.gcf().axes))
        plt.close()
        # boundplot
        dyplot.boundplot(results,
                         dims=(0, 1)[:min(ndim, 2)],
                         it=1000,
                         prior_transform=g.prior_transform,
                         show_live=True,
                         span=[(-10, 10), (-10, 10)])
        plt.close()

        # cornerpoints
        dyplot.cornerpoints(results)
        plt.close()
        dyplot.cornerpoints(results,
                            span=[[-10, 10], .9, [-10, 10]],
                            truths=[-0.1, 0, .1])
        plt.close()

@@ -1522,29 +1522,15 @@
 ndim
-%22, %5B(False, False, 3)
+,bound%22
 ,%0A  
@@ -1544,32 +1544,61 @@
                 
+%5B(False, False, 3, 'multi'),%0A
                 
@@ -1626,61 +1626,19 @@
 e, 3
-),%0A                                                  
+, 'multi'),
  (Tr
@@ -1648,16 +1648,25 @@
  True, 3
+, 'multi'
 ),%0A     
@@ -1690,33 +1690,8 @@
     
- 
-                        
 (Tru
@@ -1701,16 +1701,25 @@
 False, 1
+, 'multi'
 )%5D)%0Adef 
@@ -1755,16 +1755,23 @@
 ic, ndim
+, bound
 ):%0A    r
@@ -2276,24 +2276,84 @@
 dic=periodic
+,%0A                                               bound=bound
 )%0A    else:%0A
@@ -2681,16 +2681,69 @@
 periodic
+,%0A                                        bound=bound
 )%0A    sa
@@ -3098,32 +3098,173 @@
     plt.close()%0A
+%0A    truths = np.zeros(ndim)%0A    truths%5B0%5D = -.1%0A    span = %5B%5B-10, 10%5D%5D * ndim%0A    if ndim %3E 1:%0A        truths%5B1%5D = .1%0A        span%5B1%5D = .9%0A%0A
     dyplot.corne
@@ -3307,20 +3307,14 @@
 ths=
-%5B-.1, 0, .1%5D
+truths
 )%0A  
@@ -3575,38 +3575,20 @@
    span=
-%5B(-10, 10), (-10, 10)%5D
+span
 )%0A      
@@ -3727,38 +3727,20 @@
    span=
-%5B(-10, 10), (-10, 10)%5D
+span
 ,%0A      
@@ -4089,30 +4089,12 @@
 pan=
-%5B(-10, 10), (-10, 10)%5D
+span
 )%0A  
@@ -4232,118 +4232,33 @@
 lts,
-%0A                            span=%5B%5B-10, 10%5D, .9, %5B-10, 10%5D%5D,%0A                            truths=%5B-0.1, 0, .1%5D
+ span=span, truths=truths
 )%0A  

make sure I can run plotting tests over different number of dims