#
# This file is part of pySMT.
#
#   Copyright 2015 Andrea Micheli and Marco Gario
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
from math import log, ceil

from pysmt.exceptions import SolverAPINotFound

try:
    import pyboolector
except ImportError:
    raise SolverAPINotFound


from pysmt.solvers.solver import (IncrementalTrackingSolver, UnsatCoreSolver,
                                  Converter, SolverOptions)
from pysmt.solvers.smtlib import SmtLibBasicSolver, SmtLibIgnoreMixin
from pysmt.solvers.eager import EagerModel
from pysmt.walkers import DagWalker
from pysmt.exceptions import (SolverReturnedUnknownResultError,
                              ConvertExpressionError, PysmtValueError)
from pysmt.decorators import clear_pending_pop, catch_conversion_error
from pysmt.logics import QF_BV, QF_UFBV, QF_ABV, QF_AUFBV, QF_AX, ARRAYS_CONST_LOGICS
from pysmt.constants import to_python_integer


class BoolectorOptions(SolverOptions):
    def __init__(self, **base_options):
        SolverOptions.__init__(self, **base_options)
        if self.random_seed is not None:
            raise PysmtValueError("BTOR Does not support Random Seed setting.")

        # Disabling Incremental usage is not allowed.
        # This needs to be set to 1
        self.incrementality = True
        self.internal_options = [pyboolector.BTOR_OPT_MODEL_GEN,
                                 pyboolector.BTOR_OPT_INCREMENTAL,
                                 pyboolector.BTOR_OPT_INCREMENTAL_SMT1,
                                 pyboolector.BTOR_OPT_INPUT_FORMAT,
                                 pyboolector.BTOR_OPT_OUTPUT_NUMBER_FORMAT,
                                 pyboolector.BTOR_OPT_OUTPUT_FORMAT,
                                 pyboolector.BTOR_OPT_ENGINE,
                                 pyboolector.BTOR_OPT_SAT_ENGINE,
                                 pyboolector.BTOR_OPT_AUTO_CLEANUP,
                                 pyboolector.BTOR_OPT_PRETTY_PRINT,
                                 pyboolector.BTOR_OPT_EXIT_CODES,
                                 pyboolector.BTOR_OPT_SEED,
                                 pyboolector.BTOR_OPT_VERBOSITY,
                                 pyboolector.BTOR_OPT_LOGLEVEL,
                                 pyboolector.BTOR_OPT_REWRITE_LEVEL,
                                 pyboolector.BTOR_OPT_SKELETON_PREPROC,
                                 pyboolector.BTOR_OPT_ACKERMANN,
                                 pyboolector.BTOR_OPT_BETA_REDUCE,
                                 pyboolector.BTOR_OPT_ELIMINATE_SLICES,
                                 pyboolector.BTOR_OPT_VAR_SUBST,
                                 pyboolector.BTOR_OPT_UCOPT,
                                 pyboolector.BTOR_OPT_MERGE_LAMBDAS,
                                 pyboolector.BTOR_OPT_EXTRACT_LAMBDAS,
                                 pyboolector.BTOR_OPT_NORMALIZE,
                                 pyboolector.BTOR_OPT_NORMALIZE_ADD,
                                 pyboolector.BTOR_OPT_FUN_PREPROP,
                                 pyboolector.BTOR_OPT_FUN_PRESLS,
                                 pyboolector.BTOR_OPT_FUN_DUAL_PROP,
                                 pyboolector.BTOR_OPT_FUN_DUAL_PROP_QSORT,
                                 pyboolector.BTOR_OPT_FUN_JUST,
                                 pyboolector.BTOR_OPT_FUN_JUST_HEURISTIC,
                                 pyboolector.BTOR_OPT_FUN_LAZY_SYNTHESIZE,
                                 pyboolector.BTOR_OPT_FUN_EAGER_LEMMAS,
                                 pyboolector.BTOR_OPT_FUN_STORE_LAMBDAS,
                                 pyboolector.BTOR_OPT_SLS_NFLIPS,
                                 pyboolector.BTOR_OPT_SLS_STRATEGY,
                                 pyboolector.BTOR_OPT_SLS_JUST,
                                 pyboolector.BTOR_OPT_SLS_MOVE_GW,
                                 pyboolector.BTOR_OPT_SLS_MOVE_RANGE,
                                 pyboolector.BTOR_OPT_SLS_MOVE_SEGMENT,
                                 pyboolector.BTOR_OPT_SLS_MOVE_RAND_WALK,
                                 pyboolector.BTOR_OPT_SLS_PROB_MOVE_RAND_WALK,
                                 pyboolector.BTOR_OPT_SLS_MOVE_RAND_ALL,
                                 pyboolector.BTOR_OPT_SLS_MOVE_RAND_RANGE,
                                 pyboolector.BTOR_OPT_SLS_MOVE_PROP,
                                 pyboolector.BTOR_OPT_SLS_MOVE_PROP_N_PROP,
                                 pyboolector.BTOR_OPT_SLS_MOVE_PROP_N_SLS,
                                 pyboolector.BTOR_OPT_SLS_MOVE_PROP_FORCE_RW,
                                 pyboolector.BTOR_OPT_SLS_MOVE_INC_MOVE_TEST,
                                 pyboolector.BTOR_OPT_SLS_USE_RESTARTS,
                                 pyboolector.BTOR_OPT_SLS_USE_BANDIT,
                                 pyboolector.BTOR_OPT_PROP_NPROPS,
                                 pyboolector.BTOR_OPT_PROP_USE_RESTARTS,
                                 pyboolector.BTOR_OPT_PROP_USE_BANDIT,
                                 pyboolector.BTOR_OPT_PROP_PATH_SEL,
                                 pyboolector.BTOR_OPT_PROP_PROB_USE_INV_VALUE,
                                 pyboolector.BTOR_OPT_PROP_PROB_FLIP_COND,
                                 pyboolector.BTOR_OPT_PROP_PROB_FLIP_COND_CONST,
                                 pyboolector.BTOR_OPT_PROP_FLIP_COND_CONST_DELTA,
                                 pyboolector.BTOR_OPT_PROP_FLIP_COND_CONST_NPATHSEL,
                                 pyboolector.BTOR_OPT_PROP_PROB_SLICE_KEEP_DC,
                                 pyboolector.BTOR_OPT_PROP_PROB_CONC_FLIP,
                                 pyboolector.BTOR_OPT_PROP_PROB_SLICE_FLIP,
                                 pyboolector.BTOR_OPT_PROP_PROB_EQ_FLIP,
                                 pyboolector.BTOR_OPT_PROP_PROB_AND_FLIP,
                                 pyboolector.BTOR_OPT_PROP_NO_MOVE_ON_CONFLICT,
                                 pyboolector.BTOR_OPT_AIGPROP_USE_RESTARTS,
                                 pyboolector.BTOR_OPT_AIGPROP_USE_BANDIT,
                                 pyboolector.BTOR_OPT_QUANT_SYNTH,
                                 pyboolector.BTOR_OPT_QUANT_DUAL_SOLVER,
                                 pyboolector.BTOR_OPT_QUANT_SYNTH_LIMIT,
                                 pyboolector.BTOR_OPT_QUANT_SYNTH_QI,
                                 pyboolector.BTOR_OPT_QUANT_DER,
                                 pyboolector.BTOR_OPT_QUANT_CER,
                                 pyboolector.BTOR_OPT_QUANT_MINISCOPE,
                                 pyboolector.BTOR_OPT_SORT_EXP,
                                 pyboolector.BTOR_OPT_SORT_AIG,
                                 pyboolector.BTOR_OPT_SORT_AIGVEC,
                                 pyboolector.BTOR_OPT_AUTO_CLEANUP_INTERNAL,
                                 pyboolector.BTOR_OPT_SIMPLIFY_CONSTRAINTS,
                                 pyboolector.BTOR_OPT_CHK_FAILED_ASSUMPTIONS,
                                 pyboolector.BTOR_OPT_CHK_MODEL,
                                 pyboolector.BTOR_OPT_CHK_UNCONSTRAINED,
                                 pyboolector.BTOR_OPT_PARSE_INTERACTIVE,
                                 pyboolector.BTOR_OPT_SAT_ENGINE_LGL_FORK,
                                 pyboolector.BTOR_OPT_SAT_ENGINE_CADICAL_FREEZE,
                                 pyboolector.BTOR_OPT_SAT_ENGINE_N_THREADS,
                                 pyboolector.BTOR_OPT_SIMP_NORMAMLIZE_ADDERS,
                                 pyboolector.BTOR_OPT_DECLSORT_BV_WIDTH,
                                 pyboolector.BTOR_OPT_QUANT_SYNTH_ITE_COMPLETE,
                                 pyboolector.BTOR_OPT_QUANT_FIXSYNTH,
                                 pyboolector.BTOR_OPT_RW_ZERO_LOWER_SLICE,
                                 pyboolector.BTOR_OPT_NONDESTR_SUBST]

    def _set_option(self, btor, name, value):
        available_options = {pyboolector.BoolectorOpt(btor, io).lng: io
                             for io in self.internal_options}
        try:
            btor.Set_opt(available_options[name], value)
        except TypeError:
            raise PysmtValueError("Error setting the option '%s=%s'"
                                  % (name, value))
        except pyboolector.BoolectorException:
            raise PysmtValueError("Error setting the option '%s=%s'"
                                  % (name, value))
        except KeyError:
            raise PysmtValueError("Unable to set non-existing option '%s'. "
                                  "The accepted options options are: %s"
                                  % (name, ", ".join(pyboolector.BoolectorOpt(btor, io).lng
                                                     for io in self.internal_options)))

    def __call__(self, solver):
        if self.generate_models:
            self._set_option(solver.btor, "model-gen", 1)
        else:
            self._set_option(solver.btor, "model-gen", 0)
        if self.incrementality:
            self._set_option(solver.btor, "incremental", 1)

        for k, v in self.solver_options.items():
            # Note: Options values in btor are mostly integers
            self._set_option(solver.btor, str(k), v)

