"""
Check to see if the model has the variable
sharing property.
"""
from itertools import product
from typing import List, Optional, Set, Tuple
from common import set_to_str
from logic import Logic, Implication
from model import (
    Model, model_closure, ModelFunction, ModelValue
)

SMT_LOADED = True
try:
    from z3 import And, Or, Implies, sat
    from smt import SMTModelEncoder, SMTLogicEncoder
except ImportError:
    SMT_LOADED = False

class VSP_Result:
    def __init__(
            self, has_vsp: bool, model_name: Optional[str] = None,
            subalgebra1: Optional[Set[ModelValue]] = None,
            subalgebra2: Optional[Set[ModelValue]] = None):
        self.has_vsp = has_vsp
        self.model_name = model_name
        self.subalgebra1 = subalgebra1
        self.subalgebra2 = subalgebra2

    def __str__(self):
        if not self.has_vsp:
            return f"Model {self.model_name} does not have the variable sharing property."
        return f"""Model {self.model_name} has the variable sharing property.
Subalgebra 1: {set_to_str(self.subalgebra1)}
Subalgebra 2: {set_to_str(self.subalgebra2)}
"""

def has_vsp_magical(model: Model, impfunction: ModelFunction,
            negation_defined: bool) -> VSP_Result:
    """
    Checks whether a MaGIC model has the variable
    sharing property.
    """
    # NOTE: No models with only one designated
    # value satisfies VSP
    if len(model.designated_values) == 1:
        return VSP_Result(False, model.name)

    assert model.ordering is not None, "Expected ordering table in model"

    top = model.ordering.top()
    bottom = model.ordering.bottom()

    # Compute I the set of tuples (x, y) where
    # x -> y does not take a designiated value
    I: List[Tuple[ModelValue, ModelValue]] = []

    for (x, y) in product(model.designated_values, model.designated_values):
        if impfunction(x, y) not in model.designated_values:
            I.append((x, y))

    # Find the subalgebras which falsify implication
    for xys in I:

        xi = xys[0]

        # Discard ({⊥} ∪ A', B) subalgebras
        if bottom is not None and xi == bottom:
            continue

        # Discard ({⊤} ∪ A', B) subalgebras when negation is defined
        if top is not None and negation_defined and xi == top:
            continue

        yi = xys[1]

        # Discard (A, {⊤} ∪ B') subalgebras
        if top is not None and yi == top:
            continue

        # Discard (A, {⊥} ∪ B') subalgebras when negation is defined
        if bottom is not None and negation_defined and yi == bottom:
            continue

        # Discard ({a} ∪ A', {b} ∪ B') subalgebras when a <= b
        if model.ordering.is_lt(xi, yi):
            continue

        # Discard ({a} ∪ A', {b} ∪ B') subalgebras when b <= a and negation is defined
        if negation_defined and model.ordering.is_lt(yi, xi):
            continue

        # Compute the left closure of the set containing xi under all the operations
        carrier_set_left: Set[ModelValue] = model_closure({xi,}, model.logical_operations, bottom)

        # Discard ({⊥} ∪ A', B) subalgebras
        if bottom is not None and bottom in carrier_set_left:
            continue

        # Discard ({⊤} ∪ A', B) subalgebras when negation is defined
        if top is not None and negation_defined and top in carrier_set_left:
            continue

        # Compute the closure of all operations
        # with just the ys
        carrier_set_right: Set[ModelValue] = model_closure({yi,}, model.logical_operations, top)

        # Discard (A, {⊤} ∪ B') subalgebras
        if top is not None and top in carrier_set_right:
            continue

        # Discard (A, {⊥} ∪ B') subalgebras when negation is defined
        if bottom is not None and negation_defined and bottom in carrier_set_right:
            continue

        # Discard subalgebras that intersect
        if not carrier_set_left.isdisjoint(carrier_set_right):
            continue

        # Check whether for all pairs in the subalgebra,
        # that implication is falsified.
        falsified = True
        for (x2, y2) in product(carrier_set_left, carrier_set_right):
            if impfunction(x2, y2) in model.designated_values:
                falsified = False
                break

        if falsified:
            return VSP_Result(True, model.name, carrier_set_left, carrier_set_right)

    return VSP_Result(False, model.name)

