/*++
Copyright (c) 2011 Microsoft Corporation

Module Name:

    expr2polynomial.cpp

Abstract:

    Translator from Z3 expressions into multivariate polynomials (and back).
    
    
Author:

    Leonardo (leonardo) 2011-12-23

Notes:

--*/
#include"expr2polynomial.h"
#include"expr2var.h"
#include"arith_decl_plugin.h"
#include"ast_smt2_pp.h"
#include"z3_exception.h"
#include"cooperate.h"

struct expr2polynomial::imp {
    struct frame {
        app *     m_curr;
        unsigned  m_idx;
        frame():m_curr(0), m_idx(0) {}
        frame(app * t):m_curr(t), m_idx(0) {}
    };
    
    expr2polynomial &                  m_wrapper;
    ast_manager &                      m_am;
    arith_util                         m_autil;
    polynomial::manager &              m_pm;
    expr2var *                         m_expr2var;
    bool                               m_expr2var_owner;
    expr_ref_vector                    m_var2expr;
    
    obj_map<expr, unsigned>            m_cache;
    expr_ref_vector                    m_cached_domain;
    polynomial::polynomial_ref_vector  m_cached_polynomials;
    polynomial::scoped_numeral_vector  m_cached_denominators;

    svector<frame>                     m_frame_stack;
    polynomial::polynomial_ref_vector  m_presult_stack;
    polynomial::scoped_numeral_vector  m_dresult_stack;

    volatile bool                      m_cancel;
    
    imp(expr2polynomial & w, ast_manager & am, polynomial::manager & pm, expr2var * e2v):
        m_wrapper(w),
        m_am(am),
        m_autil(am),
        m_pm(pm),
        m_expr2var(e2v == 0 ? alloc(expr2var, am) : e2v),
        m_expr2var_owner(e2v == 0),
        m_var2expr(am),
        m_cached_domain(am),
        m_cached_polynomials(pm),
        m_cached_denominators(pm.m()),
        m_presult_stack(pm),
        m_dresult_stack(pm.m()),
        m_cancel(false) {
    }

    ~imp() {
        if (m_expr2var_owner)
            dealloc(m_expr2var);
    }

    ast_manager & m() { return m_am; }
    polynomial::manager & pm() { return m_pm; }
    polynomial::numeral_manager & nm() { return pm().m(); }

    void reset() {
        m_frame_stack.reset();
        m_presult_stack.reset();
        m_dresult_stack.reset();
    }

    void reset_cache() {
        m_cache.reset();
        m_cached_domain.reset();
        m_cached_polynomials.reset();
        m_cached_denominators.reset();
    }

    void checkpoint() {
        if (m_cancel)
            throw default_exception("canceled");
        cooperate("expr2polynomial");
    }

    void push_frame(app * t) {
        m_frame_stack.push_back(frame(t));
    }

    void cache_result(expr * t) {
        SASSERT(!m_cache.contains(t));
        SASSERT(m_cached_denominators.size() == m_cached_polynomials.size());
        SASSERT(m_cached_denominators.size() == m_cached_domain.size());
        if (t->get_ref_count() <= 1) 
            return;
        unsigned idx = m_cached_polynomials.size();
        m_cache.insert(t, idx);
        m_cached_domain.push_back(t);
        m_cached_polynomials.push_back(m_presult_stack.back());
        m_cached_denominators.push_back(m_dresult_stack.back());
    }

    bool is_cached(expr * t) {
        return t->get_ref_count() > 1 && m_cache.contains(t);
    }

    bool is_int_real(expr * t) {
        return m_autil.is_int_real(t);
    }

    void store_result(expr * t, polynomial::polynomial * p, polynomial::numeral & d) {
        m_presult_stack.push_back(p);
        m_dresult_stack.push_back(d);
        cache_result(t);
    }

    void store_var_poly(expr * t) {
        polynomial::var x = m_expr2var->to_var(t);
        if (x == UINT_MAX) {
            bool is_int = m_autil.is_int(t);
            x = m_wrapper.mk_var(is_int);
            m_expr2var->insert(t, x);
            if (x >= m_var2expr.size())
                m_var2expr.resize(x+1, 0);
            m_var2expr.set(x, t);
        }
        polynomial::numeral one(1);
        store_result(t, pm().mk_polynomial(x), one);
    }

    void store_const_poly(app * n) {
        rational val;
        VERIFY(m_autil.is_numeral(n, val));
        polynomial::scoped_numeral d(nm());
        d = val.to_mpq().denominator();
        store_result(n, pm().mk_const(numerator(val)), d);
    }

    bool visit_arith_app(app * t) {
        switch (t->get_decl_kind()) {
        case OP_NUM:
            store_const_poly(t);
            return true;
        case OP_ADD: case OP_SUB: case OP_MUL: case OP_UMINUS: case OP_TO_REAL: 
            push_frame(t);
            return false;
        case OP_POWER: {
            rational k;
            SASSERT(t->get_num_args() == 2);
            if (!m_autil.is_numeral(t->get_arg(1), k) || !k.is_int() || !k.is_unsigned()) {
                store_var_poly(t);
                return true;
            }
            push_frame(t);
            return false;
        }
        default:
            // can't handle operator
            store_var_poly(t);
            return true;
        }
    }

