"""
Created on October 22, 2018

@author: Nicola DI Mauro
"""
from spn.algorithms.MPE import get_mpe_top_down_leaf, add_node_mpe
from spn.structure.leaves.cltree.CLTree import CLTree

import numpy as np
import logging

logger = logging.getLogger(__name__)


def cltree_bottom_up_log_ll(node, data, dtype=np.float64):
    probs = np.ones((data.shape[0], 1))
    cltree_mpe(node, data, probs)
    return probs


def cltree_top_down(node, input_vals, data, lls_per_node=None, dtype=np.float64):
    return


def add_cltree_mpe_support():
    add_node_mpe(CLTree, cltree_bottom_up_log_ll, cltree_top_down)


def cltree_mpe(node, data, logprobs):

    log_factors = np.array(node.log_factors)

    for r in range(data.shape[0]):

        messages = np.zeros((node.n_features, 2))
        states = [[0, 0] for i in range(node.n_features)]
        MAP = {}

        for i in node.post_order:
            if i != 0:
                state_evidence = data[r, node.scope[i]]
                if not np.isnan(state_evidence):
                    state_evidence = int(state_evidence)
                    states[i][0] = state_evidence
                    states[i][1] = state_evidence
                    messages[node.tree[i], 0] += log_factors[i, state_evidence, 0] + messages[i, state_evidence]
                    messages[node.tree[i], 1] += log_factors[i, state_evidence, 1] + messages[i, state_evidence]
                else:
                    state_evidence_parent = data[r, node.scope[node.tree[i]]]
                    if not np.isnan(state_evidence_parent):
                        state_evidence_parent = int(state_evidence_parent)
                        if (
                            log_factors[i, 0, state_evidence_parent] + messages[i, 0]
                            > log_factors[i, 1, state_evidence_parent] + messages[i, 1]
                        ):
                            states[i][state_evidence_parent] = 0
                            messages[node.tree[i], state_evidence_parent] += (
                                log_factors[i, 0, state_evidence_parent] + messages[i, 0]
                            )
                        else:
                            states[i][state_evidence_parent] = 1
                            messages[node.tree[i], state_evidence_parent] += (
                                log_factors[i, 1, state_evidence_parent] + messages[i, 1]
                            )
                    else:
                        for parent in range(2):
                            if log_factors[i, 0, parent] + messages[i, 0] > log_factors[i, 1, parent] + messages[i, 1]:
                                states[i][parent] = 0
                                messages[node.tree[i], parent] += log_factors[i, 0, parent] + messages[i, 0]
                            else:
                                states[i][parent] = 1
                                messages[node.tree[i], parent] += log_factors[i, 1, parent] + messages[i, 1]
        logprob = 0.0
        for i in node.df_order:
            if node.tree[i] == -1:
                state_evidence = data[r, node.scope[i]]
                if not np.isnan(state_evidence):
                    state_evidence = int(state_evidence)
                    MAP[i] = state_evidence
                    logprob += log_factors[i, int(MAP[i]), 0]
                else:
                    if log_factors[i, 0, 0] + messages[i, 0] > log_factors[i, 1, 0] + messages[i, 1]:
                        MAP[i] = 0
                        data[r, node.scope[i]] = 0
                    else:
                        MAP[i] = 1
                        data[r, node.scope[i]] = 1
                    logprob += log_factors[i, int(MAP[i]), 0]
            else:
                MAP[i] = states[i][MAP[node.tree[i]]]
                if np.isnan(data[r, node.scope[i]]):
                    data[r, node.scope[i]] = MAP[i]
                logprob += log_factors[i, int(MAP[i]), int(MAP[node.tree[i]])]

        logprobs[r] = logprob