# EOC BoolectorOptions


class BoolectorSolver(IncrementalTrackingSolver, UnsatCoreSolver,
                      SmtLibBasicSolver, SmtLibIgnoreMixin):

    LOGICS = [QF_BV, QF_UFBV, QF_ABV, QF_AUFBV, QF_AX] + \
        list(filter(lambda l: l.name in {
             'QF_ABV*', 'QF_AUFBV*', 'QF_AX*'}, ARRAYS_CONST_LOGICS))
    OptionsClass = BoolectorOptions

    def __init__(self, environment, logic, **options):
        IncrementalTrackingSolver.__init__(self,
                                           environment=environment,
                                           logic=logic,
                                           **options)
        self.btor = pyboolector.Boolector()
        self.options(self)
        self.converter = BTORConverter(environment, self.btor)
        self.mgr = environment.formula_manager
        self.declarations = {}
        self._named_assertions = {}
        return

# EOC BoolectorOptions

        pass

    @clear_pending_pop
    def _reset_assertions(self):
        self.btor = pyboolector.Boolector()
        self.options(self)
        self.converter = BTORConverter(self.environment, self.btor)
        self.declarations = {}

    @clear_pending_pop
    def declare_variable(self, var):
        raise NotImplementedError

    @clear_pending_pop
    def _add_assertion(self, formula, named=None):
        self._assert_is_boolean(formula)
        term = self.converter.convert(formula)
        if self.options.unsat_cores_mode is None:
            self.btor.Assert(term)
        else:
            if self.options.unsat_cores_mode == "named" and \
               named is not None:
                self._named_assertions[formula] = named
            # need to use assumptions to get unsat cores
            self.btor.Assume(term)
        return formula

    def get_model(self):
        assignment = {}
        for s, _ in self.converter.declared_vars.items():
            assignment[s] = self.get_value(s)
        return EagerModel(assignment=assignment, environment=self.environment)

    @clear_pending_pop
    def _solve(self, assumptions=None):
        if assumptions is not None:
            btor_assumptions = [self.converter.convert(a) for a in assumptions]
            self.btor.Assume(*btor_assumptions)

        res = self.btor.Sat()

        # need to re-add assumptions if in unsat-core mode
        # which uses Assume instead of Assert
        if self.options.unsat_cores_mode is not None:
            for a in self._assertion_stack:
                self._add_assertion(a)

        if res == self.btor.SAT:
            return True
        elif res == self.btor.UNSAT:
            return False
        else:
            raise SolverReturnedUnknownResultError

    def get_unsat_core(self):
        """After a call to solve() yielding UNSAT, returns the unsat core as a
        set of formulae"""
        self._check_unsat_core_config()

        if self.options.unsat_cores_mode == 'all':
            unsat_core = set()
            # relies on this assertion stack being ordered
            assert isinstance(self._assertion_stack, list)
            btor_assertions = [self.converter.convert(
                a) for a in self._assertion_stack]
            in_unsat_core = self.btor.Failed(*btor_assertions)
            for a, in_core in zip(self._assertion_stack, in_unsat_core):
                if in_core:
                    unsat_core.add(a)
            return unsat_core
        else:
            return self.get_named_unsat_core().values()

    def get_named_unsat_core(self):
        """After a call to solve() yielding UNSAT, returns the unsat core as a
        dict of names to formulae"""
        self._check_unsat_core_config()

        if self.options.unsat_cores_mode == "named":
            unsat_core = {}
            # relies on this assertion stack being ordered
            assert isinstance(self._assertion_stack, list)
            btor_named_assertions = [self.converter.convert(
                a) for a in self._named_assertions.keys()]
            in_unsat_core = self.btor.Failed(*btor_named_assertions)
            for a, in_core in zip(self._assertion_stack, in_unsat_core):
                if in_core:
                    name = self._named_assertions[a]
                    unsat_core[name] = a
            return unsat_core
        else:
            return dict(("_a%d" % i, f)
                        for i, f in enumerate(self.get_unsat_core()))

    @clear_pending_pop
    def _push(self, levels=1):
        self.btor.Push(levels)

    @clear_pending_pop
    def _pop(self, levels=1):
        self.btor.Pop(levels)

    def print_model(self, name_filter=None):
        for var in self.declarations:
            if name_filter is None or not var.symbol_name().startswith(name_filter):
                print("%s = %s" % (var.symbol_name(), self.get_value(var)))

    def get_value(self, item):
        self._assert_no_function_type(item)
        itype = item.get_type()
        titem = self.converter.convert(item)

        def _back_bv_func(width):
            def _back_bv(assignment):
                return self.mgr.BV(assignment, width)
            return _back_bv

        def _back_bool(assignment):
            return self.mgr.Bool(bool(int(assignment)))

        if itype.is_bv_type():
            return _back_bv_func(item.bv_width())(titem.assignment)
        elif itype.is_bool_type():
            return _back_bool(titem.assignment)
        else:
            assert itype.is_array_type()
            assert itype.index_type.is_bv_type() or itype.index_type.is_bool_type()
            assert itype.elem_type.is_bv_type() or itype.elem_type.is_bool_type()

            if itype.index_type.is_bv_type():
                _back_index = _back_bv_func(itype.index_type.width)
            else:
                _back_index = _back_bool

            if itype.elem_type.is_bv_type():
                _back_elem = _back_bv_func(itype.elem_type.width)
                default_elem = self.mgr.BV(0, itype.elem_type.width)
            else:
                _back_elem = _back_bool
                default_elem = self.mgr.Bool(False)

            assign = {}
            for (idx, val) in titem.assignment:
                elem_value = _back_elem(val)
                if idx == '*':
                    default_elem = elem_value
                else:
                    assign[_back_index(idx)] = elem_value

            return self.mgr.Array(itype.index_type,
                                  default_elem, assign)

    def _exit(self):
        del self.btor


