from typing import NewType

import z3
from pysmt.shortcuts import GE, LE, Real, And, GT, LT, Equals, Or, Not

from assignments import Assignment
from DPGM import Node
from distribution import Distribution
from solver_instance import SolverInstance
from z3_instance import Z3Instance

# label constants
Label = NewType("Label", str)
l_in = Label("in")
l_out = Label("out")
l_undec = Label("undec")


class LabelingScheme:

    num_args = 0  # The number of arguments each labeling scheme takes

    def get_label(self, argument: Node, distribution: Distribution) -> Label:
        """
        Get the label of the given argument according the given model and the labeling approach.
        :return: l_in, l_out or l_undec
        """
        # implemented by subclasses
        pass

    def get_constraints(self, node: Node, si: SolverInstance):
        """
        Return a list of constraints on the probabilistic variable of the node where each constraint corresponds to a
        different label. I.e. for cautious labeling and node A, return [p_A == 1, p_A == 0, And(p_A < 1, p_A > 0)].
        Used to compute all labelings.
        """
        # implemented by subclasses
        return []

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        """
        Return a list of constraints on the probabilistic variable of the node where each constraint corresponds to a
        different label. I.e. for cautious labeling and node A, return [p_A == 1, p_A == 0, And(p_A < 1, p_A > 0)].
        Used to compute all labelings.
        """
        # implemented by subclasses
        return []


class Threshold(LabelingScheme):
    """
    Given two thresholds t_in and t_out, all nodes with probability
     - greater equal t_in are labeled 'in',
     - less than or equal to t_out are labeled 'out',
     - and between t_in and t_out are labeled 'undec'.
     """
    num_args = 2

    def __init__(self, t_in: float, t_out: float):
        self.t_in = t_in
        self.t_out = t_out

    def get_label(self, node: Node, distribution: Distribution):
        prob = distribution.get_marginal_probability(node)
        if prob >= self.t_in:
            return l_in
        elif prob <= self.t_out:
            return l_out
        else:
            return l_undec

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        in_constraint = GE(node_prob, Real(self.t_in))
        out_constraint = LE(node_prob, Real(self.t_out))
        undec_constraint = And(LT(node_prob, Real(self.t_in)), GT(node_prob, Real(self.t_out)))
        return [in_constraint, out_constraint, undec_constraint]

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        node_prob = z3i.get_node_prob_var(node)
        in_constraint = node_prob >= self.t_in
        out_constraint = node_prob <= self.t_out
        undec_constraint = z3.And(node_prob < self.t_in, node_prob > self.t_out, z3i.context)
        return [in_constraint, out_constraint, undec_constraint]


class Cautious(Threshold):
    """
    Arguments with probability 1 are labeled 'in', those with probability 0 are labeled 'out' and all others are
    labeled 'undec'.
    """
    num_args = 0

    def __init__(self):
        super().__init__(1.0, 0.0)


class Firm(LabelingScheme):
    """
    Arguments with probability 0.5 are labeled 'undec', those with probability >0.5 are labeled 'in' and all those with
    probability <0.5 are labeled 'out'.
    """
    def get_label(self, node: Node, distribution: Distribution):
        prob = distribution.get_marginal_probability(node)
        if prob > 0.5:
            return l_in
        elif prob < 0.5:
            return l_out
        else:
            return l_undec

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        in_constraint = GT(node_prob, Real(0.5))
        out_constraint = LT(node_prob, Real(0.5))
        undec_constraint = Equals(node_prob, Real(0.5))
        return [undec_constraint, in_constraint, out_constraint]
        # undec first because checking for equality is probably easier than inequalities

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        node_prob = z3i.get_node_prob_var(node)
        in_constraint = node_prob > 0.5
        out_constraint = node_prob < 0.5
        undec_constraint = node_prob == 0.5
        return [undec_constraint, in_constraint, out_constraint]
        # undec first because checking for equality is probably easier than inequalities


