import math
import pickle

import h5py
import numpy as np
import pytest

from aspire.samples import (
    BaseSamples,
    MCMCSamples,
    PTMCMCSamples,
    Samples,
    SMCSamples,
)


def make_simple_samples(n=5, d=2, a=1.234):
    """
    Construct a simple, consistent set of arrays:
    - x: shape (n, d)
    - log_likelihood = a (constant)
    - log_prior = 0
    - log_q = 0
    This makes the math easy: evidence = exp(a), ESS = n, etc.
    """
    x = np.arange(n * d, dtype=float).reshape(n, d)
    ll = np.full(n, a, dtype=float)
    lp = np.zeros(n, dtype=float)
    lq = np.zeros(n, dtype=float)
    return x, ll, lp, lq


@pytest.fixture
def base_samples(xp):
    x = np.arange(6.0).reshape(3, 2)
    ll = np.array([0.1, 0.2, 0.3])
    return BaseSamples(x=x, log_likelihood=ll, parameters=["a", "b"], xp=xp)


def test_basesamples_defaults_and_dims_numpy_namespace():
    x = np.zeros((3, 4))
    bs = BaseSamples(x=x)
    # xp inferred from x
    assert bs.xp is not None
    assert bs.dtype == x.dtype
    assert bs.x.dtype == x.dtype
    # parameters default names
    assert bs.parameters == ["x_0", "x_1", "x_2", "x_3"]
    # dims computed correctly
    assert bs.dims == 4
    # device set for numpy
    assert bs.device is None
    # __len__
    assert len(bs) == 3


def test_basesamples_to_dict_flat_and_nested(base_samples):
    bs = base_samples
    assert bs.log_likelihood.dtype == bs.dtype

    d_flat = bs.to_dict(flat=True)
    assert "a" in d_flat and "b" in d_flat
    assert "log_likelihood" in d_flat and np.allclose(
        d_flat["log_likelihood"], bs.log_likelihood
    )

    d_nested = bs.to_dict(flat=False)
    assert "samples" in d_nested and isinstance(d_nested["samples"], dict)
    assert set(d_nested["samples"].keys()) == {"a", "b"}


def test_basesamples_to_numpy_and_namespace_roundtrip():
    x = np.arange(6.0).reshape(3, 2)
    ll = np.array([0.1, 0.2, 0.3])
    bs = BaseSamples(x=x, log_likelihood=ll, parameters=["a", "b"])

    bs_np = bs.to_numpy()
    assert isinstance(bs_np.x, np.ndarray)
    assert isinstance(bs_np.log_likelihood, np.ndarray)

    # Convert to the same namespace (numpy) explicitly
    bs_ns = bs.to_namespace(np)
    assert isinstance(bs_ns.x, np.ndarray)
    assert isinstance(bs_ns.log_likelihood, np.ndarray)


def test_basesamples_getitem_slice_returns_same_type_and_fields():
    x = np.arange(12.0).reshape(6, 2)
    ll = np.linspace(0, 1, 6)
    lp = np.linspace(1, 2, 6)
    lq = np.linspace(-1, 0, 6)
    bs = BaseSamples(
        x=x, log_likelihood=ll, log_prior=lp, log_q=lq, parameters=["a", "b"]
    )

    one = bs[2]
    assert isinstance(one, BaseSamples.__mro__[0]) or isinstance(
        one, BaseSamples
    )  # basic sanity
    assert one.x.shape == (2,)  # row
    assert one.parameters == ["a", "b"]
    assert np.isclose(one.log_likelihood, ll[2])
    assert np.isclose(one.log_prior, lp[2])
    assert np.isclose(one.log_q, lq[2])
    assert one.dtype == bs.dtype


def test_basesamples_respects_explicit_dtype_string():
    x = np.arange(6).reshape(3, 2)
    ll = np.linspace(0.0, 1.0, 3)
    bs = BaseSamples(x=x, log_likelihood=ll, dtype="float32")
    assert bs.x.dtype == np.dtype("float32")
    assert bs.dtype == np.dtype("float32")
    assert bs.log_likelihood.dtype == np.dtype("float32")


