#
# This file is part of pySMT.
#
#   Copyright 2014 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 io import StringIO

import pysmt.logics as logics
import pysmt.smtlib.commands as smtcmd
from pysmt.shortcuts import (
    Real,
    Plus,
    Symbol,
    Equals,
    And,
    Bool,
    Or,
    Not,
    Div,
    LT,
    LE,
    Int,
    ToReal,
    Iff,
    Exists,
    Times,
    FALSE,
    BVLShr,
    BVLShl,
    BVAShr,
    BV,
    BVAdd,
    BVULT,
    BVMul,
    Select,
    Array,
    Ite,
    String,
    Function,
    to_smtlib,
    ForAll,
)
from pysmt.shortcuts import Solver, get_env, qelim, get_model, TRUE, ExactlyOne
from pysmt.typing import REAL, BOOL, INT, BVType, FunctionType, ArrayType
from pysmt.test import (
    TestCase,
    skipIfSolverNotAvailable,
    skipIfNoSolverForLogic,
    skipIfNoQEForLogic,
)
from pysmt.test import main
from pysmt.exceptions import (
    ConvertExpressionError,
    PysmtValueError,
    PysmtTypeError,
    InternalSolverError,
)
from pysmt.test.examples import get_example_formulae
from pysmt.environment import Environment
from pysmt.rewritings import cnf_as_set
from pysmt.smtlib.parser import SmtLibParser
from pysmt.smtlib.commands import DECLARE_FUN
from pysmt.smtlib.script import SmtLibCommand
from pysmt.logics import get_closer_smtlib_logic
from pysmt.constants import Fraction