class BTORConverter(Converter, DagWalker):

    def __init__(self, environment, btor):
        DagWalker.__init__(self, environment)
        self.mgr = environment.formula_manager
        self._get_type = environment.stc.get_type
        self._back_memoization = {}
        self._btor = btor
        self.declared_vars = {}
        self.function_declarations = {}
        return

    @catch_conversion_error
    def convert(self, formula):
        return self.walk(formula)

    def back(self, expr):
        raise NotImplementedError

    def declare_function(self, formula):
        name = formula.function_name()
        if name not in self.function_declarations:
            tp = self._type_to_btor(name.symbol_type())
            decl = self._btor.UF(tp)
            self.function_declarations[name] = decl
        else:
            decl = self.function_declarations[name]
        return decl

    def walk_and(self, formula, args, **kwargs):
        assert len(args) >= 2
        res = self._btor.And(args[0], args[1])
        for conj in args[2:]:
            res = self._btor.And(res, conj)
        return res

    def walk_or(self, formula, args, **kwargs):
        assert len(args) >= 2
        res = self._btor.Or(args[0], args[1])
        for disj in args[2:]:
            res = self._btor.Or(res, disj)
        return res

    def walk_not(self, formula, args, **kwargs):
        return self._btor.Not(args[0])

    def walk_symbol(self, formula, **kwargs):
        symbol_type = formula.symbol_type()
        if symbol_type.is_bool_type():
            res = self._btor.Var(self._btor.BitVecSort(1),
                                 formula.symbol_name())
        elif symbol_type.is_real_type():
            raise ConvertExpressionError
        elif symbol_type.is_int_type():
            raise ConvertExpressionError
        elif symbol_type.is_array_type():
            # BTOR supports only Arrays of Type (BV, BV)
            index_type = symbol_type.index_type
            elem_type = symbol_type.elem_type
            if not (index_type.is_bv_type() or index_type.is_bool_type()) or not (elem_type.is_bv_type() or elem_type.is_bool_type()):
                raise ConvertExpressionError("BTOR supports only Array(BV,BV). "
                                             "Type '%s' was given." % str(symbol_type))
            index_width = index_type.width if index_type.is_bv_type() else 1
            elem_width = elem_type.width if elem_type.is_bv_type() else 1
            res = self._btor.Array(self._btor.ArraySort(self._btor.BitVecSort(index_width),
                                                        self._btor.BitVecSort(elem_width)),
                                   formula.symbol_name())
        elif symbol_type.is_bv_type():
            res = self._btor.Var(self._btor.BitVecSort(formula.bv_width()),
                                 formula.symbol_name())
        else:
            raise SolverReturnedUnknownResultError("Unknown type for BTOR")
        self.declared_vars[formula] = res
        return res

    def walk_iff(self, formula, args, **kwargs):
        return self._btor.Iff(*args)

    def walk_implies(self, formula, args, **kwargs):
        return self._btor.Implies(*args)

    def walk_ite(self, formula, args, **kwargs):
        return self._btor.Cond(*args)

    def walk_bool_constant(self, formula, **kwargs):
        return self._btor.Const(formula.constant_value())

    def walk_equals(self, formula, args, **kwargs):
        args = self._extend_bv_equal_width(*args)
        return self._btor.Eq(*args)

    def walk_function(self, formula, args, **kwargs):
        _uf = self.declare_function(formula)
        return _uf(*args)

    def walk_bv_constant(self, formula, **kwargs):
        value = to_python_integer(formula.constant_value())
        width = formula.bv_width()
        return self._btor.Const(value, width)

    def walk_bv_ult(self, formula, args, **kwargs):
        return self._btor.Ult(args[0], args[1])

    def walk_bv_ule(self, formula, args, **kwargs):
        return self._btor.Ulte(args[0], args[1])

    def walk_bv_concat(self, formula, args, **kwargs):
        return self._btor.Concat(args[0], args[1])

    def walk_bv_extract(self, formula, args, **kwargs):
        start = formula.bv_extract_start()
        end = formula.bv_extract_end()
        return self._btor.Slice(args[0], end, start)

    def walk_bv_or(self, formula, args, **kwargs):
        return self.walk_or(formula, args, **kwargs)

    def walk_bv_not(self, formula, args, **kwargs):
        return self.walk_not(formula, args, **kwargs)

    def walk_bv_and(self, formula, args, **kwargs):
        return self.walk_and(formula, args, **kwargs)

    def walk_bv_xor(self, formula, args, **kwargs):
        return self._btor.Xor(*args)

    def walk_bv_add(self, formula, args, **kwargs):
        return self._btor.Add(args[0], args[1])

    def walk_bv_sub(self, formula, args, **kwargs):
        return self._btor.Sub(args[0], args[1])

    def walk_bv_neg(self, formula, args, **kwargs):
        return -args[0]

    def walk_bv_mul(self, formula, args, **kwargs):
        return args[0]*args[1]

    def walk_bv_udiv(self, formula, args, **kwargs):
        return args[0] / args[1]

    def walk_bv_urem(self, formula, args, **kwargs):
        return self._btor.Urem(args[0], args[1])

    def walk_bv_lshl(self, formula, args, **kwargs):
        return self._btor.Sll(args[0], args[1])

    def walk_bv_lshr(self, formula, args, **kwargs):
        return self._btor.Srl(args[0], args[1])

    def walk_bv_rol(self, formula, args, **kwargs):
        return self._btor.Rol(args[0],
                              formula.bv_rotation_step())

    def walk_bv_ror(self, formula, args, **kwargs):
        return self._btor.Ror(args[0],
                              formula.bv_rotation_step())

    def walk_bv_zext(self, formula, args, **kwargs):
        return self._btor.Uext(args[0], formula.bv_extend_step())

    def walk_bv_sext(self, formula, args, **kwargs):
        return self._btor.Sext(args[0], formula.bv_extend_step())

    def walk_bv_slt(self, formula, args, **kwargs):
        return self._btor.Slt(args[0], args[1])

    def walk_bv_sle(self, formula, args, **kwargs):
        return self._btor.Slte(args[0], args[1])

    def walk_bv_comp(self, formula, args, **kwargs):
        return self._btor.Eq(args[0], args[1])

    def walk_bv_sdiv(self, formula, args, **kwargs):
        return self._btor.Sdiv(args[0], args[1])

    def walk_bv_srem(self, formula, args, **kwargs):
        return self._btor.Srem(args[0], args[1])

    def walk_bv_ashr(self, formula, args, **kwargs):
        return self._btor.Sra(args[0], args[1])

    def walk_array_store(self, formula, args, **kwargs):
        return self._btor.Write(args[0], args[1], args[2])

    def walk_array_select(self, formula, args, **kwargs):
        return self._btor.Read(args[0], args[1])

    def walk_array_value(self, formula, args, **kwargs):
        arr_type = self.env.stc.get_type(formula)
        arr_sort = self._type_to_btor(arr_type)
        term = self._btor.ConstArray(arr_sort, args[0])

        for i in range(1, len(args), 2):
            term = self.walk_array_store(None, (term, args[i], args[i+1]))

        return term

    def _type_to_btor(self, tp):
        if tp.is_bool_type():
            return self._btor.BoolSort()
        elif tp.is_real_type():
            raise ConvertExpressionError
        elif tp.is_int_type():
            raise ConvertExpressionError
        elif tp.is_bv_type():
            return self._btor.BitVecSort(tp.width)
        elif tp.is_array_type():
            return self._btor.ArraySort(self._type_to_btor(tp.index_type),
                                        self._type_to_btor(tp.elem_type))
        else:
            assert tp.is_function_type(), "Unsupported type '%s'" % tp
            stps = [self._type_to_btor(x) for x in tp.param_types]
            rtp = self._type_to_btor(tp.return_type)
            return self._btor.FunSort(stps, rtp)

    def _extend_bv_pow2(self, btor_formula, signed=False):
        """BTOR requires that many operands have width that is a power of 2"""
        w = btor_formula.width
        target_w = 2**int(ceil(log(w, 2)))
        # Skip if width is ok
        if target_w == w:
            return btor_formula
        if signed:
            return self._btor.Sext(btor_formula, (target_w-w))
        else:
            return self._btor.Uext(btor_formula, (target_w-w))

    def _extend_bv_equal_width(self, arg1, arg2):
        if arg1.width == arg2.width:
            return (arg1, arg2)
        elif arg1.width > arg2.width:
            ext = arg1.width - arg2.width
            return (arg1,
                    self._btor.Uext(arg2, ext))
        elif arg1.width < arg2.width:
            ext = arg2.width - arg1.width
            return (self._btor.Uext(arg1, ext),
                    arg2)