def has_vsp_smt(model: Model, impfn: ModelFunction) -> VSP_Result:
    """
    Checks whether a given model satisfies the variable
    sharing property via SMT
    """
    if not SMT_LOADED:
        raise Exception("Z3 is not property installed, cannot check via SMT")

    encoder = SMTModelEncoder(model)

    # Create predicates for our two subalgebras
    IsInK1 = encoder.create_predicate("IsInK1", 1)
    IsInK2 = encoder.create_predicate("IsInK2", 1)

    # Enforce that our two subalgebras are non-empty
    encoder.solver.add(Or([IsInK1(x) for x in encoder.smt_carrier_set]))
    encoder.solver.add(Or([IsInK2(x) for x in encoder.smt_carrier_set]))

    # K1/K2 are closed under the operations
    for model_fn, smt_fn in encoder.model_function_map.items():
        for xs in product(encoder.smt_carrier_set, repeat=model_fn.arity):
            encoder.solver.add(
                Implies(
                    And([IsInK1(x) for x in xs]),
                    IsInK1(smt_fn(*xs))
                )
            )
            encoder.solver.add(
                Implies(
                    And([IsInK2(x) for x in xs]),
                    IsInK2(smt_fn(*xs))
                )
            )

    # x -> y is non-designated
    smt_imp = encoder.model_function_map[impfn]
    for (x, y) in product(encoder.smt_carrier_set, encoder.smt_carrier_set):
        encoder.solver.add(
            Implies(
                And(IsInK1(x), IsInK2(y)),
                encoder.is_designated(smt_imp(x, y)) == False
            )
        )

    # Execute solver
    if encoder.solver.check() == sat:
        # Extract subalgebras
        smt_model = encoder.solver.model()
        K1_smt = [x for x in encoder.smt_carrier_set if smt_model.evaluate(IsInK1(x))]
        K1 = {ModelValue(str(x)) for x in K1_smt}

        K2_smt = [x for x in encoder.smt_carrier_set if smt_model.evaluate(IsInK2(x))]
        K2 = {ModelValue(str(x)) for x in K2_smt}

        return VSP_Result(True, model.name, K1, K2)
    else:
        return VSP_Result(False, model.name)


def has_vsp(model: Model, impfunction: ModelFunction,
            negation_defined: bool) -> VSP_Result:
    if model.is_magical:
        return has_vsp_magical(model, impfunction, negation_defined)

    return has_vsp_smt(model)


def logic_has_vsp(logic: Logic, size: int) -> Optional[Tuple[Model, VSP_Result]]:
    """
    Checks whether a given logic satisfies
    the variable sharing property by looking
    for a many-valued matrix of a specific size.

    If the logic does witness the VSP, then
    this function will return the matrix model
    and the subalgebras that witness it.

    Otherwise, if no matrix model of that given
    size can be found, it will return None
    """
    assert size > 0

    encoder = SMTLogicEncoder(logic, size)

    ## The following adds constraints which satisfy the VSP

    # Membership Predicates for K1/K2
    IsInK1 = encoder.create_predicate("IsInK1", 1)
    IsInK2 = encoder.create_predicate("IsInK2", 1)

    # K1 and K2 are non-empty
    encoder.solver.add(Or([IsInK1(x) for x in encoder.smt_carrier_set]))
    encoder.solver.add(Or([IsInK2(x) for x in encoder.smt_carrier_set]))

    # K1/K2 are closed under the operations
    for op, SmtOp in encoder.operation_function_map.items():
        for xs in product(encoder.smt_carrier_set, repeat=op.arity):
            encoder.solver.add(
                Implies(
                    And([IsInK1(x) for x in xs]),
                    IsInK1(SmtOp(*xs))
                )
            )
            encoder.solver.add(
                Implies(
                    And([IsInK2(x) for x in xs]),
                    IsInK2(SmtOp(*xs))
                )
            )

    # x -> y is non-designated
    Impfn = encoder.operation_function_map[Implication]
    for (x, y) in product(encoder.smt_carrier_set, encoder.smt_carrier_set):
        encoder.solver.add(
            Implies(
                And(IsInK1(x), IsInK2(y)),
                encoder.is_designated(Impfn(x, y)) == False
            )
        )

    model = encoder.find_model()

    # We failed to find a VSP witness
    if model is None:
        return None

    # Otherwise, a matrix model and correspoding
    # subalgebras exist.

    smt_model = encoder.solver.model()
    K1_smt = [x for x in encoder.smt_carrier_set if smt_model.evaluate(IsInK1(x))]
    K1 = {ModelValue(str(x)) for x in K1_smt}

    K2_smt = [x for x in encoder.smt_carrier_set if smt_model.evaluate(IsInK2(x))]
    K2 = {ModelValue(str(x)) for x in K2_smt}

    return model, VSP_Result(True, model.name, K1, K2)