// -*- Mode: c++; c-basic-offset: 4 -*-

#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <string>
#include <tuple>
#include <map>
#include <unordered_map>
#include <boost/format.hpp>
#include <boost/program_options.hpp>

using namespace std;
using boost::format;

const bool debug = false;

typedef long long Cost;

struct wcsptuple {
    vector<size_t> tup;
    Cost cost;

    wcsptuple(vector<size_t> const &t, Cost c) : tup(t), cost(c) {}
};

struct wcspfunc {
    Cost defcost;
    vector<size_t> scope;
    vector< wcsptuple > specs;

    size_t arity() const { return scope.size(); }
};

struct wcsp {
    Cost ub;

    vector<size_t> domains;
    size_t nvars() const { return domains.size(); }

    vector<wcspfunc> functions;
};

template<typename T>
vector<T> read_vec(istream& is)
{
    vector<T> r;
    T s;
    is >> s;
    while(is) {
        r.push_back(s);
        is >> s;
    }
    return r;
}

template<typename T>
vector<T> read_vec(string const& line)
{
    istringstream iss(line);
    return read_vec<T>(iss);
}

tuple<string, size_t, size_t, size_t, Cost> read_header(string const& line)
{
    istringstream iss(line);

    string name;
    size_t nvars;
    size_t domsize;
    size_t nfun;
    Cost ub;

    iss >> name >> nvars >> domsize >> nfun >> ub;
    return make_tuple(name, nvars, domsize, nfun, ub);
}

wcspfunc read_fun(istream& is)
{
    string line;
    getline(is, line);
    vector<Cost> hd = read_vec<Cost>(line);
    size_t arity = hd[0];
    Cost defcost = hd[hd.size()-2];
    size_t nspec = hd[hd.size()-1];

    vector<wcsptuple> specs;
    for(size_t i = 0; i != nspec; ++i) {
        getline(is, line);
        vector<Cost> v = read_vec<Cost>(line);
        specs.push_back( {vector<size_t>(v.begin(), v.begin()+arity), v[v.size()-1]} );
    }

    return { defcost, vector<size_t>(hd.begin()+1, hd.begin()+1+arity),
            specs };
}

wcsp readwcsp(istream& is)
{
    wcsp w;

    string name;
    size_t nvars;
    size_t domsize;
    size_t nfun;

    string line;

    getline(is, line);
    tie(name, nvars, domsize, nfun, w.ub) = read_header(line);

    getline(is, line);
    w.domains = read_vec<size_t>(line);

    for(size_t i = 0; i != nfun; ++i)
        w.functions.push_back(read_fun(is));

    return w;
}

struct wclause {
    vector<int> c;
    Cost w;
};

typedef unordered_map<string, int> varmap;
typedef unordered_map<int, string> rvarmap;

struct wcnf {
    int numvars;
    mutable int numclauses;
    Cost ub;
    varmap v;
    varmap curfuncmap;
    rvarmap rv;

    wcnf() : numvars(0), numclauses(0), ub(0) {}
};

void write_clause(ostream& os, wcnf const& c, wclause const& cl)
{
    ++c.numclauses;
    os << cl.w << ' ';
    for(int l : cl.c)
        if( debug ) {
            if( l < 0 )
                os << '-';
            string const& n = c.rv.find(abs(l))->second;
            os << n << ' ';
        }
        else
            os << l << ' ';
    os << "0\n";
}

int var(wcnf& w, string const& name, varmap& vmap)
{
    auto i = vmap.find(name);
    if( i != vmap.end() )
        return i->second;
    ++w.numvars;
    vmap[name] = w.numvars;
    if(debug) w.rv[w.numvars] = name;
    return w.numvars;
}

namespace std {
    template<typename T>
    ostream& operator<<(ostream& os, vector<T> const& v)
    {
        os << '[';
        bool comma = false;
        for(auto& e: v) {
            if(comma)
                os << ',';
            comma = true;
            os << e;
        }
        os << ']';
        return os;
    }
}

int dvar(wcnf& w, int v, int d)
{
    return var(w, (format("x%s_%s") % v % d).str(), w.v);
}