def test_basesamples_pickle_restores_namespace_and_fields():
    x = np.arange(6.0).reshape(3, 2)
    bs = BaseSamples(x=x)
    blob = pickle.dumps(bs)
    bs2 = pickle.loads(blob)
    # xp restored from x
    assert bs2.xp is not None
    # functionality intact
    assert bs2.dims == 2
    assert np.allclose(bs2.x, x)


def test_basesamples_save_load(tmp_path, base_samples):
    with h5py.File(tmp_path / "basesamples.pkl", "w") as f:
        base_samples.save(f)

    with h5py.File(tmp_path / "basesamples.pkl", "r") as f:
        loaded_bs = BaseSamples.load(f)

    assert loaded_bs.xp == base_samples.xp
    assert isinstance(loaded_bs, BaseSamples)
    assert loaded_bs.parameters == ["a", "b"]
    assert np.allclose(loaded_bs.log_likelihood, base_samples.log_likelihood)
    assert np.allclose(loaded_bs.x, base_samples.x)


@pytest.mark.parametrize("flat", [True, False])
def test_basesamples_from_dict(base_samples, flat):
    d = base_samples.to_dict(flat=flat)
    print(d)
    bs2 = BaseSamples.from_dict(d)
    assert isinstance(bs2, BaseSamples)
    assert bs2.parameters == base_samples.parameters
    assert np.allclose(bs2.log_likelihood, base_samples.log_likelihood)
    assert np.allclose(bs2.x, base_samples.x)


def test_basesamples_to_dataframe_include(base_samples):
    df = base_samples.to_dataframe(include=["log_likelihood"])
    assert "log_likelihood" in df.columns
    assert "log_prior" not in df.columns
    assert "a" in df.columns and "b" in df.columns
    assert len(df) == len(base_samples)


def test_samples_compute_weights_constant_case():
    n, d = 10, 3
    a = 1.5
    x, ll, lp, lq = make_simple_samples(n=n, d=d, a=a)
    s = Samples(
        x=x,
        log_likelihood=ll,
        log_prior=lp,
        log_q=lq,
        parameters=[f"p{i}" for i in range(d)],
    )

    # log_w = a; evidence = exp(a); evidence_error = 0; ESS = n
    assert np.allclose(s.log_w, a)
    assert math.isclose(float(s.evidence), math.exp(a), rel_tol=1e-12)
    assert math.isclose(float(s.log_evidence), a, rel_tol=1e-12)
    assert math.isclose(float(s.evidence_error), 0.0, abs_tol=1e-15)
    assert math.isclose(float(s.log_evidence_error), 0.0, abs_tol=1e-15)
    assert math.isclose(float(s.effective_sample_size), n, rel_tol=1e-12)
    assert math.isclose(float(s.efficiency), 1.0, rel_tol=1e-12)

    # scaled_weights ∝ exp(log_w - max)
    sw = s.scaled_weights
    assert np.allclose(sw, np.exp(a - a))  # all ones


def test_samples_to_dict_includes_weight_fields():
    x, ll, lp, lq = make_simple_samples(n=4, d=2, a=0.3)
    s = Samples(
        x=x, log_likelihood=ll, log_prior=lp, log_q=lq, parameters=["a", "b"]
    )
    d = s.to_dict(flat=True)
    for k in [
        "log_w",
        "weights",
        "evidence",
        "log_evidence",
        "evidence_error",
        "log_evidence_error",
        "effective_sample_size",
        "a",
        "b",
    ]:
        assert k in d


def test_samples_getitem_preserves_evidence_fields():
    x, ll, lp, lq = make_simple_samples(n=6, d=2, a=0.2)
    s = Samples(
        x=x, log_likelihood=ll, log_prior=lp, log_q=lq, parameters=["a", "b"]
    )
    s2 = s[1:4]
    assert isinstance(s2, Samples)
    assert s2.x.shape[0] == 3
    # evidence fields were set on s; __getitem__ preserves them via from_samples
    assert math.isclose(
        float(s2.log_evidence), float(s.log_evidence), rel_tol=1e-12
    )
    assert math.isclose(
        float(s2.log_evidence_error),
        float(s.log_evidence_error),
        rel_tol=1e-12,
    )


