/*++
Copyright (c) 2013 Microsoft Corporation

Module Name:

    pb_decl_plugin.cpp

Abstract:

    Cardinality Constraints plugin

Author:

    Nikolaj Bjorner (nbjorner) 2013-05-11

Revision History:

--*/

#include "ast/pb_decl_plugin.h"
#include "ast/ast_util.h"
#include "ast/ast_pp.h"

pb_decl_plugin::pb_decl_plugin():
    m_at_most_sym("at-most"),
    m_at_least_sym("at-least"),
    m_pble_sym("pble"),
    m_pbge_sym("pbge"),
    m_pbeq_sym("pbeq")
{}

func_decl * pb_decl_plugin::mk_func_decl(decl_kind k, unsigned num_parameters, parameter const * parameters, 
                                         unsigned arity, sort * const * domain, sort * range) {
    SASSERT(m_manager);
    ast_manager& m = *m_manager;
    for (unsigned i = 0; i < arity; ++i) 
        if (!m.is_bool(domain[i])) 
            m.raise_exception("invalid non-Boolean sort applied to Pseudo-Boolean relation");
    symbol sym;
    switch(k) {
    case OP_AT_LEAST_K: sym = m_at_least_sym; break;
    case OP_AT_MOST_K: sym = m_at_most_sym; break;
    case OP_PB_LE: sym = m_pble_sym; break;
    case OP_PB_GE: sym = m_pbge_sym; break;
    case OP_PB_EQ: sym = m_pbeq_sym; break;
    default: break;
    }
    switch(k) {
    case OP_AT_LEAST_K:
    case OP_AT_MOST_K: {
        if (num_parameters != 1 || !parameters[0].is_int() || parameters[0].get_int() < 0) 
            m.raise_exception("function expects one non-negative integer parameter");
        func_decl_info info(m_family_id, k, 1, parameters);
        return m.mk_func_decl(sym, arity, domain, m.mk_bool_sort(), info);
    }      
    case OP_PB_GE:
    case OP_PB_LE: 
    case OP_PB_EQ: {
        if (num_parameters != 1 + arity) {
            m.raise_exception("function expects arity+1 rational parameters");
        }
        vector<parameter> params;
        for (unsigned i = 0; i < num_parameters; ++i) {
            parameter const& p = parameters[i];
            if (p.is_int()) {
                params.push_back(p);
            }
            else if (p.is_rational()) {
                // HACK: ast pretty printer does not work with rationals.
                rational const& r = p.get_rational();
                if (r.is_int32()) {
                    params.push_back(parameter(r.get_int32()));
                }
                else {
                    params.push_back(p);
                }
            }
            else {
                m.raise_exception("functions 'pble/pbge/pbeq' expect arity+1 integer parameters");
            }
        }
        func_decl_info info(m_family_id, k, num_parameters, params.data());
        return m.mk_func_decl(sym, arity, domain, m.mk_bool_sort(), info);
    }
    default:
        UNREACHABLE();
        return nullptr;
    }
}

void pb_decl_plugin::get_op_names(svector<builtin_name> & op_names, symbol const & logic) {
    if (logic == symbol::null || logic == "QF_FD" || logic == "ALL" || logic == "HORN") {
        op_names.push_back(builtin_name(m_at_most_sym.str(), OP_AT_MOST_K));
        op_names.push_back(builtin_name(m_at_least_sym.str(), OP_AT_LEAST_K));
        op_names.push_back(builtin_name(m_pble_sym.str(), OP_PB_LE));
        op_names.push_back(builtin_name(m_pbge_sym.str(), OP_PB_GE));
        op_names.push_back(builtin_name(m_pbeq_sym.str(), OP_PB_EQ));
    }
}

void pb_util::normalize(unsigned num_args, rational const* coeffs, rational const& k) {
    m_coeffs.reset();
    bool all_ones = true;
    for (unsigned i = 0; i < num_args && all_ones; ++i) {
        all_ones = denominator(coeffs[i]).is_one();
    }
    if (all_ones) {
        for (unsigned i = 0; i < num_args; ++i) {
            m_coeffs.push_back(coeffs[i]);
        }
        m_k = k;        
    }
    else {
        rational d(1);
        for (unsigned i = 0; i < num_args; ++i) {
            d = lcm(d, denominator(coeffs[i]));
        }
        for (unsigned i = 0; i < num_args; ++i) {
            m_coeffs.push_back(d*coeffs[i]);
        }
        m_k = d*k;
    }
}

app * pb_util::mk_le(unsigned num_args, rational const * coeffs, expr * const * args, rational const& k) {
    normalize(num_args, coeffs, k);
    m_params.reset();
    m_params.push_back(parameter(floor(m_k)));
    bool all_ones = true;
    for (unsigned i = 0; i < num_args; ++i) {
        all_ones &= m_coeffs[i].is_one();
        m_params.push_back(parameter(m_coeffs[i]));
    }
    if (all_ones && k.is_unsigned() && floor(m_k).is_int32()) {
        m_params[0] = parameter(floor(m_k).get_int32());
        return m.mk_app(m_fid, OP_AT_MOST_K, 1, m_params.data(), num_args, args, m.mk_bool_sort());
    }
    return m.mk_app(m_fid, OP_PB_LE, m_params.size(), m_params.data(), num_args, args, m.mk_bool_sort());
}