int tvar(wcnf& w, wcspfunc const& f, wcsptuple const& t)
{
    return var(w, (format("t%s_%s") % f.scope % t.tup).str(), w.curfuncmap);
}

// write direct encoding of a single cost function. Also used by the
// tuple encoding for unary cost functions.
template<typename P>
void foreach_tuple(wcsp const &w, wcspfunc const &f,
                   wcnf & c, vector<size_t>& current,
                   P proc)
{
    if( current.size() == f.arity() ) {
        proc(current);
        return;
    }
    int idx = current.size();
    current.push_back(0);
    int var = f.scope[idx];
    for(size_t j = 0; j != w.domains[var]; ++j) {
        current[idx] = j;
        foreach_tuple(w, f, c, current, proc);
    }
    current.pop_back();
}

int write_cf_clauses(ostream& os,
                      wcsp const &w, wcspfunc const &f,
                      map< vector<size_t>, Cost > & fexplicit,
                      wcnf & c, vector<size_t>& current)
{
    int nc = 0;
    foreach_tuple(w, f, c, current, [&](vector<size_t>& current) {
            auto i = fexplicit.find(current);
            Cost tcost;
            if( i != fexplicit.end() )
                tcost = min(i->second, w.ub);
            else
                tcost = min(f.defcost, w.ub);
            if( tcost == 0 )
                return;
            wclause cl{ {}, tcost };
            for(size_t j = 0; j != current.size(); ++j)
                cl.c.push_back(-dvar(c, f.scope[j], current[j]));
            write_clause(os, c, cl);
            ++nc;
            return;
        });
    return nc;
}

int realnumtups(wcsp const& w, wcspfunc const& f)
{
    int inf = count_if(f.specs.begin(), f.specs.end(),
                       [&w](wcsptuple const& t) -> bool {
                           return t.cost >= w.ub;
                       });
    if( f.defcost >= w.ub )
        return f.specs.size() - inf;
    int numtups = accumulate(f.scope.begin(), f.scope.end(), 1,
                             [&w](int x, size_t var) -> int {
                                 return x*w.domains[var];
                             });
    return numtups - inf;
}

int nonzerotups(wcsp const& w, wcspfunc const& f)
{
    if( f.defcost == 0 ) {
        int nonzero = count_if(f.specs.begin(), f.specs.end(),
                               [&w](wcsptuple const& t) -> bool {
                                   return t.cost > 0 && t.cost < w.ub;
                               });
        return nonzero;
    } else {
        int zero = count_if(f.specs.begin(), f.specs.end(),
                               [](wcsptuple const& t) -> bool {
                                   return t.cost == 0;
                               });
        return realnumtups(w, f) - zero;
    }
}

void write_domains(ostream& os, wcsp const& w, wcnf& c)
{
    for(size_t i = 0; i != w.nvars(); ++i) {
        if( w.domains[i] <= 2 ) {
            // we treat unary domain as a special case of binary
            int bvar = dvar(c, i, 1);
            c.v[(format("x%s_%s") % i % 0).str()] = -bvar;
            if( w.domains[i] == 1 )
                write_clause( os, c, {{ dvar(c, i, 0) }, w.ub });
        } else {
            wclause pos{ {}, w.ub };
            for(size_t j1 = 0; j1 != w.domains[i]; ++j1) {
                pos.c.push_back(dvar(c, i, j1));
                for(size_t j2 = j1+1; j2 != w.domains[i]; ++j2) {
                    write_clause(os, c, {{-dvar(c, i, j1), -dvar(c, i, j2)}, w.ub});
                }
            }
            write_clause(os, c, pos);
        }
    }
}