def test_samples_concatenate_success_and_mismatch_errors():
    x1, ll1, lp1, lq1 = make_simple_samples(n=3, d=2, a=0.0)
    x2, ll2, lp2, lq2 = make_simple_samples(n=4, d=2, a=0.1)
    s1 = Samples(
        x=x1,
        log_likelihood=ll1,
        log_prior=lp1,
        log_q=lq1,
        parameters=["a", "b"],
    )
    s2 = Samples(
        x=x2,
        log_likelihood=ll2,
        log_prior=lp2,
        log_q=lq2,
        parameters=["a", "b"],
    )

    sc = Samples.concatenate([s1, s2])
    assert isinstance(sc, Samples)
    assert sc.x.shape == (7, 2)
    assert sc.parameters == ["a", "b"]

    # Mismatching dtype -> error
    x3 = x1.astype(np.float32)
    s3 = Samples(
        x=x3,
        log_likelihood=ll1,
        log_prior=lp1,
        log_q=lq1,
        parameters=["a", "b"],
        dtype="float32",
    )
    with pytest.raises(ValueError):
        Samples.concatenate([s1, s3])

    # Mismatching parameters -> error
    t_bad = Samples(
        x=x2,
        log_likelihood=ll2,
        log_prior=lp2,
        log_q=lq2,
        parameters=["a", "c"],
    )
    with pytest.raises(ValueError):
        Samples.concatenate([s1, t_bad])

    # Mismatching namespaces -> error (simulate by converting to a different namespace if available)
    # Here both are numpy, so simulate by monkeypatching xp check with a subclass of BaseSamples? Easiest:
    # Create a dummy with same params but fake xp attr.
    s_fake = Samples(
        x=x2,
        log_likelihood=ll2,
        log_prior=lp2,
        log_q=lq2,
        parameters=["a", "b"],
    )
    s_fake.xp = object()  # break equality
    with pytest.raises(ValueError):
        Samples.concatenate([s1, s_fake])


def test_samples_rejection_sample_accepts_proportionally():
    # Make one point much heavier
    n, d = 100, 1
    x = np.arange(n, dtype=float).reshape(n, d)
    ll = np.zeros(n)
    lp = np.zeros(n)
    lq = np.zeros(n)
    ll[0] = 10.0  # very large log-likelihood for index 0
    s = Samples(
        x=x, log_likelihood=ll, log_prior=lp, log_q=lq, parameters=["p0"]
    )
    rng = np.random.default_rng(1234)
    rs = s.rejection_sample(rng=rng)
    assert isinstance(rs, Samples)
    assert len(rs) <= len(s)
    # Heaviest point should be present with high probability
    if len(rs) > 0:
        assert 0 in rs.x.squeeze().astype(int)


def test_samples_to_dataframe():
    x, ll, lp, lq = make_simple_samples(n=4, d=2, a=0.5)
    s = Samples(
        x=x, log_likelihood=ll, log_prior=lp, log_q=lq, parameters=["a", "b"]
    )
    df = s.to_dataframe()
    assert "a" in df.columns and "b" in df.columns
    assert "log_w" in df.columns
    assert "log_likelihood" in df.columns
    assert len(df) == 4


def test_mcmc_samples_from_chain(rng):
    n_steps = 10
    n_walkers = 5
    dims = 3
    chain = rng.normal(size=(n_steps, n_walkers, dims))
    parameters = [f"p{i}" for i in range(dims)]
    samples = MCMCSamples.from_chain(
        chain=chain,
        parameters=parameters,
    )
    assert isinstance(samples, MCMCSamples)
    assert samples.x.shape == (n_steps * n_walkers, dims)
    assert samples.chain.shape == (n_steps, n_walkers, dims)


def test_ptmcmc_samples_from_chain(rng):
    n_steps = 10
    n_walkers = 5
    dims = 3
    n_temps = 4
    chain = rng.normal(size=(n_temps, n_steps, n_walkers, dims))
    parameters = [f"p{i}" for i in range(dims)]
    samples = PTMCMCSamples.from_chain(
        chain=chain,
        parameters=parameters,
    )
    assert isinstance(samples, MCMCSamples)
    assert samples.x.shape == (n_temps * n_steps * n_walkers, dims)
    assert samples.chain.shape == (n_temps, n_steps, n_walkers, dims)
    assert samples.cold_chain().x.shape == (n_steps * n_walkers, dims)