    bool visit(expr * t) {
        SASSERT(is_int_real(t));
        if (is_cached(t)) {
            unsigned idx = m_cache.find(t);
            m_presult_stack.push_back(m_cached_polynomials.get(idx));
            m_dresult_stack.push_back(m_cached_denominators.get(idx));
            return true;
        }

        SASSERT(!is_quantifier(t));
        if (::is_var(t)) {
            store_var_poly(t);
            return true;
        }
        
        SASSERT(is_app(t));
        if (!m_autil.is_arith_expr(t)) {
            store_var_poly(t);
            return true;
        }
        
        return visit_arith_app(to_app(t));
    }

    void pop(unsigned num_args) {
        SASSERT(m_presult_stack.size() == m_dresult_stack.size());
        SASSERT(m_presult_stack.size() >= num_args);
        m_presult_stack.shrink(m_presult_stack.size() - num_args);
        m_dresult_stack.shrink(m_dresult_stack.size() - num_args);
    }

    polynomial::polynomial * const * polynomial_args(unsigned num_args) {
        SASSERT(m_presult_stack.size() >= num_args);
        return m_presult_stack.c_ptr() + m_presult_stack.size() - num_args;
    } 

    polynomial::numeral const * denominator_args(unsigned num_args) {
        SASSERT(m_dresult_stack.size() >= num_args);
        return m_dresult_stack.c_ptr() + m_dresult_stack.size() - num_args;
    }

    template<bool is_add>
    void process_add_sub(app * t) {
        SASSERT(t->get_num_args() <= m_presult_stack.size());
        unsigned num_args = t->get_num_args();
        polynomial::polynomial * const * p_args = polynomial_args(num_args);
        polynomial::numeral const * d_args = denominator_args(num_args);
        polynomial::polynomial_ref p(pm());
        polynomial::polynomial_ref p_aux(pm());
        polynomial::scoped_numeral d(nm());
        polynomial::scoped_numeral d_aux(nm());
        d = 1;
        for (unsigned i = 0; i < num_args; i++) {
            nm().lcm(d, d_args[i], d);
        }
        p = pm().mk_zero();
        for (unsigned i = 0; i < num_args; i++) {
            checkpoint();
            nm().div(d, d_args[i], d_aux);
            p_aux = pm().mul(d_aux, p_args[i]);
            if (i == 0)
                p = p_aux;
            else if (is_add)
                p = pm().add(p, p_aux);
            else
                p = pm().sub(p, p_aux);
        }
        pop(num_args);
        store_result(t, p.get(), d.get());
    }

    void process_add(app * t) {
        process_add_sub<true>(t);
    }

    void process_sub(app * t) {
        process_add_sub<false>(t);
    }

    void process_mul(app * t) {
        SASSERT(t->get_num_args() <= m_presult_stack.size());
        unsigned num_args = t->get_num_args();
        polynomial::polynomial * const * p_args = polynomial_args(num_args);
        polynomial::numeral const * d_args = denominator_args(num_args);
        polynomial::polynomial_ref p(pm());
        polynomial::scoped_numeral d(nm());
        p = pm().mk_const(rational(1));
        d = 1;
        for (unsigned i = 0; i < num_args; i++) {
            checkpoint();
            p = pm().mul(p, p_args[i]);
            d = d * d_args[i];
        }
        pop(num_args);
        store_result(t, p.get(), d.get());
    }

    void process_uminus(app * t) {
        SASSERT(t->get_num_args() <= m_presult_stack.size());
        polynomial::polynomial_ref neg_p(pm());
        neg_p = pm().neg(m_presult_stack.back());
        m_presult_stack.pop_back();
        m_presult_stack.push_back(neg_p);
        cache_result(t);
    }

    void process_power(app * t) {
        SASSERT(t->get_num_args() <= m_presult_stack.size());
        rational _k;
        VERIFY(m_autil.is_numeral(t->get_arg(1), _k));
        SASSERT(_k.is_int() && _k.is_unsigned());
        unsigned k = _k.get_unsigned();
        polynomial::polynomial_ref p(pm());
        polynomial::scoped_numeral d(nm());
        unsigned num_args = t->get_num_args();
        polynomial::polynomial * const * p_args = polynomial_args(num_args);
        polynomial::numeral const * d_args = denominator_args(num_args);
        pm().pw(p_args[0], k, p);
        nm().power(d_args[0], k, d);
        pop(num_args);
        store_result(t, p.get(), d.get());
    }

    void process_to_real(app * t) {
        // do nothing
        cache_result(t);
    }
 
    void process_app(app * t) {
        SASSERT(m_presult_stack.size() == m_dresult_stack.size());
        
        switch (t->get_decl_kind()) {
        case OP_ADD: 
            process_add(t);
            return;
        case OP_SUB:
            process_sub(t);
            return;
        case OP_MUL: 
            process_mul(t);
            return;
        case OP_POWER:
            process_power(t);
            return;
        case OP_UMINUS: 
            process_uminus(t);
            return;
        case OP_TO_REAL: 
            process_to_real(t);
            return;
        default:
            UNREACHABLE();
        }
    }
    