app * pb_util::mk_ge(unsigned num_args, rational const * coeffs, expr * const * args, rational const& k) {
    normalize(num_args, coeffs, k);
    m_params.reset();
    m_params.push_back(parameter(ceil(m_k)));
    bool all_ones = true;
    for (unsigned i = 0; i < num_args; ++i) {
        all_ones &= m_coeffs[i].is_one();
        m_params.push_back(parameter(m_coeffs[i]));
    }
    if (all_ones && k.is_unsigned()) {
        m_params[0] = parameter(ceil(m_k).get_unsigned());
        return m.mk_app(m_fid, OP_AT_LEAST_K, 1, m_params.data(), num_args, args, m.mk_bool_sort());
    }
    return m.mk_app(m_fid, OP_PB_GE, m_params.size(), m_params.data(), num_args, args, m.mk_bool_sort());
}

app * pb_util::mk_eq(unsigned num_args, rational const * coeffs, expr * const * args, rational const& k) {
    normalize(num_args, coeffs, k);
    if (!m_k.is_int()) {
        return m.mk_false();
    }
    if (num_args == 0) {
        return m_k.is_zero() ? m.mk_true() : m.mk_false();
    }
    m_params.reset();
    m_params.push_back(parameter(m_k));
    for (unsigned i = 0; i < num_args; ++i) {
        m_params.push_back(parameter(m_coeffs[i]));
    }
    return m.mk_app(m_fid, OP_PB_EQ, m_params.size(), m_params.data(), num_args, args, m.mk_bool_sort());
}

// ax + by < k
// <=>
// -ax - by >= -k + 1
// <=>
// a(1-x) + b(1-y) >= -k + a + b + 1
app * pb_util::mk_lt(unsigned num_args, rational const * _coeffs, expr * const * _args, rational const& _k) {
    normalize(num_args, _coeffs, _k);
    expr_ref_vector args(m);
    for (unsigned i = 0; i < num_args; ++i) {
        args.push_back(mk_not(m, _args[i]));
    }
    m_k = floor(m_k);
    m_k.neg();
    m_k += rational::one();
    for (unsigned i = 0; i < num_args; ++i) {
        m_k += m_coeffs[i];
    }
    return mk_ge(num_args, m_coeffs.data(), args.data(), m_k);
}


app * pb_util::mk_at_most_k(unsigned num_args, expr * const * args, unsigned k) {
    parameter param(k);
    return m.mk_app(m_fid, OP_AT_MOST_K, 1, &param, num_args, args, m.mk_bool_sort());
}

bool pb_util::is_at_most_k(func_decl *a) const {
    return is_decl_of(a, m_fid, OP_AT_MOST_K);
}

bool pb_util::is_at_most_k(expr *a, rational& k) const {
    if (is_at_most_k(a)) {
        k = get_k(a);
        return true;
    }
    else {
        return false;
    }
}

app * pb_util::mk_at_least_k(unsigned num_args, expr * const * args, unsigned k) {
    parameter param(k);
    return m.mk_app(m_fid, OP_AT_LEAST_K, 1, &param, num_args, args, m.mk_bool_sort());
}

bool pb_util::is_at_least_k(func_decl *a) const {
    return is_decl_of(a, m_fid, OP_AT_LEAST_K);
}

bool pb_util::is_at_least_k(expr *a, rational& k) const {
    if (is_at_least_k(a)) {
        k = get_k(a);
        return true;
    }
    else {
        return false;
    }
}

rational pb_util::get_k(func_decl *a) const {
    parameter const& p = a->get_parameter(0);
    if (is_at_most_k(a) || is_at_least_k(a)) {
        return to_rational(p);
    }
    else {
        SASSERT(is_le(a) || is_ge(a) || is_eq(a));
        return to_rational(p);
    }
}


bool pb_util::is_le(func_decl *a) const {
    return is_decl_of(a, m_fid, OP_PB_LE);
}

bool pb_util::is_le(expr* a, rational& k) const {
    if (is_le(a)) {
        k = get_k(a);
        return true;
    }
    else {
        return false;
    }
}

bool pb_util::is_ge(func_decl *a) const {
    return is_decl_of(a, m_fid, OP_PB_GE);
}

bool pb_util::is_ge(expr* a, rational& k) const {
    if (is_ge(a)) {
        k = get_k(a);
        return true;
    }
    else {
        return false;
    }
}


bool pb_util::is_eq(func_decl *a) const {
    return is_decl_of(a, m_fid, OP_PB_EQ);
}

bool pb_util::is_eq(expr* a, rational& k) const {
    if (is_eq(a)) {
        k = get_k(a);
        return true;
    }
    else {
        return false;
    }
}

rational pb_util::get_coeff(func_decl* a, unsigned index) const {
    if (is_at_most_k(a) || is_at_least_k(a)) {
        return rational::one();
    }
    SASSERT(is_le(a) || is_ge(a) || is_eq(a));
    SASSERT(1 + index < a->get_num_parameters());
    return to_rational(a->get_parameter(index + 1));
}

rational pb_util::to_rational(parameter const& p) const {
    if (p.is_int()) {
        return rational(p.get_int());
    }
    SASSERT(p.is_rational());
    return p.get_rational();
}

bool pb_util::has_unit_coefficients(func_decl* f) const {
    if (is_at_most_k(f) || is_at_least_k(f)) return true;
    unsigned sz = f->get_arity();
    for (unsigned i = 0; i < sz; ++i) {
        if (!get_coeff(f, i).is_one()) return false;
    }
    return true;
}

app* pb_util::mk_fresh_bool() {
    symbol name = m.mk_fresh_var_name("pb");
    func_decl_info info(m_fid, OP_PB_AUX_BOOL, 0, nullptr);
    return m.mk_const(m.mk_func_decl(name, 0, (sort *const*)nullptr, m.mk_bool_sort(), info));
}