@pytest.mark.parametrize("samples_cls", [MCMCSamples, PTMCMCSamples])
def test_chain_samples_save_load_preserves_chain_shape(
    samples_cls, rng, tmp_path
):
    dims = 2
    if samples_cls is MCMCSamples:
        chain = rng.normal(size=(7, 3, dims))
        kwargs = {}
    else:
        chain = rng.normal(size=(4, 7, 3, dims))
        kwargs = {"betas": np.array([1.0, 0.5, 0.2, 0.0])}

    samples = samples_cls.from_chain(
        chain=chain,
        thin=2,
        burn_in=1,
        autocorrelation_time=np.ones(chain.shape[0]),
        **kwargs,
    )

    with h5py.File(tmp_path / "chain_samples.h5", "w") as f:
        samples.save(f, path="samples")
    with h5py.File(tmp_path / "chain_samples.h5", "r") as f:
        loaded = samples_cls.load(f, path="samples")

    assert tuple(loaded.chain_shape) == tuple(samples.chain_shape)
    assert loaded.chain.shape == chain.shape
    if samples_cls is PTMCMCSamples:
        assert np.allclose(loaded.betas, samples.betas)


@pytest.mark.parametrize("samples_cls", [MCMCSamples, PTMCMCSamples])
def test_chain_samples_post_process_shared_behavior(samples_cls, rng):
    dims = 2
    if samples_cls is MCMCSamples:
        chain = rng.normal(size=(8, 3, dims))
        values = rng.normal(size=(8, 3))
        kwargs = {}
    else:
        chain = rng.normal(size=(4, 8, 3, dims))
        values = rng.normal(size=(4, 8, 3))
        kwargs = {"betas": np.array([1.0, 0.5, 0.2, 0.0])}

    samples = samples_cls.from_chain(
        chain=chain,
        log_likelihood=values,
        log_prior=values + 1.0,
        log_q=values - 1.0,
        thin=2,
        burn_in=3,
        autocorrelation_time=np.ones(chain.shape[0]),
        **kwargs,
    )

    out = samples.post_process(burn_in=1, thin=2)
    expected_chain = (
        chain[1::2] if samples_cls is MCMCSamples else chain[:, 1::2]
    )
    expected_values = (
        values[1::2] if samples_cls is MCMCSamples else values[:, 1::2]
    )

    assert isinstance(out, samples_cls)
    assert out.chain.shape == expected_chain.shape
    assert out.x.shape == (np.prod(expected_chain.shape[:-1]), dims)
    assert np.allclose(out.log_likelihood, expected_values.reshape(-1))
    assert out.thin == 4
    assert out.burn_in == 4
    assert np.allclose(out.autocorrelation_time, samples.autocorrelation_time)

    if samples_cls is PTMCMCSamples:
        assert np.allclose(out.betas, samples.betas)


@pytest.mark.parametrize("samples_cls", [MCMCSamples, PTMCMCSamples])
def test_chain_samples_post_process_noop_returns_self(samples_cls, rng):
    dims = 2
    if samples_cls is MCMCSamples:
        chain = rng.normal(size=(5, 2, dims))
        kwargs = {}
    else:
        chain = rng.normal(size=(3, 5, 2, dims))
        kwargs = {"betas": np.array([1.0, 0.5, 0.0])}
    samples = samples_cls.from_chain(chain=chain, **kwargs)

    assert samples.post_process(burn_in=0, thin=1) is samples


def test_mcmc_samples_getitem_preserves_metadata(rng):
    dims = 2
    chain = rng.normal(size=(7, 2, dims))
    samples = MCMCSamples.from_chain(
        chain=chain,
        thin=3,
        burn_in=5,
        autocorrelation_time=np.ones(chain.shape[0]),
    )
    out = samples[:4]

    assert isinstance(out, MCMCSamples)
    assert out.chain_shape == (len(out.x),)
    assert out.x.shape == (4, dims)
    assert out.thin == samples.thin
    assert out.burn_in == samples.burn_in
    assert np.allclose(out.autocorrelation_time, samples.autocorrelation_time)


