#   Copyright 2024 - present The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
import arviz as az
import numpy as np
import pytest

from pytensor import config
from pytensor.compile import SharedVariable
from pytensor.graph import Constant

import pymc as pm

from pymc import sample_posterior_predictive, set_data
from pymc.distributions.transforms import logodds
from pymc.model.transform.conditioning import (
    change_value_transforms,
    do,
    observe,
    remove_value_transforms,
)
from pymc.model.transform.optimization import freeze_dims_and_data
from pymc.variational.minibatch_rv import create_minibatch_rv


def test_observe():
    with pm.Model() as m_old:
        x = pm.Normal("x")
        y = pm.Normal("y", x)
        z = pm.Normal("z", y)

    m_new = observe(m_old, {y: 0.5})

    assert len(m_new.free_RVs) == 2
    assert len(m_new.observed_RVs) == 1
    assert m_new["x"] in m_new.free_RVs
    assert m_new["y"] in m_new.observed_RVs
    assert m_new["z"] in m_new.free_RVs

    np.testing.assert_allclose(
        m_old.compile_logp()({"x": 0.9, "y": 0.5, "z": 1.4}),
        m_new.compile_logp()({"x": 0.9, "z": 1.4}),
    )

    # Test two substitutions
    m_new = observe(m_old, {y: 0.5, z: 1.4})

    assert len(m_new.free_RVs) == 1
    assert len(m_new.observed_RVs) == 2
    assert m_new["x"] in m_new.free_RVs
    assert m_new["y"] in m_new.observed_RVs
    assert m_new["z"] in m_new.observed_RVs

    np.testing.assert_allclose(
        m_old.compile_logp()({"x": 0.9, "y": 0.5, "z": 1.4}),
        m_new.compile_logp()({"x": 0.9}),
    )


def test_observe_minibatch():
    data = np.zeros((100,), dtype=config.floatX)
    batch_size = 10
    with pm.Model() as m_old:
        x = pm.Normal("x")
        y = pm.Normal("y", x)
        # Minibatch RVs are usually created with `total_size` kwarg
        z_raw = pm.Normal.dist(y, shape=batch_size)
        mb_z = create_minibatch_rv(z_raw, total_size=data.shape)
        m_old.register_rv(mb_z, name="mb_z")

    mb_data = pm.Minibatch(data, batch_size=batch_size)
    m_new = observe(m_old, {mb_z: mb_data})

    assert len(m_new.free_RVs) == 2
    assert len(m_new.observed_RVs) == 1
    assert m_new["x"] in m_new.free_RVs
    assert m_new["y"] in m_new.free_RVs
    assert m_new["mb_z"] in m_new.observed_RVs

    np.testing.assert_allclose(
        m_old.compile_logp()({"x": 0.9, "y": 0.5, "mb_z": np.zeros(10)}),
        m_new.compile_logp()({"x": 0.9, "y": 0.5}),
    )


def test_observe_deterministic():
    y_censored_obs = np.array([0.9, 0.5, 0.3, 1, 1], dtype=config.floatX)

    with pm.Model() as m_old:
        x = pm.Normal("x")
        y = pm.Normal.dist(x, shape=(5,))
        y_censored = pm.Deterministic("y_censored", pm.math.clip(y, -1, 1))

    m_new = observe(m_old, {y_censored: y_censored_obs})

    with pm.Model() as m_ref:
        x = pm.Normal("x")
        pm.Censored("y_censored", pm.Normal.dist(x), lower=-1, upper=1, observed=y_censored_obs)


def test_observe_dims():
    with pm.Model(coords={"test_dim": range(5)}) as m_old:
        x = pm.Normal("x", dims="test_dim")

    m_new = observe(m_old, {x: np.arange(5, dtype=config.floatX)})
    assert m_new.named_vars_to_dims["x"] == ["test_dim"]


