
#include "LogicBlock/SMTLibLogicBlock.hpp"
#include "LogicBlock/SMTLibLogicModel.hpp"
#include "LogicTerm/Logic.hpp"
#include "LogicTerm/LogicTerm.hpp"

#include "gmock/gmock.h"

using namespace logicbase;

class TestSMTLib: public testing::TestWithParam<logicbase::OpType> {};

TEST(TestSMTLib, TestPrint) {
    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, std::cout);

    const LogicTerm a = smtLibLogic->makeVariable("a", CType::BOOL);
    const LogicTerm b = smtLibLogic->makeVariable("b", CType::INT);
    const LogicTerm c = smtLibLogic->makeVariable("c", CType::REAL);
    const LogicTerm d = smtLibLogic->makeVariable("d", CType::BITVECTOR, 8);

    const LogicTerm t = a && (b == LogicTerm(1));
    t.print(std::cout);
    t.prettyPrint(std::cout);

    const LogicTerm f = (c == LogicTerm(1.0)) && (d == LogicTerm(1, 8));
    f.print(std::cout);
    f.prettyPrint(std::cout);

    smtLibLogic->assertFormula(t);
    smtLibLogic->assertFormula(f);

    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_BV);
    smtLibLogic->produceInstance();
}

TEST(TestSMTLib, TestBasicBooleanAnd) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_UF)\n"
           << "(declare-const p_0 Bool)\n"
           << "(assert (and p_0 (not p_0)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm p = smtLibLogic->makeVariable("p", CType::BOOL);
    const LogicTerm t = p && !p;
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_UF);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestBasicBooleanOrXor) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_UF)\n"
           << "(declare-const p_0 Bool)\n"
           << "(assert (xor p_0 (not p_0)))\n"
           << "(assert (or p_0 (not p_0)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm p  = smtLibLogic->makeVariable("p", CType::BOOL);
    const LogicTerm t  = p || !p;
    const LogicTerm t2 = p != !p;
    smtLibLogic->assertFormula(t);
    smtLibLogic->assertFormula(t2);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_UF);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestBasicBooleanImpl) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_UF)\n"
           << "(declare-const p_0 Bool)\n"
           << "(assert (=> p_0 (not p_0)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm p = smtLibLogic->makeVariable("p", CType::BOOL);
    const LogicTerm t = LogicTerm::implies(p, !p);
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_UF);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestBasicBooleanIte) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_UF)\n"
           << "(declare-const p_0 Bool)\n"
           << "(assert (ite p_0 (not p_0) p_0))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm p = smtLibLogic->makeVariable("p", CType::BOOL);
    const LogicTerm t = LogicTerm::ite(p, !p, p);
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_UF);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestIntegerArithmetic) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 Int)\n"
           << "(declare-const y_1 Int)\n"
           << "(assert (= (- x_0 y_1) (+ (+ x_0 (- 0 y_1)) 1)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x = smtLibLogic->makeVariable("x", CType::INT);
    const LogicTerm y = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm t = (x - y) == (x + (LogicTerm(0) - y) + LogicTerm(1));
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestIntegerArithmeticMul) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 Int)\n"
           << "(declare-const y_1 Int)\n"
           << "(assert (= (+ x_0 (* y_1 2)) 20))\n"
           << "(assert (= (- x_0 y_1) 2))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x  = smtLibLogic->makeVariable("x", CType::INT);
    const LogicTerm y  = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm t  = (x + (LogicTerm(2) * y)) == LogicTerm(20);
    const LogicTerm t2 = (x - y) == LogicTerm(2);
    smtLibLogic->assertFormula(t);
    smtLibLogic->assertFormula(t2);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestIntegerArithmeticDiv) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 Int)\n"
           << "(declare-const y_1 Int)\n"
           << "(assert (= (+ x_0 (* y_1 2)) 20))\n"
           << "(assert (= (div x_0 y_1) 2))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x  = smtLibLogic->makeVariable("x", CType::INT);
    const LogicTerm y  = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm t  = (x + (LogicTerm(2) * y)) == LogicTerm(20);
    const LogicTerm t2 = (x / y) == LogicTerm(2);
    smtLibLogic->assertFormula(t);
    smtLibLogic->assertFormula(t2);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestIntegerArithmeticGTLT) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 Int)\n"
           << "(declare-const y_1 Int)\n"
           << "(assert (< (+ x_0 (* y_1 2)) 20))\n"
           << "(assert (> (div x_0 y_1) 2))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x  = smtLibLogic->makeVariable("x", CType::INT);
    const LogicTerm y  = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm t  = (x + (LogicTerm(2) * y)) < LogicTerm(20);
    const LogicTerm t2 = (x / y) > LogicTerm(2);
    smtLibLogic->assertFormula(t);
    smtLibLogic->assertFormula(t2);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestIntegerArithmeticGTELTE) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 Int)\n"
           << "(declare-const y_1 Int)\n"
           << "(assert (<= (+ x_0 (* y_1 2)) 20))\n"
           << "(assert (>= (div x_0 y_1) 2))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x  = smtLibLogic->makeVariable("x", CType::INT);
    const LogicTerm y  = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm t  = (x + (LogicTerm(2) * y)) <= LogicTerm(20);
    const LogicTerm t2 = (x / y) >= LogicTerm(2);
    smtLibLogic->assertFormula(t);
    smtLibLogic->assertFormula(t2);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestBitVectorArithmeticAnd) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 (_ BitVec 8))\n"
           << "(declare-const y_1 (_ BitVec 8))\n"
           << "(assert (= x_0 (bvand x_0 y_1)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x = smtLibLogic->makeVariable("x", CType::BITVECTOR, 8);
    const LogicTerm y = smtLibLogic->makeVariable("y", CType::BITVECTOR, 8);
    const LogicTerm t = x == (x & y);
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestBitVectorArithmeticOr) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 (_ BitVec 8))\n"
           << "(declare-const y_1 (_ BitVec 8))\n"
           << "(assert (= x_0 (bvor x_0 y_1)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x = smtLibLogic->makeVariable("x", CType::BITVECTOR, 8);
    const LogicTerm y = smtLibLogic->makeVariable("y", CType::BITVECTOR, 8);
    const LogicTerm t = x == (x | y);
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestBitVectorArithmeticXor) {
    std::stringstream ss;

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    std::stringstream output{};
    output << "; Generated by " << PROJECT_NAME << " " << PROJECT_VER << "\n"
           << "(set-option :produce-models true)\n"
           << "(set-logic QF_LIA)\n"
           << "(declare-const x_0 (_ BitVec 8))\n"
           << "(declare-const y_1 (_ BitVec 8))\n"
           << "(assert (= x_0 (bvxor x_0 y_1)))\n"
           << "(check-sat)\n"
           << "(get-model)\n"
           << "(exit)\n";

    const LogicTerm x = smtLibLogic->makeVariable("x", CType::BITVECTOR, 8);
    const LogicTerm y = smtLibLogic->makeVariable("y", CType::BITVECTOR, 8);
    const LogicTerm t = x == (x ^ y);
    smtLibLogic->assertFormula(t);
    smtLibLogic->setOutputLogic(smtliblogic::SMTLibLogic::QF_LIA);
    smtLibLogic->produceInstance();

    ASSERT_EQ(ss.str(), output.str());
}

