import sys

from problog.cycles import break_cycles
from problog.engine import DefaultEngine
from problog.formula import LogicDAG
from problog.logic import Constant, Term
from problog.program import PrologFile, ExtendedPrologFactory
from problog.sdd_formula import SDD


class ConstraintFactory(ExtendedPrologFactory):
    """Extends standard ProbLog input reader to process ensure_true, ensure_false and ensure_prob."""

    def __init__(self, *args, **kwargs):
        super(ConstraintFactory, self).__init__(*args, **kwargs)
        self.constraint_count = 0

    def build_clause(self, functor, operand1, operand2, location=None, **extra):
        if len(operand1) == 1 and operand1[0].functor in (
            "ensure_true",
            "ensure_false",
            "ensure_prob",
        ):
            operand1 = [
                Term("constraint", Constant(self.constraint_count), operand1[0])
            ]
            self.constraint_count += 1
        return super(ConstraintFactory, self).build_clause(
            functor, operand1, operand2, location=location, **extra
        )


def formula_to_flatzinc_float(formula):
    """Converts a formula to a MiniZinc model using floats."""
    # print (formula)
    vardefs = []
    constraints = []
    boolvars = []

    for i, n, t in formula:
        if t == "atom":
            if str(n.probability) == "?":
                vardefs.append("var bool: B%s :: output_var;" % i)
                vardefs.append("var int: I%s;" % i)
                vardefs.append("var 0.0 .. 1.0: F%s;" % i)
                constraints.append("constraint bool2int(B%s,I%s);" % (i, i))
                constraints.append("constraint int2float(I%s,F%s);" % (i, i))
                boolvars.append("B%s" % i)
                vardefs.append("var 0.0 .. 1.0: V%s;" % i)
                vardefs.append("var 0.0 .. 1.0: Vn%s;" % i)
                constraints.append(
                    "constraint float_lin_eq([1.0, 1.0], [V%s, Vn%s], 1.0);" % (i, i)
                )
                constraints.append("constraint float_eq(F%s,V%s);" % (i, i))
            else:
                vardefs.append("float: V%s = %s;" % (i, n.probability))
                vardefs.append("float: Vn%s = %s;" % (i, 1 - float(n.probability)))
        elif t == "disj":
            vardefs.append("var 0.0 .. 1.0: V%s;" % i)
            vardefs.append("var 0.0 .. 1.0: Vn%s;" % i)
            constraints.append(
                "constraint float_lin_eq([1.0, 1.0], [V%s, Vn%s], 1.0);" % (i, i)
            )

            children = ["V%s" % str(c).replace("-", "n") for c in n.children] + [
                "V%s" % i
            ]
            coef = ["1.0"] * (len(children) - 1) + ["-1.0"]
            constraints.append(
                "constraint float_lin_eq([%s], [%s], 0.0);"
                % (", ".join(coef), ", ".join(children))
            )
        elif t == "conj":
            vardefs.append("var 0.0 .. 1.0: V%s;" % i)
            vardefs.append("var 0.0 .. 1.0: Vn%s;" % i)
            constraints.append(
                "constraint float_lin_eq([1.0, 1.0], [V%s, Vn%s], 1.0);" % (i, i)
            )

            children = [str(c).replace("-", "n") for c in n.children]
            constraints.append(
                "constraint float_times(V%s, V%s, V%s);" % (children[0], children[1], i)
            )

    for name, index in formula.get_names(label="constraint"):
        index = str(index).replace("-", "n")
        con = name.args[1]
        if con.functor == "ensure_prob":
            constraints.append(
                "constraint float_lin_le([1.0],[V%s], %s);" % (index, con.args[1])
            )
            constraints.append(
                "constraint float_lin_le([-1.0], [V%s], -%s);" % (index, con.args[0])
            )
        elif con.functor == "ensure_true":
            constraints.append("constraint float_eq(V%s, 1.0);" % (index,))
        elif con.functor == "ensure_false":
            constraints.append("constraint float_eq(V%s, 0.0);" % (index,))

    solve = [
        "solve :: bool_search([%s], input_order, indomain_min, complete) satisfy;"
        % ", ".join(boolvars)
    ]

    fzn = "\n".join(vardefs + constraints + solve)

    return fzn