def test_do():
    rng = np.random.default_rng(seed=435)
    with pm.Model() as m_old:
        x = pm.Normal("x", 0, 1e-3)
        y = pm.Normal("y", x, 1e-3)
        z = pm.Normal("z", y + x, 1e-3)

    assert -5 < pm.draw(z, random_seed=rng) < 5

    m_new = do(m_old, {y: x + 100})

    assert len(m_new.free_RVs) == 2
    assert m_new["x"] in m_new.free_RVs
    assert m_new["y"] in m_new.deterministics
    assert m_new["z"] in m_new.free_RVs

    assert 95 < pm.draw(m_new["z"], random_seed=rng) < 105

    # Test two substitutions
    with m_old:
        switch = pm.Data("switch", 1)
    m_new = do(m_old, {y: 100 * switch, x: 100 * switch})

    assert len(m_new.free_RVs) == 1
    assert m_new["y"] not in m_new.deterministics
    assert m_new["x"] not in m_new.deterministics
    assert m_new["z"] in m_new.free_RVs

    assert 195 < pm.draw(m_new["z"], random_seed=rng) < 205
    with m_new:
        pm.set_data({"switch": 0})
    assert -5 < pm.draw(m_new["z"], random_seed=rng) < 5


def test_do_posterior_predictive():
    with pm.Model() as m:
        x = pm.Normal("x", 0, 1)
        y = pm.Normal("y", x, 1)
        z = pm.Normal("z", y + x, 1e-3)

    # Dummy posterior
    idata_m = az.from_dict(
        {
            "x": np.full((2, 500), 25),
            "y": np.full((2, 500), np.nan),
            "z": np.full((2, 500), np.nan),
        }
    )

    # Replace `y` by a constant `100.0`
    m_do = do(m, {y: 100.0})
    with m_do:
        idata_do = pm.sample_posterior_predictive(idata_m, var_names="z")

    assert 120 < idata_do.posterior_predictive["z"].mean() < 130


@pytest.mark.parametrize("mutable", (False, True))
def test_do_constant(mutable):
    rng = np.random.default_rng(seed=122)
    with pm.Model() as m:
        x = pm.Data("x", 0)
        y = pm.Normal("y", x, 1e-3)

    if not mutable:
        m = freeze_dims_and_data(m, data=["x"])

    do_m = do(m, {m["x"]: 105})
    assert pm.draw(do_m["y"], random_seed=rng) > 100


def test_do_deterministic():
    rng = np.random.default_rng(seed=435)
    with pm.Model() as m:
        x = pm.Normal("x", 0, 1e-3)
        y = pm.Deterministic("y", x + 105)
        z = pm.Normal("z", y, 1e-3)

    do_m = do(m, {"z": x - 105})
    assert pm.draw(do_m["z"], random_seed=rng) < 100


def test_do_dims():
    coords = {"test_dim": range(10)}
    with pm.Model(coords=coords) as m:
        x = pm.Normal("x", dims="test_dim")
        y = pm.Deterministic("y", x + 5, dims="test_dim")

    do_m = do(
        m,
        {"x": np.zeros(10, dtype=config.floatX)},
    )
    assert do_m.named_vars_to_dims["x"] == ["test_dim"]

    do_m = do(
        m,
        {"y": np.zeros(10, dtype=config.floatX)},
    )
    assert do_m.named_vars_to_dims["y"] == ["test_dim"]


@pytest.mark.parametrize("prune", (False, True))
def test_do_prune(prune):
    with pm.Model() as m:
        x0 = pm.Data("x0", 0)
        x1 = pm.Data("x1", 0)
        y = pm.Normal("y")
        y_det = pm.Deterministic("y_det", y + x0)
        z = pm.Normal("z", y_det)
        llike = pm.Normal("llike", z + x1, observed=0)

    orig_named_vars = {"x0", "x1", "y", "y_det", "z", "llike"}
    assert set(m.named_vars) == orig_named_vars

    do_m = do(m, {y_det: x0 + 5}, prune_vars=prune)
    if prune:
        assert set(do_m.named_vars) == {"x0", "x1", "y_det", "z", "llike"}
    else:
        assert set(do_m.named_vars) == orig_named_vars

    do_m = do(m, {z: 0.5}, prune_vars=prune)
    if prune:
        assert set(do_m.named_vars) == {"x1", "z", "llike"}
    else:
        assert set(do_m.named_vars) == orig_named_vars


def test_do_self_reference():
    """Check we can replace a variable by an expression that refers to the same variable."""
    with pm.Model() as m:
        x = pm.Normal("x", 0, 1)

    with pytest.warns(
        UserWarning,
        match="Intervention expression references the variable that is being intervened",
    ):
        new_m = do(m, {x: x + 100})

    x = new_m["x"]
    do_x = new_m["do_x"]
    draw_x, draw_do_x = pm.draw([x, do_x], draws=5)
    np.testing.assert_allclose(draw_x + 100, draw_do_x)