TEST(TestSMTLib, TestModelInt) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () Int 8)\n"
                               "(define-fun y () Int 6)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::INT);
    const LogicTerm y           = smtLibLogic->makeVariable("y", CType::INT);

    ASSERT_EQ(model.getIntValue(x, nullptr), 8);
    ASSERT_EQ(model.getIntValue(y, nullptr), 6);
}

TEST(TestSMTLib, TestModelBool) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () Bool true)\n"
                               "(define-fun y () Bool false)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::BOOL);
    const LogicTerm y           = smtLibLogic->makeVariable("y", CType::BOOL);

    ASSERT_EQ(model.getBoolValue(x, nullptr), true);
    ASSERT_EQ(model.getBoolValue(y, nullptr), false);
}

TEST(TestSMTLib, TestModelFloat) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () Real 1.0)\n"
                               "(define-fun y () Real 0.7)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::BOOL);
    const LogicTerm y           = smtLibLogic->makeVariable("y", CType::BOOL);

    ASSERT_EQ(model.getRealValue(x, nullptr), 1);
    ASSERT_EQ(model.getRealValue(y, nullptr), 0.7);
}

TEST(TestSMTLib, TestModelBitVector) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () (_ BitVector 8) #x00)\n"
                               "(define-fun y () (_ BitVector 8) #x10)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::BITVECTOR, 8);
    const LogicTerm y           = smtLibLogic->makeVariable("y", CType::BITVECTOR, 8);

    ASSERT_EQ(model.getBitvectorValue(x, nullptr), 0);
    ASSERT_EQ(model.getBitvectorValue(y, nullptr), 2);
}