def test_ptmcmc_getitem_raises_not_implemented(rng):
    chain = rng.normal(size=(3, 7, 2, 2))
    samples = PTMCMCSamples.from_chain(
        chain=chain, betas=np.array([1.0, 0.5, 0.0])
    )
    with pytest.raises(NotImplementedError):
        _ = samples[:4]


def test_ptmcmc_subsample_shapes_and_metadata(rng):
    n_temps, n_steps, n_walkers, dims = 4, 20, 3, 2
    chain = rng.normal(size=(n_temps, n_steps, n_walkers, dims))
    ll = rng.normal(size=(n_temps, n_steps, n_walkers))
    lp = rng.normal(size=(n_temps, n_steps, n_walkers))
    betas = np.array([1.0, 0.5, 0.2, 0.0])
    samples = PTMCMCSamples.from_chain(
        chain=chain,
        betas=betas,
        log_likelihood=ll,
        log_prior=lp,
        thin=2,
        burn_in=5,
        autocorrelation_time=np.ones(n_temps),
    )

    n_sub = 10
    out = samples.subsample(n_sub, rng=rng)

    assert isinstance(out, PTMCMCSamples)
    assert out.chain_shape == (n_temps, n_sub)
    assert out.x.shape == (n_temps * n_sub, dims)
    assert out.log_likelihood.shape == (n_temps * n_sub,)
    assert out.log_prior.shape == (n_temps * n_sub,)
    assert np.allclose(out.betas, betas)
    assert out.thin == samples.thin
    assert out.burn_in == samples.burn_in
    assert np.allclose(out.autocorrelation_time, samples.autocorrelation_time)


def test_ptmcmc_subsample_samples_are_from_original(rng):
    n_temps, n_steps, dims = 3, 15, 2
    chain = rng.normal(size=(n_temps, n_steps, dims))
    betas = np.array([1.0, 0.5, 0.0])
    samples = PTMCMCSamples.from_chain(chain=chain, betas=betas)

    out = samples.subsample(5, rng=rng)

    # Each subsampled row must appear in the original flat chain
    original_x = samples.x
    for row in out.x:
        assert any(
            np.allclose(row, original_x[i]) for i in range(len(original_x))
        )


def test_ptmcmc_subsample_no_replacement_within_temperature(rng):
    n_temps, n_steps, dims = 2, 20, 1
    chain = np.arange(n_temps * n_steps * dims, dtype=float).reshape(
        n_temps, n_steps, dims
    )
    betas = np.array([1.0, 0.0])
    samples = PTMCMCSamples.from_chain(chain=chain, betas=betas)

    out = samples.subsample(10, rng=rng)
    out_chain = out.chain  # (n_temps, 10, dims)

    for t in range(n_temps):
        rows = [out_chain[t, i, 0] for i in range(10)]
        assert len(rows) == len(set(rows)), (
            "Duplicate samples within a temperature"
        )


def test_ptmcmc_subsample_too_many_raises(rng):
    chain = rng.normal(size=(3, 5, 2))
    samples = PTMCMCSamples.from_chain(
        chain=chain, betas=np.array([1.0, 0.5, 0.0])
    )
    with pytest.raises(ValueError, match="exceeds"):
        samples.subsample(6, rng=rng)


def test_ptmcmc_ti_matches_equations_35_to_37():
    chain = np.zeros((3, 2, 2, 1), dtype=float)
    betas = np.array([1.0, 0.5, 0.0], dtype=float)
    logl = np.array(
        [
            [[1.0, 0.0], [2.0, -1.0]],  # beta=1.0
            [[0.5, -0.5], [1.5, -1.5]],  # beta=0.5
            [[0.0, -1.0], [1.0, -2.0]],  # beta=0.0
        ],
        dtype=float,
    )
    samples = PTMCMCSamples.from_chain(
        chain=chain,
        betas=betas,
        log_likelihood=logl,
    )

    logz, err = samples.log_evidence_thermodynamic_integration(
        burn_in_fraction=0.0
    )

    # Manual Eq. (35)-(37).
    order = np.argsort(betas)
    t = betas[order]
    logl_flat = logl.reshape(3, -1)[order]
    mean_logl = np.mean(logl_flat, axis=1)
    expected_logz = np.trapezoid(mean_logl, t)
    ti_per_sample = np.trapezoid(logl_flat, t, axis=0)
    expected_err = math.sqrt(np.var(ti_per_sample) / len(ti_per_sample))

    assert math.isclose(logz, expected_logz, rel_tol=1e-12, abs_tol=1e-12)
    assert math.isclose(err, expected_err, rel_tol=1e-12, abs_tol=1e-12)


