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

from math import log

import numpy

from sppl.distributions import bernoulli
from sppl.distributions import norm
from sppl.spn import ProductSPN
from sppl.spn import SumSPN
from sppl.spn import spn_cache_duplicate_subtrees
from sppl.transforms import Id

rng = numpy.random.RandomState(1)

W = Id('W')
Y = Id('Y')
X = [Id('X[0]'), Id('X[1]')]
Z = [Id('Z[0]'), Id('Z[1]')]

def test_cache_simple_leaf():
    spn = .5 * (W >> norm(loc=0, scale=1)) | .5 * (W >> norm(loc=0, scale=1))
    assert spn.children[0] is not spn.children[1]
    spn_cached = spn_cache_duplicate_subtrees(spn, {})
    assert spn_cached.children[0] is spn_cached.children[1]

def test_cache_simple_sum_of_product():
    spn \
        = 0.3 * ((W >> norm(loc=0, scale=1)) & (Y >> norm(loc=0, scale=1))) \
        | 0.7 * ((W >> norm(loc=0, scale=1)) & (Y >> norm(loc=0, scale=2)))
    spn_cached = spn_cache_duplicate_subtrees(spn, {})
    assert spn_cached.children[0].children[0] is spn_cached.children[1].children[0]

def test_cache_complex_sum_of_product():
    # Test case adapted from the SPN generated by
    # test_repeat.make_model_repeat(n=2)
    duplicate_subtrees = [None, None]
    for i in range(2):
        duplicate_subtrees[i] = SumSPN([
            ProductSPN([
                (X[0] >> bernoulli(p=.1)),
                SumSPN([
                    (Z[0] >> bernoulli(p=.5))
                        & (Y >> {'0':.1, '1': .9}),
                    (Z[0] >> bernoulli(p=.1))
                        & (Y >> {'0':.9, '1': .1})
                ], weights=[log(.730), log(.270)])
            ]),
            ProductSPN([
                Z[0] >> bernoulli(p=.1),
                Y >> {'0':.9, '1':.1},
                X[0] >> bernoulli(p=.5),
            ]),
        ], weights=[log(.925), log(.075)])

    assert duplicate_subtrees[0] == duplicate_subtrees[1]
    assert duplicate_subtrees[0] is not duplicate_subtrees[1]

    left_subtree = ProductSPN([
        X[1] >> bernoulli(p=.5),
        SumSPN([
            ProductSPN([
                duplicate_subtrees[0],
                Z[1] >> bernoulli(p=.5),
            ]),
            ProductSPN([
                Z[1] >> bernoulli(p=.7),
                SumSPN([
                    Y >> {'0':.3, '1':.7}
                        & X[0] >> bernoulli(p=.1)
                        & Z[0] >> bernoulli(p=.1),
                    Y >> {'0':.7, '1':.3}
                        & X[0] >> bernoulli(p=.5)
                        & Z[0] >> bernoulli(p=.5),
                ], weights=[log(.9), log(.1)])
            ])
        ], weights=[log(.783), log(.217)])
    ])

    right_subtree = ProductSPN([
        Z[1] >> bernoulli(p=.8),
        X[1] >> bernoulli(p=.1),
        duplicate_subtrees[1]
    ])

    spn = .92 * left_subtree | .08 * right_subtree

    spn_cached = spn_cache_duplicate_subtrees(spn, {})
    assert spn_cached.children[0].children[1].children[0].children[0] is duplicate_subtrees[0]
    assert spn_cached.children[1].children[2] is duplicate_subtrees[0]
