# Copyright 2020 MIT Probabilistic Computing Project.
# See LICENSE.txt

from math import log

from sppl.compilers.ast_to_spn import For
from sppl.compilers.ast_to_spn import Id
from sppl.compilers.ast_to_spn import IdArray
from sppl.compilers.ast_to_spn import IfElse
from sppl.compilers.ast_to_spn import Otherwise
from sppl.compilers.ast_to_spn import Sample
from sppl.compilers.ast_to_spn import Sequence
from sppl.distributions import bernoulli
from sppl.distributions import choice
from sppl.math_util import allclose

Y = Id('Y')
X = IdArray('X', 5)
Z = IdArray('Z', 5)

def test_simple_model():
    command = Sequence(
        Sample(Y, bernoulli(p=0.5)),
        For(0, 5, lambda i:
            Sample(X[i], bernoulli(p=1/(i+1)))))
    model = command.interpret()

    symbols = model.get_symbols()
    assert len(symbols) == 6
    assert Y in symbols
    assert X[0] in symbols
    assert X[1] in symbols
    assert X[2] in symbols
    assert X[3] in symbols
    assert X[4] in symbols
    assert model.logprob(X[0] << {1}) == log(1/1)
    assert model.logprob(X[1] << {1}) == log(1/2)
    assert model.logprob(X[2] << {1}) == log(1/3)
    assert model.logprob(X[3] << {1}) == log(1/4)
    assert model.logprob(X[4] << {1}) == log(1/5)

def test_complex_model():
    # Slow for larger number of repetitions
    # https://github.com/probcomp/sum-product-dsl/issues/43
    command = Sequence(
        Sample(Y, choice({'0': .2, '1': .2, '2': .2, '3': .2, '4': .2})),
        For(0, 3, lambda i: Sequence(
            Sample(Z[i], bernoulli(p=0.1)),
            IfElse(
                Y << {str(i)} | Z[i] << {0}, Sample(X[i], bernoulli(p=1/(i+1))),
                Otherwise,                   Sample(X[i], bernoulli(p=0.1))))))
    model = command.interpret()
    assert allclose(model.prob(Y << {'0'}), 0.2)

def test_complex_model_reorder():
    command = Sequence(
        Sample(Y, choice({'0': .2, '1': .2, '2': .2, '3': .2, '4': .2})),
        For(0, 3, lambda i:
            Sample(Z[i], bernoulli(p=0.1))),
        For(0, 3, lambda i:
            IfElse(
                Y << {str(i)}, Sample(X[i], bernoulli(p=1/(i+1))),
                Z[i] << {0},   Sample(X[i], bernoulli(p=1/(i+1))),
                Otherwise,     Sample(X[i], bernoulli(p=0.1)))))
    model = command.interpret()
    assert(allclose(model.prob(Y << {'0'}), 0.2))

def test_repeat_handcode_equivalence():
    model_repeat = make_model_for()
    model_hand = make_model_handcode()

    assert allclose(model_repeat.prob(Y << {'0', '1'}), 0.4)
    assert allclose(model_repeat.prob(Z[0] << {0}), 0.5)
    assert allclose(model_repeat.prob(Z[0] << {1}), 0.5)

    event_condition = (X[0] << {1}) | (Y << {'1'})
    model_repeat_condition = model_repeat.condition(event_condition)
    model_hand_condition = model_hand.condition(event_condition)

    for event in [
            Y << {'0','1'},
            Z[0] << {0},
            Z[1] << {0},
            X[0] << {0},
            X[1] << {0},
        ]:
        lp_repeat = model_repeat.logprob(event)
        lp_hand = model_hand.logprob(event)
        assert allclose(lp_hand, lp_repeat)

        lp_repeat_condition = model_repeat_condition.logprob(event)
        lp_hand_condition = model_hand_condition.logprob(event)
        assert allclose(lp_hand_condition, lp_repeat_condition)

    # This test case ensures that the duplicate subtrees in the mixture
    # components are pointers to the same object, and is obtained by
    # manually inspecting the rendering of the network generated by the
    # following code:
    #
    #   from sppl.magics.render import render_graphviz
    #   render_graphviz(model_repeat, show=True)
    #
    # See also test_cache_duplicate_subtrees.test_cache_complex_sum_of_product
    a = model_repeat.children[0].children[0].children[1].children[0]
    b = model_repeat.children[1].children[0]
    assert a is b

# ==============================================================================
# Helper functions.

def make_model_for(n=2):
    command = Sequence(
        Sample(Y, choice({'0': .2, '1': .2, '2': .2, '3': .2, '4': .2})),
        For(0, n, lambda i: Sequence(
            Sample(Z[i], bernoulli(p=.5)),
            IfElse(
                (Y << {str(i)}) | (Z[i] << {0}), Sample(X[i], bernoulli(p=.1)),
                Otherwise,                       Sample(X[i], bernoulli(p=.5))))))
    return command.interpret()

def make_model_handcode():
    command = Sequence(
        Sample(Y, choice({'0': .2, '1': .2, '2': .2, '3': .2, '4': .2})),
        Sample(Z[0], bernoulli(p=.5)),
        Sample(Z[1], bernoulli(p=.5)),
        IfElse(
            Y << {str(0)}, Sequence(
                Sample(X[0], bernoulli(p=.1)),
                IfElse(
                    Z[1] << {0},    Sample(X[1], bernoulli(p=.1)),
                    Otherwise,      Sample(X[1], bernoulli(p=.5)))),
            Y << {str(1)}, Sequence(
                Sample(X[1], bernoulli(p=.1)),
                IfElse(
                    Z[0] << {0},    Sample(X[0], bernoulli(p=.1)),
                    Otherwise,      Sample(X[0], bernoulli(p=.5)))),
            Otherwise, Sequence(
                IfElse(
                    Z[0] << {0},    Sample(X[0], bernoulli(p=.1)),
                    Otherwise,      Sample(X[0], bernoulli(p=.5))),
                IfElse(
                    Z[1] << {0},    Sample(X[1], bernoulli(p=.1)),
                    Otherwise,      Sample(X[1], bernoulli(p=.5))))))
    return command.interpret()