    bool to_polynomial(expr * t, polynomial::polynomial_ref & p, polynomial::scoped_numeral & d) {
        if (!is_int_real(t))
            return false;
        reset();
        if (!visit(t)) {
            while (!m_frame_stack.empty()) {
            begin_loop:
                checkpoint();
                frame & fr = m_frame_stack.back();
                app * t = fr.m_curr;
                TRACE("expr2polynomial", tout << "processing: " << fr.m_idx << "\n" << mk_ismt2_pp(t, m()) << "\n";);
                unsigned num_args = t->get_num_args();
                while (fr.m_idx < num_args) {
                    expr * arg = t->get_arg(fr.m_idx);
                    fr.m_idx++;
                    if (!visit(arg)) 
                        goto begin_loop;
                }
                process_app(t);
                m_frame_stack.pop_back();
            }
        }
        p = m_presult_stack.back();
        d = m_dresult_stack.back();
        reset();
        return true;
    }

    bool is_int_poly(polynomial::polynomial_ref const & p) {
        unsigned sz = size(p);
        for (unsigned i = 0; i < sz; i++) {
            polynomial::monomial * m = pm().get_monomial(p, i);
            unsigned msz = pm().size(m);
            for (unsigned j = 0; j < msz; j++) {
                polynomial::var x = pm().get_var(m, j);
                if (!m_wrapper.is_int(x))
                    return false;
            }
        }
        return true;
    }

    void to_expr(polynomial::polynomial_ref const & p, bool use_power, expr_ref & r) {
        expr_ref_buffer args(m());
        expr_ref_buffer margs(m());
        unsigned sz = size(p);
        bool is_int = is_int_poly(p);

        for (unsigned i = 0; i < sz; i++) {
            margs.reset();
            polynomial::monomial * m = pm().get_monomial(p, i);
            polynomial::numeral const & a = pm().coeff(p, i);
            if (!nm().is_one(a)) {
                margs.push_back(m_autil.mk_numeral(rational(a), is_int));
            }
            unsigned msz = pm().size(m);
            for (unsigned j = 0; j < msz; j++) {
                polynomial::var x = pm().get_var(m, j);
                expr * t = m_var2expr.get(x);
                if (m_wrapper.is_int(x) && !is_int) {
                    t = m_autil.mk_to_real(t);
                }
                unsigned d = pm().degree(m, j);
                if (use_power && d > 1) {
                    margs.push_back(m_autil.mk_power(t, m_autil.mk_numeral(rational(d), is_int)));
                }
                else {
                    for (unsigned k = 0; k < d; k++)
                        margs.push_back(t);
                }
            }
            if (margs.size() == 0) {
                args.push_back(m_autil.mk_numeral(rational(1), is_int));
            }
            else if (margs.size() == 1) {
                args.push_back(margs[0]);
            }
            else {
                args.push_back(m_autil.mk_mul(margs.size(), margs.c_ptr()));
            }
        }

        if (args.size() == 0) {
            r = m_autil.mk_numeral(rational(0), is_int);
        }
        else if (args.size() == 1) {
            r = args[0];
        }
        else {
            r = m_autil.mk_add(args.size(), args.c_ptr());
        }
    }

    void set_cancel(bool f) {
        m_cancel = f;
    }
};

expr2polynomial::expr2polynomial(ast_manager & am, polynomial::manager & pm, expr2var * e2v) {
    m_imp = alloc(imp, *this, am, pm, e2v);
}

expr2polynomial::~expr2polynomial() {
    dealloc(m_imp);
}

ast_manager & expr2polynomial::m() const { 
    return m_imp->m_am; 
}

polynomial::manager & expr2polynomial::pm() const { 
    return m_imp->m_pm; 
}

bool expr2polynomial::to_polynomial(expr * t, polynomial::polynomial_ref & p, polynomial::scoped_numeral & d) {
    return m_imp->to_polynomial(t, p, d);
}

void expr2polynomial::to_expr(polynomial::polynomial_ref const & p, bool use_power, expr_ref & r) {
    m_imp->to_expr(p, use_power, r);
}

bool expr2polynomial::is_var(expr * t) const { 
    return m_imp->m_expr2var->is_var(t); 
}
    
expr2var const & expr2polynomial::get_mapping() const { 
    return *(m_imp->m_expr2var); 
}

void expr2polynomial::set_cancel(bool f) {
    m_imp->set_cancel(f);
}

default_expr2polynomial::default_expr2polynomial(ast_manager & am, polynomial::manager & pm):
    expr2polynomial(am, pm, 0) {
}

default_expr2polynomial::~default_expr2polynomial() {
}

bool default_expr2polynomial::is_int(polynomial::var x) const {
    return m_is_int[x];
}

polynomial::var default_expr2polynomial::mk_var(bool is_int) { 
    polynomial::var x = pm().mk_var(); 
    m_is_int.reserve(x+1, false);
    m_is_int[x] = is_int;
    return x;
}