def formula_to_flatzinc_bool(formula):
    """Converts a formula to a MiniZinc model using booleans."""
    vardefs = []
    constraints = []
    boolvars = []

    for i, n, t in formula:
        if t == "atom":
            if str(n.probability) == "?":
                vardefs.append("var bool: B%s :: output_var;" % i)
            else:
                vardefs.append("bool: B%s = true;" % (i,))
        elif t == "disj":
            vardefs.append("var bool: B%s;" % i)
            children = ["B%s" % str(c).replace("-", "n") for c in n.children]
            constraints.append(
                "constraint array_bool_or([%s], %s);" % (", ".join(children), "B%s" % i)
            )
        elif t == "conj":
            vardefs.append("var bool: B%s;" % i)
            children = ["B%s" % str(c).replace("-", "n") for c in n.children]
            constraints.append(
                "constraint array_bool_and([%s], %s);"
                % (", ".join(children), "B%s" % i)
            )
        vardefs.append("var bool: Bn%s;" % i)
        constraints.append("constraint bool_not(B%s, Bn%s);" % (i, i))

    for name, index in formula.get_names(label="constraint"):
        if index is not None:
            index = str(index).replace("-", "n")
            con = name.args[1]
            assert con.functor != "ensure_prob"
            if con.functor == "ensure_true":
                constraints.append("constraint bool_eq(B%s, true);" % (index,))
            elif con.functor == "ensure_false":
                constraints.append("constraint bool_eq(B%s, false);" % (index,))

    solve = [
        "solve :: bool_search([%s], input_order, indomain_min, complete) satisfy;"
        % ", ".join(boolvars)
    ]

    fzn = "\n".join(vardefs + constraints + solve)

    return fzn


def debug(*args):
    """Helper function for printing debugging output."""
    print(*args, file=sys.stderr)


def find_solver():
    import subprocess

    try:
        subprocess.check_output(["fzn-gecode", "-help"])
    except FileNotFoundError:
        raise ImportError(
            "Please install 'fzn-gecode' to use the constraint task: 'apt install flatzinc'."
        )


def solve(fzn):
    """Call FlatZinc solver (Gecode) and process output."""
    import subprocess

    from problog.util import mktempfile

    fznfile = mktempfile(".fzn")
    with open(fznfile, "w") as f:
        f.write(fzn)
    cmd = ["fzn-gecode", "-a", fznfile]
    result = subprocess.check_output(cmd).decode()

    current = {}
    for line in result.split("\n"):
        if line.startswith("----"):
            yield current
            current = {}
        elif line.startswith("===="):
            return
        else:
            vr, vl = line.strip(";").split(" = ")
            vr = int(vr[1:])
            vl = 1 if vl == "true" else 0
            current[vr] = vl


def compress(formula, atoms):
    """Compress the formula by setting the values of the given atoms."""

    # TODO support cyclic programs

    formula.clear_labeled("constraint")

    out = LogicDAG()

    relevant = formula.extract_relevant()

    cyclic = set()

    translate = {}
    for i, n, t in formula:

        if not relevant[i]:
            continue
        elif t == "atom":
            if i in atoms:
                if atoms[i] == 1:
                    translate[i] = 0
                else:
                    translate[i] = None
            else:
                translate[i] = out.add_atom(*n)
        else:
            children = []
            for c in n.children:
                if abs(c) in translate:
                    tr = translate[abs(c)]
                    if c > 0:
                        children.append(tr)
                    else:
                        children.append(out.negate(tr))
                else:
                    cyclic.add(c)
            if t == "conj":
                translate[i] = out.add_and(children)
            else:
                translate[i] = out.add_or(children)

    if cyclic:
        debug("Original formula:\n", formula)
        debug("Cyclic nodes:", cyclic)
        debug("New formula so far:\n", out)
        raise RuntimeError(
            "Cycle detected in program. Cycles are currently not supported."
        )

    for q, n, l in formula.labeled():
        if l != "constraint":
            if n < 0:
                out.add_name(q, out.negate(translate[-n]), l)
            else:
                out.add_name(q, translate[n], l)
    return out