def test_do_make_intervenstions_shared():
    with pm.Model(coords={"obs": [1]}) as m:
        x = pm.Normal("x", dims="obs")
        y = pm.Normal("y", dims="obs")

    constant_m = do(m, {x: [0.5]}, make_interventions_shared=False)
    constant_x = constant_m["x"]
    assert isinstance(constant_x, Constant)
    np.testing.assert_array_equal(constant_x.data, [0.5])

    shared_m = do(m, {x: [0.5]}, make_interventions_shared=True)
    shared_x = shared_m["x"]
    assert isinstance(shared_x, SharedVariable)
    np.testing.assert_array_equal(shared_x.get_value(borrow=True), [0.5])

    with shared_m:
        set_data({"x": [0.6, 0.9]}, coords={"obs": [2, 3]})
        pp_y = pm.sample_prior_predictive(draws=3).prior["y"]
    assert pp_y.sizes == {"chain": 1, "draw": 3, "obs": 2}
    assert pp_y.shape == (1, 3, 2)


@pytest.mark.parametrize(
    "make_interventions_shared",
    [True, pytest.param(False, marks=pytest.mark.xfail(reason="#6876"))],
)
def test_do_sample_posterior_predictive(make_interventions_shared):
    # Regression test for https://github.com/pymc-devs/pymc/issues/6977
    with pm.Model() as model:
        a = pm.Normal("a")
        b = pm.Deterministic("b", a * 2)
        c = pm.Normal("c", b / 2)

    idata = az.from_dict({"a": [[1.0]], "b": [[2.0]], "c": [[1.0]]})

    with do(model, {a: 1000}, make_interventions_shared=make_interventions_shared):
        pp = sample_posterior_predictive(idata, var_names=["c"], predictions=True).predictions
    assert (pp["c"] > 500).all()


def test_change_value_transforms():
    with pm.Model() as base_m:
        p = pm.Uniform("p", 0, 1, default_transform=None)
        w = pm.Binomial("w", n=9, p=p, observed=6)
        assert base_m.rvs_to_transforms[p] is None
        assert base_m.rvs_to_values[p].name == "p"

    with change_value_transforms(base_m, {"p": logodds}) as transformed_p:
        new_p = transformed_p["p"]
        assert transformed_p.rvs_to_transforms[new_p] == logodds
        assert transformed_p.rvs_to_values[new_p].name == "p_logodds__"
        mean_q = pm.find_MAP(progressbar=False)

    with change_value_transforms(transformed_p, {"p": None}) as untransformed_p:
        new_p = untransformed_p["p"]
        assert untransformed_p.rvs_to_transforms[new_p] is None
        assert untransformed_p.rvs_to_values[new_p].name == "p"
        std_q = ((1 / pm.find_hessian(mean_q, vars=[new_p])) ** 0.5)[0]

    np.testing.assert_allclose(np.round(mean_q["p"], 2), 0.67)
    np.testing.assert_allclose(np.round(std_q[0], 2), 0.16)


def test_change_value_transforms_error():
    with pm.Model() as m:
        x = pm.Uniform("x", observed=5.0)

    with pytest.raises(ValueError, match="All keys must be free variables in the model"):
        change_value_transforms(m, {x: logodds})


def test_remove_value_transforms():
    with pm.Model() as base_m:
        p = pm.Uniform("p", transform=logodds, default_transform=None)
        q = pm.Uniform("q", transform=logodds, default_transform=None)

    new_m = remove_value_transforms(base_m)
    new_p = new_m["p"]
    new_q = new_m["q"]
    assert new_m.rvs_to_transforms == {new_p: None, new_q: None}

    new_m = remove_value_transforms(base_m, [p, q])
    new_p = new_m["p"]
    new_q = new_m["q"]
    assert new_m.rvs_to_transforms == {new_p: None, new_q: None}

    new_m = remove_value_transforms(base_m, [p])
    new_p = new_m["p"]
    new_q = new_m["q"]
    assert new_m.rvs_to_transforms == {new_p: None, new_q: logodds}

    new_m = remove_value_transforms(base_m, ["q"])
    new_p = new_m["p"]
    new_q = new_m["q"]
    assert new_m.rvs_to_transforms == {new_p: logodds, new_q: None}