class ThresholdClassicalOld(LabelingScheme):
    """
    Requires two thresholds t_in and t_min_attack_out. A node is labeled
     - 'in' if its probability is greater or equal t_in,
     - 'out' if it is not labeled 'in' and one of its attackers has probability greater or equal t_min_attack_out, and
     - 'undec' in all other cases.

     Uses old implementation of get_constraints_z3 which does not necessarily produce unique labelings, though might be
     faster in some cases.
     """
    num_args = 2

    def __init__(self, t_in: float, t_min_attack_out: float):
        self.t_in = t_in
        self.t_att_out = t_min_attack_out

    def get_label(self, node: Node, distribution: Distribution):
        node_prob = distribution.get_marginal_probability(node)
        if node_prob >= self.t_in:
            return l_in
        else:
            label = l_undec
            for attacker in node.get_parents():
                attacker_prob = distribution.get_marginal_probability(attacker)
                if attacker_prob >= self.t_att_out:
                    label = l_out
                    break
            return label

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        if self.t_in == self.t_att_out:
            in_constraint = GE(node_prob, Real(self.t_in))
            not_in_constraint = LT(node_prob, Real(self.t_in))
            return [in_constraint, not_in_constraint]
        else:
            t_max = max(self.t_in, self.t_att_out)
            t_min = min(self.t_in, self.t_att_out)
            first_constraint = GE(node_prob, Real(t_max))
            second_constraint = LT(node_prob, Real(t_min))
            third_constraint = And(LT(node_prob, Real(t_max)), GE(node_prob, Real(t_min)))
            return [first_constraint, second_constraint, third_constraint]

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        node_prob = z3i.get_node_prob_var(node)
        if self.t_in == self.t_att_out:
            in_constraint = node_prob >= self.t_in
            not_in_constraint = node_prob < self.t_in
            return [in_constraint, not_in_constraint]
        else:
            t_max = max(self.t_in, self.t_att_out)
            t_min = min(self.t_in, self.t_att_out)
            first_constraint = node_prob >= t_max
            second_constraint = node_prob < t_min
            third_constraint = z3.And(node_prob < t_max, node_prob >= t_min, z3i.context)
            return [first_constraint, second_constraint, third_constraint]


class ClassicalOld(ThresholdClassicalOld):
    """
    A node is labeled
     - 'in' if it holds with probability 1,
     - 'out' if it is not labeled 'in' and one of its attackers holds with probability 1, and
     - 'undec' in all other cases.

     Uses old implementation of get_constraints_z3 which does not necessarily produce unique labelings, though might be
     faster in some cases.
    """
    num_args = 0

    def __init__(self):
        super().__init__(1.0, 1.0)


class ThresholdClassical(LabelingScheme):
    """
    Requires two thresholds t_in and t_min_attack_out. A node is labeled
     - 'in' if its probability is greater or equal t_in,
     - 'out' if it is not labeled 'in' and one of its attackers has probability greater or equal t_min_attack_out, and
     - 'undec' in all other cases.
     """
    num_args = 2

    def __init__(self, t_in: float, t_min_attack_out: float):
        self.t_in = t_in
        self.t_att_out = t_min_attack_out

    def get_label(self, node: Node, distribution: Distribution):
        node_prob = distribution.get_marginal_probability(node)
        if node_prob >= self.t_in:
            return l_in
        else:
            label = l_undec
            for attacker in node.get_parents():
                attacker_prob = distribution.get_marginal_probability(attacker)
                if attacker_prob >= self.t_att_out:
                    label = l_out
                    break
            return label

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        in_constraint = GE(node_prob, Real(self.t_in))
        parent_out_constraint = Or(
            [GE(si.get_node_symbol(parent), Real(self.t_att_out)) for parent in node.get_parents()])
        out_constraint = And(LT(node_prob, Real(self.t_in)), parent_out_constraint)
        parent_undec_constraint = And(
            [LT(si.get_node_symbol(parent), Real(self.t_att_out)) for parent in node.get_parents()])
        undec_constraint = And(LT(node_prob, Real(self.t_in)), parent_undec_constraint)
        return [in_constraint, out_constraint, undec_constraint]

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        node_prob = z3i.get_node_prob_var(node)
        in_constraint = node_prob >= self.t_in
        parent_out_constraint = z3.Or(
            [z3i.get_node_prob_var(parent) >= self.t_att_out for parent in node.get_parents()])
        out_constraint = z3.And(node_prob < self.t_in, parent_out_constraint)
        parent_undec_constraint = z3.And(
            [z3i.get_node_prob_var(parent) < self.t_att_out for parent in node.get_parents()])
        undec_constraint = z3.And(node_prob < self.t_in, parent_undec_constraint)

        return [in_constraint, out_constraint, undec_constraint]