class TestRegressions(TestCase):
    @skipIfSolverNotAvailable("msat")
    @skipIfSolverNotAvailable("z3")
    def test_plus_converts_correctly_n_ary_functions(self):
        """Handling of Plus n-ary functionality.

        Only the first two elements were translated to the solver
        """

        a = Symbol("a", REAL)
        b = Symbol("b", REAL)
        c = Symbol("c", REAL)

        p1 = Plus(
            a,
            Real((1, 6)),
            b,
            c,
        )
        p2 = Plus(a, b, c, Real((1, 6)))

        self.assertValid(Equals(p1, p2))
        self.assertValid(Equals(p1, p2), solver_name="z3")
        self.assertValid(Equals(p1, p2), solver_name="msat")

    def test_substitute_memoization(self):
        a = Symbol("A", BOOL)
        b = Symbol("B", BOOL)

        f = And(a, b)
        g = f.substitute({a: Bool(True)})
        h = f.substitute({a: Bool(False)})

        self.assertNotEqual(h, g)

    @skipIfSolverNotAvailable("msat")
    def test_msat_bool_back_conversion(self):
        f = Symbol("A")
        with Solver(name="msat") as solver:
            solver.solve()
            val = solver.get_value(Symbol("A"))
            self.assertTrue(val.is_bool_constant())

    @skipIfSolverNotAvailable("msat")
    @skipIfSolverNotAvailable("z3")
    def test_conversion_of_fractions_in_z3(self):
        self.assertValid(
            Equals(Real(Fraction(1, 9)), Div(Real(1), Real(9))), solver_name="msat"
        )
        self.assertValid(
            Equals(Real(Fraction(1, 9)), Div(Real(1), Real(9))), solver_name="z3"
        )

    def test_simplifying_int_plus_changes_type_of_expression(self):
        varA = Symbol("At", INT)
        varB = Symbol("Bt", INT)
        get_type = get_env().stc.get_type
        f = Plus(varB, Int(1))
        old_type = get_type(f)
        f = f.simplify()
        new_type = get_type(f)
        self.assertEqual(new_type, old_type)

    @skipIfNoSolverForLogic(logics.QF_UFLIRA)
    def test_nary_operators_in_solver_converter(self):
        """Conversion of n-ary operators was not handled correctly by converters."""
        x = Symbol("x")
        r = Symbol("p", REAL)
        f_and_one = And(x)
        f_or_one = Or(x)
        f_plus_one = LT(Plus(r), Real(0))

        ten_x = [x, x, x, x, x, x, x, x, x, x]
        f_and_many = And(ten_x)
        f_or_many = Or(ten_x)
        f_plus_many = LT(Plus(r, r, r, r, r, r, r, r, r, r, r), Real(0))

        for name in get_env().factory.all_solvers(logic=logics.QF_BOOL):
            self.assertSat(f_and_one, solver_name=name)
            self.assertSat(f_or_one, solver_name=name)
            self.assertSat(f_and_many, solver_name=name)
            self.assertSat(f_or_many, solver_name=name)

        for name in get_env().factory.all_solvers(logic=logics.QF_UFLIRA):
            self.assertSat(f_plus_one, solver_name=name)
            self.assertSat(f_plus_many, solver_name=name)

    def test_dependencies_not_includes_toreal(self):
        p = Symbol("p", INT)
        r = ToReal(p)
        deps = r.get_free_variables()

        self.assertIn(p, deps)
        self.assertNotIn(r, deps)

    def test_infix_notation_wrong_le(self):
        p = Symbol("p", INT)
        get_env().enable_infix_notation = True
        self.assertEqual(LE(p, Int(2)), p <= Int(2))

    def test_multiple_declaration_w_same_functiontype(self):
        ft1 = FunctionType(REAL, [REAL])
        ft2 = FunctionType(REAL, [REAL])

        f1 = Symbol("f1", ft1)
        # The following raises an exception if not (ft1 == ft2)
        # since the same symbol has already been defined with
        # a "different" type.
        f1 = Symbol("f1", ft2)

    @skipIfSolverNotAvailable("z3")
    def test_z3_iff(self):
        z3 = Solver(name="z3")
        conv = z3.converter

        x, y = Symbol("x"), Symbol("y")
        term = conv.convert(Iff(x, y))
        back = conv.back(term)
        self.assertEqual(Iff(x, y), back)

    @skipIfSolverNotAvailable("msat")
    def test_msat_iff(self):
        msat = Solver(name="msat")
        conv = msat.converter

        x, y = Symbol("x"), Symbol("y")
        term = conv.convert(Iff(x, y))
        back = conv.back(term)
        # Mathsat can reorder variables...
        self.assertTrue(Iff(x, y) == back or Iff(y, x) == back)

    def test_multiple_exit(self):
        for sname in get_env().factory.all_solvers():
            # Multiple exits should be ignored
            s = Solver(name=sname)
            s.exit()
            s.exit()
            self.assertTrue(True)

    @skipIfNoQEForLogic(logics.LIA)
    def test_lia_qe_requiring_modulus(self):
        x = Symbol("x", INT)
        y = Symbol("y", INT)
        f = Exists([x], Equals(y, Times(x, Int(2))))
        with self.assertRaises(ConvertExpressionError):
            qelim(f)

        try:
            qelim(f)
        except ConvertExpressionError as ex:
            # The modulus operator must be there
            self.assertTrue(
                "%2" in str(ex.expression) or "int_mod_congr" in str(ex.expression)
            )

    @skipIfSolverNotAvailable("msat")
    def test_msat_partial_model(self):
        msat = Solver(name="msat")
        x, y = Symbol("x"), Symbol("y")
        msat.add_assertion(x)
        c = msat.solve()
        self.assertTrue(c)

        model = msat.get_model()
        self.assertNotIn(y, model)
        self.assertIn(x, model)
        msat.exit()

    @skipIfSolverNotAvailable("z3")
    def test_z3_model_iteration(self):
        x, y = Symbol("x"), Symbol("y")
        m = get_model(And(x, y), solver_name="z3")
        self.assertIsNotNone(m)

        for _, v in m:
            self.assertEqual(v, TRUE())

    def test_exactlyone_w_generator(self):
        x, y = Symbol("x"), Symbol("y")

        elems = [x, y]
        f1 = ExactlyOne(elems)
        f2 = ExactlyOne(e for e in elems)

        self.assertEqual(f1, f2)

    def test_determinism(self):
        def get_set(env):
            mgr = env.formula_manager
            r = set(mgr.Symbol("x%d" % i) for i in range(1000))
            for (f, _, _, _) in get_example_formulae(env):
                r |= set([f])
            return r

        # As first thing on the environment we build the set of formulae
        l1 = list(get_set(get_env()))

        # We try this ten times...
        for _ in range(10):
            # Do something to screw up memory layout...
            for y in (Symbol("y%d" % i) for i in range(1000)):
                self.assertIsNotNone(y)

            with Environment() as new_env:
                # As first thing on the environment we build the set of formulae
                l_test = list(get_set(new_env))

                # The ordering of the sets should be the same...
                for i, f in enumerate(l1):
                    nf = new_env.formula_manager.normalize(f)
                    self.assertEqual(nf, l_test[i])

    def test_is_one(self):
        self.assertTrue(Int(1).is_one())
        self.assertTrue(Real(1).is_one())
        self.assertTrue(Int(0).is_zero())
        self.assertTrue(Real(0).is_zero())

    def test_cnf_as_set(self):
        r = cnf_as_set(Symbol("x"))
        self.assertTrue(type(r) == frozenset)

    def test_substitute_to_real(self):
        p = Symbol("p", INT)
        f = LT(ToReal(p), Real(0))

        new_f = f.substitute({p: Real(1)}).simplify()
        self.assertEqual(new_f, Bool(False))

    def test_empty_string_symbol(self):
        with self.assertRaises(PysmtValueError):
            Symbol("")

    def test_smtlib_info_quoting(self):
        cmd = SmtLibCommand(smtcmd.SET_INFO, [":source", "This\nis\nmultiline!"])
        output = cmd.serialize_to_string()
        self.assertEqual(output, "(set-info :source |This\nis\nmultiline!|)")

    def test_parse_define_fun(self):
        smtlib_input = (
            "(declare-fun z () Bool)" "(define-fun .def_1 ((z Bool)) Bool (and z z))"
        )
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        parser.get_script(buffer_)

    def test_parse_define_fun_bind(self):
        smtlib_input = (
            "(declare-fun y () Bool)" "(define-fun .def_1 ((z Bool)) Bool (and z z))"
        )
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        parser.get_script(buffer_)

    def test_parse_bvx_var(self):
        """bvX is a valid identifier."""
        smtlib_input = """
        (declare-fun bv1 () (_ BitVec 8))
        (assert (bvult (_ bv0 8) (bvmul (bvadd bv1 (_ bv1 8)) (_ bv5 8))))
        (check-sat)"""
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        script = parser.get_script(buffer_)
        # Check Parsed result
        iscript = iter(script)
        cmd = next(iscript)
        self.assertEqual(cmd.name, DECLARE_FUN)
        bv1 = cmd.args[0]
        self.assertEqual(bv1.symbol_type().width, 8)
        cmd = next(iscript)
        parsed_f = cmd.args[0]
        target_f = BVULT(BV(0, 8), BVMul(BVAdd(bv1, BV(1, 8)), BV(5, 8)))
        self.assertEqual(parsed_f, target_f)

    def test_simplify_times(self):
        a, b = Real(5), Real((1, 5))
        f = Times(a, b).simplify()
        self.assertEqual(f.constant_value(), 1)

    @skipIfSolverNotAvailable("yices")
    def test_yices_push(self):
        with Solver(name="yices") as solver:
            solver.add_assertion(FALSE())
            res = solver.solve()
            self.assertFalse(res)
            solver.push()
            solver.add_assertion(TRUE())
            res = solver.solve()
            self.assertFalse(res)
            solver.pop()

    def test_qf_bool_smt2(self):
        # QF_BOOL does not exist in SMT-LIB
        # This test is to enforce the consistent choice of QF_UF
        close_l = get_closer_smtlib_logic(logics.QF_BOOL)
        self.assertEqual(close_l, logics.QF_UF)
        # For BOOL we use LRA
        close_l = get_closer_smtlib_logic(logics.BOOL)
        self.assertEqual(close_l, logics.LRA)

    def test_exactly_one_unpacking(self):
        s1, s2 = Symbol("x"), Symbol("y")
        f1 = ExactlyOne((s for s in [s1, s2]))
        f2 = ExactlyOne([s1, s2])
        f3 = ExactlyOne(s1, s2)

        self.assertEqual(f1, f2)
        self.assertEqual(f2, f3)

    @skipIfSolverNotAvailable("btor")
    def test_btor_bitwidth_bug_in_shift(self):
        # (384, 384, 9)
        # (x69 >> 1_384)
        s = Solver(name="btor")
        x69 = Symbol("x69", BVType(384))
        # BVLShr
        f = BVLShr(x69, BV(1, 384))
        c = s.converter.convert(f)
        self.assertIsNotNone(c)
        # BVLShl
        f = BVLShl(x69, BV(1, 384))
        c = s.converter.convert(f)
        self.assertIsNotNone(c)
        # BVAShr
        f = BVAShr(x69, BV(1, 384))
        c = s.converter.convert(f)
        self.assertIsNotNone(c)

    @skipIfSolverNotAvailable("btor")
    def test_btor_get_non_bool_value(self):
        with Solver(name="btor") as s:
            x = Symbol("x", BVType(16))
            s.add_assertion(Equals(x, BV(1, 16)))
            self.assertTrue(s.solve())
            self.assertEqual(s.get_value(Equals(x, BV(1, 16))), TRUE())
            self.assertEqual(s.get_value(BVAdd(x, BV(1, 16))), BV(2, 16))

    @skipIfSolverNotAvailable("btor")
    def test_btor_get_array_element(self):
        with Solver(name="btor") as s:
            x = Symbol("a", ArrayType(BVType(16), BVType(16)))
            s.add_assertion(Equals(Select(x, BV(1, 16)), BV(1, 16)))
            s.add_assertion(Equals(Select(x, BV(2, 16)), BV(3, 16)))
            self.assertTrue(s.solve())
            self.assertEqual(s.get_value(Select(x, BV(1, 16))), BV(1, 16))
            self.assertIsNotNone(s.get_value(x))

    def test_smtlib_define_fun_serialization(self):
        smtlib_input = "(define-fun init ((x Bool)) Bool (and x (and x (and x (and x (and x (and x x)))))))"
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        s = parser.get_script(buffer_)
        for c in s:
            res = c.serialize_to_string(daggify=False)
        self.assertEqual(res.replace("__x0", "x"), smtlib_input)

    @skipIfSolverNotAvailable("z3")
    def test_z3_nary_back(self):
        from z3 import Tactic

        r = Symbol("r", REAL)
        s = Symbol("s", REAL)
        t = Symbol("t", REAL)
        f = Equals(Times(r, s, t), Real(0))

        with Solver(name="z3") as solver:
            z3_f = solver.converter.convert(f)
            z3_f = Tactic("simplify", solver.z3.ctx)(z3_f).as_expr()
            fp = solver.converter.back(z3_f)
            self.assertValid(Iff(f, fp), (f, fp))

    def test_array_initialization_printing(self):
        self.assertEqual(
            str(Array(INT, Int(0), {Int(1): Int(2)})), "Array{Int, Int}(0)[1 := 2]"
        )

    @skipIfSolverNotAvailable("btor")
    def test_boolector_assumptions(self):
        with Solver(name="btor") as solver:
            x = Symbol("x")
            y = Symbol("y")
            solver.add_assertion(Or(x, y))
            solver.solve([Not(x), Not(y)])
            btor_notx = solver.converter.convert(Not(x))
            btor_noty = solver.converter.convert(Not(y))
            self.assertEqual(solver.btor.Failed(btor_notx, btor_noty), [True, True])

    def test_parse_declare_const(self):
        smtlib_input = """
        (declare-const s Int)
        (check-sat)"""
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        script = parser.get_script(buffer_)
        self.assertIsNotNone(script)

    @skipIfSolverNotAvailable("z3")
    def test_z3_conversion_ite(self):
        with Solver(name="z3") as solver:
            x = Symbol("x")
            y = Symbol("y")
            f = Ite(x, y, FALSE())
            solver.add_assertion(f)
            self.assertTrue(solver.solve())

    def test_parse_exception(self):
        from pysmt.exceptions import PysmtSyntaxError

        smtlib_input = "(declare-const x x x Int)" + "(check-sat)"
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        try:
            parser.get_script(buffer_)
            self.assertFalse(True)
        except PysmtSyntaxError as ex:
            self.assertEqual(ex.pos_info[0], 0)
            self.assertEqual(ex.pos_info[1], 19)

    def test_parse_bvconst_width(self):
        smtlib_input = "(assert (> #x10 #x10))"
        parser = SmtLibParser()
        buffer_ = StringIO(smtlib_input)
        expr = parser.get_script(buffer_).get_last_formula()
        const = expr.args()[0]
        self.assertEqual(const.bv_width(), 8, const.bv_width())

    def test_equality_typing(self):
        x = Symbol("x", BOOL)
        y = Symbol("y", BOOL)
        with self.assertRaises(PysmtTypeError):
            Equals(x, y)
        with self.assertRaises(PysmtTypeError):
            LE(x, y)

    def test_string_constant_quote_escaping_hr_printer(self):
        x = String('a"b')
        y = String('""')
        z = String('"')
        self.assertEqual('"a""b"', str(x))
        self.assertEqual('""""""', str(y))
        self.assertEqual('""""', str(z))

    def test_string_constant_quote_escaping_smtlib_printer(self):
        x = String('a"b')
        y = String('""')
        z = String('"')
        self.assertEqual('"a""b"', x.to_smtlib())
        self.assertEqual('""""""', y.to_smtlib())
        self.assertEqual('""""', z.to_smtlib())

    def test_string_constant_quote_escaping_parsing(self):
        script = """
        (declare-const x String)
        (assert (= x "a""b"))
        (assert (= x "\"\""))
        (assert (= x "\"\"\"\""))
        """

        p = SmtLibParser()
        buffer = StringIO(script)
        s = p.get_script(buffer)
        self.assertEqual('a"b', s.commands[1].args[0].arg(1).constant_value())
        self.assertEqual('"', s.commands[2].args[0].arg(1).constant_value())
        self.assertEqual('""', s.commands[3].args[0].arg(1).constant_value())

    def test_pysmt_syntax_error(self):
        from pysmt.exceptions import PysmtSyntaxError

        try:
            raise PysmtSyntaxError("'define-fun' expected", (5, 5))
        except PysmtSyntaxError as ex:
            self.assertEqual(str(ex), "Line 5, Col 5: 'define-fun' expected")

    def test_function_smtlib_print(self):
        f_t = FunctionType(BOOL, [BOOL])
        f0 = Symbol("f 0", f_t)
        f0_of_false = Function(f0, [Bool(False)])
        s = to_smtlib(f0_of_false, False)
        self.assertEqual(s, "(|f 0| false)")

    @skipIfSolverNotAvailable("yices")
    def test_yices_bv_overflow(self):
        smt_script = """
        (set-logic QF_BV)
        (declare-fun s0 () (_ BitVec 64))
        (define-fun s1 () (_ BitVec 64) #xFFFFFFFFFFFFFFFF)
        (define-fun s2 () Bool (= s1 s0))
        (assert s2)
        (check-sat)
        """
        from pysmt.smtlib.parser import get_formula_strict

        f = get_formula_strict(StringIO(smt_script))
        self.assertSat(f, solver_name="yices")

    @skipIfSolverNotAvailable("yices")
    def test_yices_bv_no_overflow(self):
        smt_script = """
        (set-logic QF_BV)
        (declare-fun s0 () (_ BitVec 64))
        (define-fun s1 () (_ BitVec 64) #x7FFFFFFFFFFFFFFF)
        (define-fun s2 () Bool (= s1 s0))
        (assert s2)
        (check-sat)
        """
        from pysmt.smtlib.parser import get_formula_strict

        f = get_formula_strict(StringIO(smt_script))
        self.assertSat(f, solver_name="yices")

    def test_get_atoms_array_select(self):
        a = Symbol("a", ArrayType(INT, BOOL))
        x = Symbol("x", INT)
        p = Symbol("p", BOOL)

        phi = And(Iff(Select(a, x), p), Equals(x, Int(1)))

        atoms = phi.get_atoms()

        self.assertEqual(len(atoms), 3)
        self.assertIn(Select(a, x), atoms)
        self.assertIn(p, atoms)
        self.assertIn(Equals(x, Int(1)), atoms)

    @skipIfSolverNotAvailable("yices")
    def test_yices_quantifier(self):
        x = Symbol("x", REAL)
        f = ForAll([x], LE(x, Real(0)))
        with self.assertRaises(InternalSolverError):
            with Solver(name="yices") as s:
                s.add_assertion(f)
                self.assertFalse(s.solve())


if __name__ == "__main__":
    main()
