import aesara.tensor as at
import pytest

from aemcmc.utils import ModelInfo


def test_ModelInfo_maps():
    """Test maps generated by `ModelInfo`'s constructor."""
    srng = at.random.RandomStream()

    beta_rv = srng.normal(0, 1, name="beta")
    beta_vv = beta_rv.clone()
    beta_vv.name = "beta_vv"

    X = at.matrix("X")

    mu = at.exp(X.dot(beta_rv))
    mu.name = "mu"

    Y_rv = srng.poisson(mu, name="Y")
    y_vv = Y_rv.clone()
    y_vv.name = "y"

    observed_rvs = (Y_rv,)
    rvs_to_values = {Y_rv: y_vv, beta_rv: beta_vv}
    deterministic_vars = (mu,)

    amodel = ModelInfo(observed_rvs, rvs_to_values, deterministic_vars, {})

    assert amodel.values_to_rvs == {y_vv: Y_rv, beta_vv: beta_rv}
    assert amodel.names_to_vars == {
        "Y": Y_rv,
        "beta": beta_rv,
        "beta_vv": beta_vv,
        "mu": mu,
        "y": y_vv,
    }
    assert amodel.observed_values == (y_vv,)
    assert amodel.unobserved_rvs == (beta_rv,)
    assert amodel.unobserved_values == (beta_vv,)
    assert amodel.updates == {}
    assert set(amodel.names_to_vars.values()) == set((Y_rv, beta_rv, mu, y_vv, beta_vv))

    smodel = ModelInfo(observed_rvs, rvs_to_values)
    assert smodel.deterministic_vars == ()
    assert smodel.updates == {}
    assert set(smodel.names_to_vars.values()) == set((Y_rv, beta_rv, y_vv, beta_vv))


def test_ModelInfo_errors():
    with pytest.raises(ValueError, match=".*non-empty names.*"):
        ModelInfo((), {at.vector("a"): at.vector()}, (), {})

    with pytest.raises(ValueError, match=".*unique names.*"):
        ModelInfo((), {at.vector("a"): at.vector("a")}, (), {})