def test_ptmcmc_stepping_stone_error_matches_formula():
    # Three temperatures (including beta=0), two steps, two walkers.
    chain = np.zeros((3, 2, 2, 1), dtype=float)
    betas = np.array([1.0, 0.5, 0.0], dtype=float)
    logl = np.array(
        [
            [[0.0, 0.0], [0.0, 0.0]],
            [[1.0, 0.0], [2.0, -1.0]],
            [[0.5, -0.5], [1.5, -1.0]],
        ],
        dtype=float,
    )
    samples = PTMCMCSamples.from_chain(
        chain=chain,
        betas=betas,
        log_likelihood=logl,
    )

    logz, err = samples.log_evidence_stepping_stone(burn_in_fraction=0.0)

    # Manual stepping-stone and error computation.
    logl_flat = logl.reshape(3, -1)
    expected_logz = 0.0
    var = 0.0
    n = logl_flat.shape[1]
    for i in range(len(betas) - 1):
        dbeta = betas[i] - betas[i + 1]
        a = dbeta * logl_flat[i + 1]
        expected_logz += np.log(np.mean(np.exp(a)))
        ratio = np.exp(a) / np.mean(np.exp(a))
        var += np.sum(ratio**2)
    expected_log_err = math.sqrt(var / (n**2))

    assert math.isclose(logz, expected_logz, rel_tol=1e-12, abs_tol=1e-12)
    assert math.isclose(err, expected_log_err, rel_tol=1e-12, abs_tol=1e-12)


def test_ptmcmc_stepping_stone_error_constant_logl_matches_formula():
    chain = np.zeros((3, 2, 2, 1), dtype=float)
    betas = np.array([1.0, 0.5, 0.0], dtype=float)
    # Constant per temperature => exp(a_n)/r_j = 1 for all n.
    logl = np.array(
        [
            [[1.0, 1.0], [1.0, 1.0]],
            [[2.0, 2.0], [2.0, 2.0]],
            [[3.0, 3.0], [3.0, 3.0]],
        ],
        dtype=float,
    )
    samples = PTMCMCSamples.from_chain(
        chain=chain,
        betas=betas,
        log_likelihood=logl,
    )

    _, err = samples.log_evidence_stepping_stone(burn_in_fraction=0.0)
    n = logl.reshape(3, -1).shape[1]
    expected_err = math.sqrt((len(betas) - 1) / n)
    assert math.isclose(err, expected_err, rel_tol=1e-12, abs_tol=1e-12)


def test_smc_unnormalized_and_normalized_log_weights_and_ratio():
    n, d = 6, 2
    x = np.arange(n * d, dtype=float).reshape(n, d)
    ll = np.linspace(0.0, 1.0, n)
    lp = np.zeros(n)
    lq = np.zeros(n)
    beta0 = 0.3
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=lp,
        log_q=lq,
        beta=beta0,
        parameters=["a", "b"],
    )

    beta1 = 0.7
    # unnormalized weights (beta-beta0) * (ll+lp) since lq=0
    uw = smc.unnormalized_log_weights(beta1)
    assert np.allclose(uw, (beta1 - beta0) * (ll + lp))

    # log_evidence_ratio = logsumexp(uw) - log(n)
    import scipy.special as sps  # optional; if unavailable, reimplement quickly

    lse = sps.logsumexp(uw) - math.log(n)
    assert math.isclose(
        float(smc.log_evidence_ratio(beta1)), float(lse), rel_tol=1e-12
    )

    # normalized (shifted) log weights should not contain NaN and sum of exp(log_w - logsumexp(log_w)) = 1
    lw = smc.log_weights(beta1)
    w = np.exp(lw - sps.logsumexp(lw))
    assert np.isfinite(w).all()
    assert math.isclose(w.sum(), 1.0, rel_tol=1e-12)


