# Path: elfi/methods/mcmc.py
# def eff_sample_size(chains):
# def gelman_rubin_statistic(chains):
# def nuts(n_iter,
#          params0,
#          target,
#          grad_target,
#          n_adapt=None,
#          target_prob=0.6,
#          max_depth=5,
#          seed=0,
#          info_freq=100,
#          max_retry_inits=20,
#          stepsize=None):
# def _build_tree_nuts(params, momentum, log_slicevar, step, depth, log_joint0, target, grad_target,
#                      random_state):
# def metropolis(n_samples, params0, target, sigma_proposals, warmup=0, seed=0):


import numpy as np
import pytest
from elfi.methods import mcmc
        n_samples = 100000
        n_adapt = 10000
        x_init = np.random.rand(n)
        samples = mcmc.nuts(n_samples, x_init, log_pdf, grad_log_pdf, n_adapt=n_adapt)
        assert samples.shape == (n_samples, n)
        cov = np.cov(samples[n_adapt:, :].T)
        assert np.allclose(cov, true_cov, atol=0.1, rtol=0.1)


# some data generated in PyStan
chains_Stan = np.array([[0.2955857, 1.27937191, 1.05884099, 0.91236858], [
    0.38128885, 1.34242613, 0.49102573, 0.76061715
], [0.38128885, 1.18404563, 0.49102573,
    0.78910512], [0.38128885, 0.72150199, 0.49102573,
                  1.13845618], [0.38128885, 0.72150199, 0.38102685,
                                0.81298041], [0.26917982, 0.72150199, 0.38102685, 0.81298041],
                        [0.26917982, 0.68149163, 0.45830605,
                         0.86364605], [0.51213898, 0.68149163, 0.29170172, 0.80734373],
                        [0.51213898, 0.85560228, 0.29170172,
                         0.48134129], [0.22711558, 0.85560228, 0.29170172, 0.48134129]]).T

ess_Stan = 4.09
Rhat_Stan = 1.714


def test_ESS():
    assert np.isclose(mcmc.eff_sample_size(chains_Stan), ess_Stan, atol=0.01)


def test_Rhat():

# Path: elfi/methods/mcmc.py
# def eff_sample_size(chains):
# def gelman_rubin_statistic(chains):
# def nuts(n_iter,
#          params0,
#          target,
#          grad_target,
#          n_adapt=None,
#          target_prob=0.6,
#          max_depth=5,
#          seed=0,
#          info_freq=100,
#          max_retry_inits=20,
#          stepsize=None):
# def _build_tree_nuts(params, momentum, log_slicevar, step, depth, log_joint0, target, grad_target,
#                      random_state):
# def metropolis(n_samples, params0, target, sigma_proposals, warmup=0, seed=0):


# Path: tests/unit/test_mcmc.py
import numpy as np
import pytest
from elfi.methods import mcmc

        n_samples = 100000
        n_adapt = 10000
        x_init = np.random.rand(n)
        samples = mcmc.nuts(n_samples, x_init, log_pdf, grad_log_pdf, n_adapt=n_adapt)
        assert samples.shape == (n_samples, n)
        cov = np.cov(samples[n_adapt:, :].T)
        assert np.allclose(cov, true_cov, atol=0.1, rtol=0.1)


# some data generated in PyStan
chains_Stan = np.array([[0.2955857, 1.27937191, 1.05884099, 0.91236858], [
    0.38128885, 1.34242613, 0.49102573, 0.76061715
], [0.38128885, 1.18404563, 0.49102573,
    0.78910512], [0.38128885, 0.72150199, 0.49102573,
                  1.13845618], [0.38128885, 0.72150199, 0.38102685,
                                0.81298041], [0.26917982, 0.72150199, 0.38102685, 0.81298041],
                        [0.26917982, 0.68149163, 0.45830605,
                         0.86364605], [0.51213898, 0.68149163, 0.29170172, 0.80734373],
                        [0.51213898, 0.85560228, 0.29170172,
                         0.48134129], [0.22711558, 0.85560228, 0.29170172, 0.48134129]]).T

ess_Stan = 4.09
Rhat_Stan = 1.714


def test_ESS():
    assert np.isclose(mcmc.eff_sample_size(chains_Stan), ess_Stan, atol=0.01)


def test_Rhat():

    assert np.isclose(mcmc.gelman_rubin_statistic(chains_Stan), Rhat_Stan, atol=0.01)