class Classical(ThresholdClassical):
    """
    A node is labeled
     - 'in' if it holds with probability 1,
     - 'out' if it is not labeled 'in' and one of its attackers holds with probability 1, and
     - 'undec' in all other cases.
    """
    num_args = 0

    def __init__(self):
        super().__init__(1.0, 1.0)


class AdHoc(LabelingScheme):
    """
    TODO
    """
    def get_label(self, node: Node, distribution: Distribution):
        prob = distribution.get_marginal_probability(node)
        if prob == 1.0:
            return l_in
        else:
            no_attacker: Assignment = tuple([(attacker, False) for attacker in node.get_parents()])
            if distribution.get_part_assignment_probability(no_attacker) == 0:
                return l_out
            else:
                return l_undec

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        in_constraint = GE(node_prob, Real(1))

        no_attacker: Assignment = tuple([(attacker, False) for attacker in node.get_parents()])
        no_attacker_prob = si.get_symbol(no_attacker)
        out_constraint = Equals(no_attacker_prob, Real(0))

        undec_constraint = Not(Or(in_constraint, out_constraint))
        return [in_constraint, out_constraint, undec_constraint]


class Optimistic(LabelingScheme):
    """
    A node is labeled
     - 'in' if it holds with probability larger than 0, and
     - 'out' otherwise.
     """

    def get_label(self, node: Node, distribution: Distribution):
        prob = distribution.get_marginal_probability(node)
        if prob > 0.0:
            return l_in
        else:
            return l_out

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        in_constraint = GT(node_prob, Real(0))
        out_constraint = Equals(node_prob, Real(0))
        return [in_constraint, out_constraint]

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        node_prob = z3i.get_node_prob_var(node)
        in_constraint = node_prob > 0
        out_constraint = node_prob == 0
        return [in_constraint, out_constraint]


class Pessimistic(LabelingScheme):
    """
    A node is labeled
     - 'in' if it holds with probability 1, and
     - 'out' otherwise.
     """

    def get_label(self, node: Node, distribution: Distribution):
        prob = distribution.get_marginal_probability(node)
        if prob == 1.0:
            return l_in
        else:
            return l_out

    def get_constraints(self, node: Node, si: SolverInstance):
        node_prob = si.get_node_symbol(node)
        in_constraint = Equals(node_prob, Real(1))
        out_constraint = LT(node_prob, Real(1))
        return [in_constraint, out_constraint]

    def get_constraints_z3(self, node: Node, z3i: Z3Instance):
        node_prob = z3i.get_node_prob_var(node)
        in_constraint = node_prob == 1
        out_constraint = node_prob < 1
        return [in_constraint, out_constraint]


###################
# Helper functions
###################

def get_labeling_scheme_class(name: str):
    """
    By dark magic, get the labeling scheme class given its name as string, e.g. "firm" to yield the class 'Firm'.

    :param name: the name of the labeling scheme class (case insensitive)
    :return: the labeling scheme class corresponding to the name, or None if no match was found
    """
    name = name.lower()
    queue = LabelingScheme.__subclasses__().copy()
    while queue:
        labeling_class = queue.pop()
        if name == labeling_class.__name__.lower():
            return labeling_class
        queue.extend(labeling_class.__subclasses__())
    return None


def get_all_labeling_scheme_names():
    """
    :return: a list of the names of all available labeling schemes
    """
    names = []
    queue = LabelingScheme.__subclasses__().copy()
    while queue:
        labeling_class = queue.pop()
        queue.extend(labeling_class.__subclasses__())
        names.append(labeling_class.__name__)
    return names