def create_sdd(solution, formula, verbose):
    if verbose:
        debug("Compressing...")
    new_formula = compress(formula, solution)

    if verbose:
        debug("Compiling...")
    current_sdd = SDD.create_from(new_formula)

    if verbose:
        debug("Evaluating...")
    return current_sdd


def enumerate_sdd(solution, formula, sdd, verbose):
    # target.lookup_evidence = {}
    weights = sdd.get_weights()
    # ev_nodes = []
    for k, v in solution.items():
        sddvar = formula.get_node(k).identifier
        sddatom = sdd.var2atom[sddvar]
        sddname = sdd.get_name(sddatom)
        weights[sddatom] = v
    return sdd


def enumerate_solutions(sols, verbose, formula, sdd=None):
    solutions = []
    for i, res in enumerate(sols):
        if verbose:
            debug("Evaluating solution %s/%s..." % (i + 1, len(sols)))

        if sdd is not None:
            current_sdd = enumerate_sdd(res, formula, sdd, verbose)
        else:
            current_sdd = create_sdd(res, formula, verbose)

        solutions.append(current_sdd.evaluate().items())
    return solutions


def check_prob_constraint(constraints):
    return any(
        index is not None and name.args[1].functor == "ensure_prob"
        for name, index in constraints
    )


def main(argv, handle_output=None):
    args = argparser().parse_args(argv)

    verbose = args.verbose
    filename = args.filename
    solutions = run(filename, verbose)

    for solution in solutions:
        for k, v in solution:
            if v > 0.0:
                print(k, v)

        print("----------")

    if solutions:
        print("==========")
    else:
        print("=== UNSATISFIABLE ===")

    if verbose:
        debug("Loading...")


def run(filename, verbose=None):
    find_solver()
    if verbose:
        debug("Loading...")
    problog_model = PrologFile(filename, factory=ConstraintFactory())

    engine = DefaultEngine()
    database = engine.prepare(problog_model)

    # Ground the constraints
    if verbose:
        debug("Grounding...")
    target = engine.ground(database, Term("constraint", None, None), label="constraint")

    queries = [q[0] for q in engine.query(database, Term("query", None))]
    for query in queries:
        target = engine.ground(database, query, label="query", target=target)

    if verbose is not None and verbose > 1:
        print(target, file=sys.stderr)

    has_prob_constraint = check_prob_constraint(target.get_names(label="constraint"))

    # Compile and turn into CP-problem
    sdd = None
    if has_prob_constraint:
        if verbose:
            debug("Probabilistic constraints detected.")
            debug("Compiling...")
        sdd = SDD.create_from(target)
        formula = sdd.to_formula()
    else:
        formula = LogicDAG()
        break_cycles(target, formula)

    if verbose:
        debug("Converting...")
    if sdd:
        fzn = formula_to_flatzinc_float(formula)
        sdd.clear_labeled("constraint")
    else:
        fzn = formula_to_flatzinc_bool(formula)

    if verbose is not None and verbose > 1:
        print(fzn, file=sys.stderr)
    if verbose:
        debug("Solving...")
    sols = list(solve(fzn))

    return enumerate_solutions(sols, verbose, formula, sdd)


def argparser():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("filename", metavar="MODEL")
    parser.add_argument("-v", "--verbose", action="count")
    return parser


if __name__ == "__main__":
    main(sys.argv[1:])