void convert2sat_tuple(wcsp const& w, ostream& os)
{
    wcnf c;
    c.ub = w.ub;

    int nv = accumulate(w.domains.begin(), w.domains.end(), 0,
                        [&](int x, int d) {
                            if( d > 2 )
                                return x + d;
                            else
                                return x + 1;
                        });
    nv = accumulate(w.functions.begin(), w.functions.end(), nv,
                    [&](int x, wcspfunc const& f) -> int {
                        if( f.arity() > 1 )
                            return x + realnumtups(w, f);
                        else
                            return x;
                    });

    // number of clauses to encode domains ...
    int nc = accumulate(w.domains.begin(), w.domains.end(), 0,
                        [](int x, int d) -> int {
                            if( d > 2 )
                                return x + d*(d-1)/2 + 1;
                            else if ( d == 1 )
                                return x + 1;
                            else
                                return x;
                        });
    // ... and number to encode cost functions
    for(auto& f : w.functions) {
        int nonzero = nonzerotups( w, f );
        if( f.arity() == 0 ) {
            nc += 1;
            continue;
        }
        if( f.arity() == 1 ) {
            int clnum = 0;
            int zero = count_if(f.specs.begin(), f.specs.end(),
                                [&w](wcsptuple const& t) -> bool {
                                    return t.cost == 0;
                                });
            if( f.defcost == 0 ) {
                clnum = f.specs.size() - zero;
            } else {
                clnum = w.domains[f.scope[0]] - zero;
            }
            nc += clnum;
            continue;
        }
        int numvals = accumulate(f.scope.begin(), f.scope.end(), 0,
                                 [&w](int x, size_t var) -> int {
                                     return x+w.domains[var];
                                 });
        int numtups = realnumtups(w, f);
        nc += numtups * f.arity() + numvals + nonzero;
    }

    os << "p wcnf " << nv << ' ' << nc << ' ' << w.ub << "\n";

    write_domains(os, w, c);

    // convert functions
    for(auto& f : w.functions) {
        if( f.arity() == 0 ) {
            write_clause(os, c, { {}, f.defcost });
            ++nc;
        } else if( f.arity() == 1 ) {
            map< vector<size_t>, Cost > tups;
            for(auto& tuple : f.specs )
                tups[tuple.tup] = tuple.cost;
            vector<size_t> current;
            write_cf_clauses(os, w, f, tups, c, current);
            assert(current.empty());
        } else {
            map< vector<size_t>, Cost > tups;
            for(auto& tuple : f.specs )
                tups[tuple.tup] = tuple.cost;
            vector<size_t> current;
            map< pair<int, int>, vector<shared_ptr<wcsptuple>> > supports;
            foreach_tuple(w, f, c, current, [&](vector<size_t>& current) {
                    Cost tcost;
                    auto i = tups.find(current);
                    if( i != tups.end() )
                        tcost = i->second;
                    else
                        tcost = f.defcost;
                    auto tuple = make_shared<wcsptuple>(current, tcost);
                    if( tcost >= w.ub )
                        return;
                    if( tcost != 0 ) {
                        write_clause(os, c, {{-tvar(c, f, *tuple)}, tcost});
                    }
                    for(size_t i = 0; i != f.arity(); ++i) {
                        write_clause(os, c, {
                                {-tvar(c, f, *tuple),
                                        dvar(c, f.scope[i], tuple->tup[i])},
                                    c.ub });
                        supports[make_pair(f.scope[i], tuple->tup[i])].push_back(tuple);
                    }
                });
            for(auto& var : f.scope) {
                for(size_t i = 0; i != w.domains[var]; ++i) {
                    wclause cl{ {-dvar(c, var, i)}, c.ub };
                    for(auto& sup : supports[make_pair(var, i)]) {
                        cl.c.push_back(tvar(c, f, *sup));
                    }
                    write_clause(os, c, cl);
                }
            }
        }
        c.curfuncmap.clear();
    }
}

