#include "ast/sls/bv_sls_eval.h"
#include "ast/rewriter/th_rewriter.h"
#include "ast/reg_decl_plugins.h"
#include "ast/ast_pp.h"

namespace bv {
    class sls_test {
        ast_manager& m;
        bv_util bv;

    public:
        sls_test(ast_manager& m):
            m(m),
            bv(m)
        {}

        void check_eval(expr* a, expr* b, unsigned j) {
            auto es = create_exprs(a, b, j);
            for (expr* e : es)
                check_eval(e);
        }

        void check_eval(expr* e) {
            std::function<bool(expr*, unsigned)> value = [](expr*, unsigned) {
                return false;
            };
            expr_ref_vector es(m);
            bv_util bv(m);
            es.push_back(e);
            sls_eval ev(m);
            ev.init_eval(es, value);
            th_rewriter rw(m);
            expr_ref r(e, m);
            rw(r);

            if (bv.is_bv(e)) {
                auto const& val = ev.wval0(e);
                rational n1, n2;

                val.get_value(val.bits, n1);

                VERIFY(bv.is_numeral(r, n2));
                if (n1 != n2) {
                    verbose_stream() << mk_pp(e, m) << " computed value " << val << "\n";
                    verbose_stream() << "should be " << n2 << "\n";
                }
                SASSERT(n1 == n2);
                VERIFY(n1 == n2);
            }
            else if (m.is_bool(e)) {
                auto val1 = ev.bval0(e);
                auto val2 = m.is_true(r);
                if (val1 != val2) {
                    verbose_stream() << mk_pp(e, m) << " computed value " << val1 << " at odds with definition\n";
                }
                SASSERT(val1 == val2);
                VERIFY(val1 == val2);
            }
        }

        expr_ref_vector create_exprs(expr* a, expr* b, unsigned j) {
            expr_ref_vector result(m);
            result.push_back(bv.mk_bv_add(a, b))
                .push_back(bv.mk_bv_mul(a, b))
                .push_back(bv.mk_bv_sub(a, b))
                .push_back(bv.mk_bv_udiv(a, b))
                .push_back(bv.mk_bv_sdiv(a, b))
                .push_back(bv.mk_bv_srem(a, b))
                .push_back(bv.mk_bv_urem(a, b))
                .push_back(bv.mk_bv_smod(a, b))
                .push_back(bv.mk_bv_shl(a, b))
                .push_back(bv.mk_bv_ashr(a, b))
                .push_back(bv.mk_bv_lshr(a, b))
                .push_back(bv.mk_bv_and(a, b))
                .push_back(bv.mk_bv_or(a, b))
                .push_back(bv.mk_bv_xor(a, b))
                .push_back(bv.mk_bv_neg(a))
                .push_back(bv.mk_bv_not(a))
                .push_back(bv.mk_bvumul_ovfl(a, b))
                .push_back(bv.mk_bvumul_no_ovfl(a, b))
                .push_back(bv.mk_zero_extend(3, a))
                .push_back(bv.mk_sign_extend(3, a))
                .push_back(bv.mk_ule(a, b))
                .push_back(bv.mk_sle(a, b))
                .push_back(bv.mk_concat(a, b))
                .push_back(bv.mk_extract(4, 2, a))
                .push_back(bv.mk_bvuadd_ovfl(a, b))
                .push_back(bv.mk_bv_rotate_left(a, j))
                .push_back(bv.mk_bv_rotate_right(a, j))
                .push_back(bv.mk_bv_rotate_left(a, b))
                .push_back(bv.mk_bv_rotate_right(a, b))
    //            .push_back(bv.mk_bvsadd_ovfl(a, b))
    //            .push_back(bv.mk_bvneg_ovfl(a))
    //            .push_back(bv.mk_bvsmul_no_ovfl(a, b))
    //            .push_back(bv.mk_bvsmul_no_udfl(a, b))
    //            .push_back(bv.mk_bvsmul_ovfl(a, b))
    //            .push_back(bv.mk_bvsdiv_ovfl(a, b))
                ;
            return result;
        }