TEST(TestSMTLib, TestModelWrongType) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () (_ BitVector 8) #x00)\n"
                               "(define-fun y () Int 1)\n"
                               "(define-fun z () Bool true)\n"
                               "(define-fun a () Real 1.0)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::BITVECTOR, 16);
    const LogicTerm y           = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm z           = smtLibLogic->makeVariable("z", CType::BOOL);
    const LogicTerm a           = smtLibLogic->makeVariable("a", CType::REAL);

    ASSERT_THROW(model.getBitvectorValue(x, nullptr), std::runtime_error);
    ASSERT_THROW(model.getBitvectorValue(y, nullptr), std::runtime_error);
    ASSERT_THROW(model.getBitvectorValue(z, nullptr), std::runtime_error);
    ASSERT_THROW(model.getBitvectorValue(a, nullptr), std::runtime_error);

    ASSERT_THROW(model.getIntValue(a, nullptr), std::runtime_error);
    ASSERT_THROW(model.getIntValue(x, nullptr), std::runtime_error);
    ASSERT_THROW(model.getIntValue(z, nullptr), std::runtime_error);

    ASSERT_THROW(model.getBoolValue(x, nullptr), std::runtime_error);
    ASSERT_THROW(model.getBoolValue(y, nullptr), std::runtime_error);
    ASSERT_THROW(model.getBoolValue(a, nullptr), std::runtime_error);

    ASSERT_THROW(model.getRealValue(x, nullptr), std::runtime_error);
    ASSERT_THROW(model.getRealValue(y, nullptr), std::runtime_error);
    ASSERT_THROW(model.getRealValue(z, nullptr), std::runtime_error);
}

TEST(TestSMTLib, TestModelBVSizeNotFound) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () (_ BitVector ) #x00)\n"
                               "(define-fun y () Int 1)\n"
                               "(define-fun z () Bool true)\n"
                               "(define-fun a () Real 1.0)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::BITVECTOR, 16);

    ASSERT_THROW(model.getBitvectorValue(x, nullptr), std::runtime_error);
}

TEST(TestSMTLib, TestModelNotAVariable) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () (_ BitVector 8) #x00)\n"
                               "(define-fun y () Int 1)\n"
                               "(define-fun z () Bool true)\n"
                               "(define-fun a () Real 1.0)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);

    ASSERT_THROW(model.getValue(LogicTerm(true), nullptr), std::runtime_error);
}

TEST(TestSMTLib, TestModelUnsupportedCType) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () (_ BitVector 8) #x00)\n"
                               "(define-fun y () Int 1)\n"
                               "(define-fun z () Bool true)\n"
                               "(define-fun a () Real 1.0)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::ARRAY);

    ASSERT_THROW(model.getValue(x, nullptr), std::runtime_error);
}

TEST(TestSMTLib, TestModelGetValue) {
    std::stringstream ss;

    const std::string output = "(\n"
                               "(define-fun x () (_ BitVector 8) #x00)\n"
                               "(define-fun y () Int 1)\n"
                               "(define-fun z () Bool true)\n"
                               "(define-fun a () Real 1.0)\n"
                               ")\n";

    smtliblogic::SMTLibLogicModel model(output);

    const auto      smtLibLogic = std::make_unique<smtliblogic::SMTLogicBlock>(false, ss);
    const LogicTerm x           = smtLibLogic->makeVariable("x", CType::BITVECTOR, 8);
    const LogicTerm y           = smtLibLogic->makeVariable("y", CType::INT);
    const LogicTerm z           = smtLibLogic->makeVariable("z", CType::BOOL);
    const LogicTerm a           = smtLibLogic->makeVariable("a", CType::REAL);

    ASSERT_EQ(model.getValue(x, nullptr).getBitVectorValue(), 0);
    ASSERT_EQ(model.getValue(y, nullptr).getIntValue(), 1);
    ASSERT_EQ(model.getValue(z, nullptr).getBoolValue(), true);
    ASSERT_EQ(model.getValue(a, nullptr).getFloatValue(), 1.0);
}