void convert2sat_direct(wcsp const& w, ostream& os)
{
    wcnf c;
    c.ub = w.ub;

    int nv = accumulate(w.domains.begin(), w.domains.end(), 0,
                        [&](int x, int d) {
                            if( d > 2 )
                                return x + d;
                            else
                                return x + 1;
                        });

    int nc = accumulate(w.domains.begin(), w.domains.end(), 0,
                        [](int x, int d) -> int {
                            if( d > 2 )
                                return x + d*(d-1)/2 + 1;
                            else if ( d == 1 )
                                return x + 1;
                            else
                                return x;
                        });
    for(auto& f : w.functions) {
        if( f.arity() == 0 ) {
            nc += 1;
            continue;
        }
        int clnum = 0;
        int zero = count_if(f.specs.begin(), f.specs.end(),
                            [&w](wcsptuple const& t) -> bool {
                                return t.cost == 0;
                            });
        if( f.defcost == 0 ) {
            clnum = f.specs.size() - zero;
        } else {
            int numtups = accumulate(f.scope.begin(), f.scope.end(), 1,
                                     [&w](int x, size_t var) -> int {
                                         return x*w.domains[var];
                                     });
            clnum = numtups - zero;
        }
        nc += clnum;
    }

    os << "p wcnf " << nv << ' ' << nc << ' ' << w.ub << "\n";

    // convert all domains
    write_domains(os, w, c);

    vector<size_t> current;
    // convert functions
    for(auto& f : w.functions) {
        if( f.arity() == 0 ) {
            write_clause(os, c, { {}, f.defcost });
            continue;
        }
        map< vector<size_t>, Cost > tups;
        for(auto& tuple : f.specs )
            tups[tuple.tup] = tuple.cost;
        current.clear();
        write_cf_clauses(os, w, f, tups, c, current);
        assert(current.empty());
    }
}

void remove_nullary_functions(wcsp& w)
{
    // sum the costs of all the nullary CFs and add that cost to all
    // elements of the first non-nullary CF.
    Cost c0cost = accumulate(w.functions.begin(), w.functions.end(), 0,
                            [](int x, wcspfunc& f) -> Cost {
                                if( f.arity() == 0 )
                                    return x+f.defcost;
                                else
                                    return x;
                            });
    w.functions.erase(remove_if(w.functions.begin(), w.functions.end(),
                                [](wcspfunc& f) -> bool {
                                    return f.arity() == 0;
                                }), w.functions.end());
    wcspfunc f{ c0cost, {0}, {}};
    w.functions.push_back(f);
}

int main(int argc, char* argv[])
{
    namespace po = boost::program_options;

    string encoding;
    po::options_description desc("Allowed options");
    desc.add_options()
        ("help", "produce help message")
        ("encoding", po::value<string>(&encoding)->default_value("tuple"),
                "use direct/tuple encoding")
        ("input-file,i", po::value<string>(), "wcsp input file")
        ("output-file,o", po::value<string>(), "wcsp output file")
        ;

    po::positional_options_description p;
    p.add("input-file", 1).add("output-file", 1);

    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).
              options(desc).positional(p).run(), vm);
    po::notify(vm);

    if (vm.count("help")) {
        cout << desc << "\n";
        return 1;
    }

    bool tenc = true;
    if (encoding == "direct" )
        tenc = false;

    if(!vm.count("input-file")) {
        cout << "must specify input file\n";
        return 1;
    }

    ifstream ifs(vm["input-file"].as<string>());
    if( !ifs ) {
        cout << "could not open " << argv[1] << "\n";
        return 1;
    }

    ofstream rofs;
    ostream *ofs;
    if(vm.count("output-file")) {
        rofs.open(vm["output-file"].as<string>());
        if( !rofs ) {
            cout << "could not open " << argv[2] << "\n";
            return 1;
        }
        ofs = &rofs;
    } else
        ofs = &cout;

    wcsp w = readwcsp(ifs);

    Cost newtop = 0;
    for(auto const& f : w.functions) {
        if( f.specs.empty() ) { // empty function
            newtop += f.defcost;
            continue;
        }
        auto me = std::max_element(f.specs.begin(), f.specs.end(),
                                   [&](wcsptuple const& m, wcsptuple const& c){
                                       return c.cost < w.ub &&
                                       (c.cost > m.cost ||
                                        m.cost >= w.ub);
                                   });
        newtop += me->cost;
    }
    if( newtop < w.ub )
        w.ub = newtop;

    // special handling for nullary cost functions because maxsat
    // solvers, including at least MaxHS and minimaxsat, are not happy
    // with empty clauses
    remove_nullary_functions(w);

    if( tenc )
        convert2sat_tuple(w, *ofs);
    else
        convert2sat_direct(w, *ofs);

    return 0;
}