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

from math import log

import numpy

from sppl.distributions import bernoulli
from sppl.distributions import choice
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 >> choice({'0':.1, '1': .9})),
                    (Z[0] >> bernoulli(p=.1))
                        & (Y >> choice({'0':.9, '1': .1}))
                ], weights=[log(.730), log(.270)])
            ]),
            ProductSPN([
                Z[0] >> bernoulli(p=.1),
                Y >> choice({'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 >> choice({'0':.3, '1':.7})
                        & X[0] >> bernoulli(p=.1)
                        & Z[0] >> bernoulli(p=.1),
                    Y >> choice({'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]