def test_smc_log_weights_nan_guard_raises():
    x = np.zeros((3, 1))
    ll = np.array([0.0, np.nan, 1.0])
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=np.zeros(3),
        log_q=np.zeros(3),
        beta=0.0,
    )
    with pytest.raises(ValueError):
        _ = smc.log_weights(0.5)


def test_smc_resample_shapes_and_beta_and_warning_on_same_beta(caplog):
    n, d = 10, 2
    x = np.arange(n * d, dtype=float).reshape(n, d)
    ll = np.linspace(0, 1, n)
    lp = np.zeros(n)
    lq = np.zeros(n)
    beta0 = 0.2
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=lp,
        log_q=lq,
        beta=beta0,
        parameters=["a", "b"],
    )

    rng = np.random.default_rng(42)
    out = smc.resample(beta=0.8, n_samples=7, rng=rng)
    assert isinstance(out, SMCSamples)
    assert out.x.shape == (7, d)
    assert math.isclose(out.beta, 0.8)

    # same beta & n_samples=None -> returns self and logs a warning
    caplog.clear()
    with caplog.at_level("WARNING"):
        same = smc.resample(beta=beta0, n_samples=None, rng=rng)
    assert same is smc
    assert any(
        "Resampling with the same beta value" in r.message
        for r in caplog.records
    )


def test_smc_variance_delta_method_nonnegative():
    x = np.arange(10.0).reshape(5, 2)
    ll = np.linspace(0, 1, 5)
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=np.zeros(5),
        log_q=np.zeros(5),
        beta=0.3,
    )
    var = smc.log_evidence_ratio_variance(beta=0.6)
    assert float(var) >= 0.0 or np.isnan(var)


def test_smc_to_standard_samples_and_getitem_preserve_fields():
    x = np.arange(8.0).reshape(4, 2)
    ll = np.linspace(0, 1, 4)
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=np.zeros(4),
        log_q=np.zeros(4),
        beta=0.4,
        parameters=["a", "b"],
        log_evidence=1.23,
    )

    std = smc.to_standard_samples()
    assert isinstance(std, Samples)
    # For std Samples, weights aren't computed (log_q missing by design),
    # but log_evidence is carried over.
    assert std.log_w is None
    assert math.isclose(float(std.log_evidence), 1.23, rel_tol=1e-12)

    # __getitem__ preserves beta and log_evidence
    smc2 = smc[1:3]
    assert isinstance(smc2, SMCSamples)
    assert smc2.x.shape[0] == 2
    assert math.isclose(float(smc2.beta), 0.4, rel_tol=1e-12)
    assert math.isclose(float(smc2.log_evidence), 1.23, rel_tol=1e-12)


def test_str_contains_counts_and_metrics():
    x, ll, lp, lq = make_simple_samples(n=4, d=2, a=0.7)
    s = Samples(
        x=x, log_likelihood=ll, log_prior=lp, log_q=lq, parameters=["a", "b"]
    )
    s_str = str(s)
    assert "No. samples: 4" in s_str
    assert "No. parameters: 2" in s_str
    assert "Log evidence:" in s_str
    assert "Effective sample size:" in s_str
    assert "Efficiency:" in s_str


def test_smcsamples_to_dict_includes_smc_fields():
    x, ll, lp, lq = make_simple_samples(n=4, d=2, a=0.5)
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=lp,
        log_q=lq,
        beta=0.3,
        parameters=["a", "b"],
        log_evidence=0.8,
        log_evidence_error=0.1,
    )
    d = smc.to_dict(flat=True)
    assert d["beta"] == 0.3
    assert d["log_evidence"] == 0.8
    assert d["log_evidence_error"] == 0.1


def test_smc_samples_to_dataframe():
    x, ll, lp, lq = make_simple_samples(n=3, d=2, a=0.4)
    smc = SMCSamples(
        x=x,
        log_likelihood=ll,
        log_prior=lp,
        log_q=lq,
        beta=0.2,
        parameters=["a", "b"],
    )
    df = smc.to_dataframe()
    assert "a" in df.columns and "b" in df.columns
    assert "log_likelihood" in df.columns
    assert len(df) == 3
