import json

import pandas as pd
from nose import with_setup

from pybbn.graph.dag import Bbn
from pybbn.graph.edge import Edge, EdgeType
from pybbn.graph.factory import Factory
from pybbn.graph.node import BbnNode
from pybbn.graph.variable import Variable
from pybbn.pptc.inferencecontroller import InferenceController
from pybbn.sampling.sampling import LogicSampler


def setup():
    """
    Setup.
    :return: None.
    """
    pass


def teardown():
    """
    Teardown.
    :return: None.
    """
    pass


def get_json_string():
    return json.dumps(get_dict())


def get_dict():
    return {
        "V": ["Letter", "Grade", "Intelligence", "SAT", "Difficulty"],
        "E": [["Difficulty", "Grade"],
              ["Intelligence", "Grade"],
              ["Intelligence", "SAT"],
              ["Grade", "Letter"]],
        "Vdata": {
            "Letter": {
                "ord": 4,
                "numoutcomes": 2,
                "vals": ["weak", "strong"],
                "parents": ["Grade"],
                "children": None,
                "cprob": {
                    "['A']": [.1, .9],
                    "['B']": [.4, .6],
                    "['C']": [.99, .01]
                }
            },
            "SAT": {
                "ord": 3,
                "numoutcomes": 2,
                "vals": ["lowscore", "highscore"],
                "parents": ["Intelligence"],
                "children": None,
                "cprob": {
                    "['low']": [.95, .05],
                    "['high']": [.2, .8]
                }
            },
            "Grade": {
                "ord": 2,
                "numoutcomes": 3,
                "vals": ["A", "B", "C"],
                "parents": ["Difficulty", "Intelligence"],
                "children": ["Letter"],
                "cprob": {
                    "['easy', 'low']": [.3, .4, .3],
                    "['easy', 'high']": [.9, .08, .02],
                    "['hard', 'low']": [.05, .25, .7],
                    "['hard', 'high']": [.5, .3, .2]
                }
            },
            "Intelligence": {
                "ord": 1,
                "numoutcomes": 2,
                "vals": ["low", "high"],
                "parents": None,
                "children": ["SAT", "Grade"],
                "cprob": [.7, .3]
            },
            "Difficulty": {
                "ord": 0,
                "numoutcomes": 2,
                "vals": ["easy", "hard"],
                "parents": None,
                "children": ["Grade"],
                "cprob": [.6, .4]
            }
        }
    }


@with_setup(setup, teardown)
def test_from_libpgm_discrete_json():
    """
    Tests create py-bbn BBN from JSON string specifying libpgm BBN.
    :return: None.
    """
    j = get_json_string()
    bbn = Factory.from_libpgm_discrete_json(j)

    assert len(bbn.nodes) == 5
    assert len(bbn.edges) == 4


@with_setup(setup, teardown)
def test_from_libpgm_discrete_dictionary():
    """
    Tests create py-bbn BBN from dictionary specifying libpgm BBN.
    :return: None.
    """
    d = get_dict()
    bbn = Factory.from_libpgm_discrete_dictionary(d)

    assert len(bbn.nodes) == 5
    assert len(bbn.edges) == 4

    join_tree = InferenceController.apply(bbn)

    __validate_posterior__({
        'Difficulty': [0.6, 0.4],
        'Intelligence': [0.7, 0.3],
        'Grade': [0.362, 0.288, 0.350],
        'SAT': [0.725, 0.275],
        'Letter': [0.498, 0.502]
    }, join_tree, debug=True)


@with_setup(setup, teardown)
def test_from_data_simple():
    """
    Tests create BBN from data.
    :return: None.
    """
    a = BbnNode(Variable(0, 'a', ['on', 'off']), [0.5, 0.5])
    b = BbnNode(Variable(1, 'b', ['on', 'off']), [0.5, 0.5, 0.4, 0.6])
    c = BbnNode(Variable(2, 'c', ['on', 'off']), [0.7, 0.3, 0.2, 0.8])

    bbn1 = Bbn() \
        .add_node(a) \
        .add_node(b) \
        .add_node(c) \
        .add_edge(Edge(a, b, EdgeType.DIRECTED)) \
        .add_edge(Edge(b, c, EdgeType.DIRECTED))

    sampler = LogicSampler(bbn1)
    samples = sampler.get_samples(n_samples=10000, seed=37)

    i2n = {n.variable.id: n.variable.name for n in bbn1.get_nodes()}
    samples = pd.DataFrame(samples).rename(columns=i2n)

    parents = {
        'a': [],
        'b': ['a'],
        'c': ['b']
    }

    bbn2 = Factory.from_data(parents, samples)

    join_tree1 = InferenceController.apply(bbn1)
    join_tree2 = InferenceController.apply(bbn2)

    posteriors1 = join_tree1.get_posteriors()
    posteriors2 = join_tree2.get_posteriors()

    for k, v1 in posteriors1.items():
        assert k in posteriors2

        v2 = posteriors2[k]
        assert len(v1) == len(v2)

        for k2 in v1:
            assert k2 in v2
            diff = abs(v1[k2] - v2[k2])
            assert diff < 0.01


def __validate_posterior__(expected, join_tree, debug=False):
    for node in join_tree.get_bbn_nodes():
        potential = join_tree.get_bbn_potential(node)
        if debug is True:
            p = ', '.join(['{}'.format(e) for e in potential.entries])
            s = '{} : {}'.format(node.variable.name, p)
            print(s)

        o = [e.value for e in potential.entries]
        e = expected[node.variable.name]

        assert len(o) == len(e)
        for ob, ex in zip(o, e):
            diff = abs(ob - ex)
            if diff > 0.001 and debug:
                print('\t**observed={}, expected={}'.format(ob, ex))
            elif debug is False:
                assert diff < 0.001