        // e = op(a, b), 
        // update value of a to "random"
        // repair a based on computed values.
        void check_repair(expr* a, expr* b, unsigned j) {
            expr_ref x(m.mk_const("x", bv.mk_sort(bv.get_bv_size(a))), m);
            expr_ref y(m.mk_const("y", bv.mk_sort(bv.get_bv_size(b))), m);
            auto es1 = create_exprs(a, b, j);
            auto es2 = create_exprs(x, b, j);
            auto es3 = create_exprs(a, y, j);
            for (unsigned i = 0; i < es1.size(); ++i) {                
                auto e1 = es1.get(i);
                auto e2 = es2.get(i);
                auto e3 = es3.get(i);
                check_repair_idx(e1, e2, 0, x);
                if (is_app(e1) && to_app(e1)->get_num_args() == 2)
                    check_repair_idx(e1, e3, 1, y);
            }
        }

        random_gen rand;

        void check_repair_idx(expr* e1, expr* e2, unsigned idx, expr* x) {            
            std::function<bool(expr*, unsigned)> value = [&](expr*, unsigned) {
                return rand() % 2 == 0;
                };
            expr_ref_vector es(m);
            bv_util bv(m);
            es.push_back(e1);
            es.push_back(e2);
            sls_eval ev(m);
            ev.init_eval(es, value);
            th_rewriter rw(m);
            expr_ref r(e1, m);
            if (m.is_bool(e1)) {
                auto val = m.is_true(r);
                auto val2 = ev.bval0(e2);
                if (val != val2) {
                    ev.set(e2, val);
                    auto rep1 = ev.try_repair(to_app(e2), idx);
                    if (!rep1) {
                        verbose_stream() << "Not repaired " << mk_pp(e2, m) << "\n";
                    }
                    SASSERT(rep1);
                }
            }
            if (bv.is_bv(e1)) {
                auto& val1 = ev.wval0(e1);
                auto& val2 = ev.wval0(e2);
                if (!val1.eq(val2)) {
                    val2.set(val1.bits);
                    auto rep2 = ev.try_repair(to_app(e2), idx);
                    if (!rep2) {
                        verbose_stream() << "Not repaired " << mk_pp(e2, m) << "\n";
                    }
                    SASSERT(rep2);
                }
            }
        }

        // todo: 
        void test_fixed() {

        }
    };
}


static void test_eval1() {
    ast_manager m;
    reg_decl_plugins(m);
    bv_util bv(m);

    expr_ref e(m);

    bv::sls_test validator(m);

    unsigned k = 0;
    unsigned bw = 6;
    for (unsigned i = 0; i < 1ul << bw; ++i) {
        expr_ref a(bv.mk_numeral(rational(i), bw), m);
        for (unsigned j = 0; j < 1ul << bw; ++j) {
            expr_ref b(bv.mk_numeral(rational(j), bw), m);
            ++k;
            if (k % 1000 == 0)
                verbose_stream() << "tests " << k << "\n";
            validator.check_eval(a, b, j);
        }
    }
}

static void test_repair1() {
    ast_manager m;
    reg_decl_plugins(m);
    bv_util bv(m);
    expr_ref e(m);
    bv::sls_test validator(m);

    unsigned k = 0;
    unsigned bw = 6;
    for (unsigned i = 0; i < 1ul << bw; ++i) {
        expr_ref a(bv.mk_numeral(rational(i), bw), m);
        for (unsigned j = 0; j < 1ul << bw; ++j) {
            expr_ref b(bv.mk_numeral(rational(j), bw), m);
            ++k;
            if (k % 1000 == 0)
                verbose_stream() << "tests " << k << "\n";
            validator.check_repair(a, b, j);
        }
    }
}

void tst_sls_test() {
    test_eval1();
}