//
// Created by kh on 3/14/20.
//

#include "Algo.h"

#include <igen/Context.h>
#include <igen/Domain.h>
#include <igen/Config.h>
#include <igen/runner/ProgramRunnerMt.h>
#include <igen/CoverageStore.h>
#include <igen/c50/CTree.h>

#include <klib/print_stl.h>
#include <klib/vecutils.h>
#include <klib/random.h>

#include <fstream>

#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string_regex.hpp>
#include <boost/container/flat_map.hpp>
#include <boost/container/flat_set.hpp>
#include <boost/lexical_cast.hpp>

#include <tsl/robin_set.h>
#include <tsl/robin_map.h>

namespace igen {

class Analyzer : public Object {
public:
    explicit Analyzer(PMutContext ctx) : Object(move(ctx)) {}

    struct LocData {
        expr e;
        PMutCTree tree;
        bool ignored;
        bool is_first;
        unsigned id_;
        vec<str> comments;

        [[nodiscard]] unsigned id() const { return id_; }
    };

    expr to_expr(const z3::model &m, double &ncex) {
        z3::expr_vector vecExpr(m.ctx());
        for (int id = 0; id < dom()->n_vars(); ++id) {
            expr v = m.get_const_interp(dom()->var(id)->zvar().decl());
            if (v) {
                dom()->var(id)->val_id_of(v);
                vecExpr.push_back(dom()->var(id)->zvar() == v);
            } else {
                ncex *= dom()->n_values(id);
            }
        }
        VLOG_IF(150, ncex > 1, "Has free variable");
        return z3::mk_and(vecExpr);
    }

    map<std::pair<unsigned, int>, double> cache_count_models;

    static constexpr int DEF_LIM_MODELS = kMAX<int>;

    double count_models(const expr &ex, int lim_n_runs = DEF_LIM_MODELS, z3::expr_vector *v_models = nullptr) {
        std::pair<unsigned, int> cache_key = {ex.hash(), lim_n_runs};
        if (v_models == nullptr) {
            auto it = cache_count_models.find(cache_key);
            if (it != cache_count_models.end()) return it->second;
        }
        double &ncex = (cache_count_models[cache_key] = 0);

        auto solver = ctx()->zscope();
        solver->add(ex);

        for (int n_run = 0; n_run < lim_n_runs; ++n_run) {
            z3::check_result checkres = solver->check();
            if (checkres == z3::unsat) {
                break;
            } else if (checkres == z3::unknown) {
                LOG(WARNING, "Z3 solver returns unknown");
                maxi(ncex, 1.0);
                break;
            }
            CHECK_EQ(checkres, z3::sat);
            z3::model m = solver->get_model();
            double this_cex = 1;
            expr mexpr = to_expr(m, this_cex);
            ncex += this_cex;
            solver->add(!mexpr);
            if (v_models != nullptr) v_models->push_back(mexpr);
            //VLOG(100, "ncex {:10G}", ncex);
        }
        return ncex;
    }

    static bool is_all(const str &s, char c) {
        if (s.empty()) return false;
        for (char x : s) if (x != c) return false;
        return true;
    }

    map<str, LocData> read_file(const str &path, vec<str> *global_comments = nullptr) {
        map<str, LocData> res;
        std::ifstream f(path);
        CHECKF(!f.fail(), "Error reading file: {}", path);
        str line;
        int read_state = 0;
        vec<str> locs, comments;
        str sexpr, stree;
        bool ignored = false;
        set<unsigned> sexprid;
        unsigned cur_id = 0;
        int cnt_ignored = 0;
        while (getline(f, line)) {
            boost::algorithm::trim(line);
            if (line.empty()) continue;
            if (line[0] == '#') {
                comments.push_back(line);
                if (line.find("IGNORED") != str::npos) ignored = true;
                continue;
            }
            if (is_all(line, '-')) {
                switch (read_state) {
                    case 0:
                        CHECKF(!locs.empty(), "Empty location set ({})", path);
                        break;
                    case 1:
                        boost::algorithm::trim(sexpr);
                        CHECKF(!sexpr.empty(), "Empty sexpr location set ({})", path);
                        break;
                    default:
                        CHECK(0);
                }
                read_state++;
                CHECK_LE(read_state, 2);
                continue;
            } else if (is_all(line, '=')) {
                if (locs.empty()) {
                    if (global_comments != nullptr)
                        vec_move_append(*global_comments, comments);
                    comments.clear();
                    continue;
                }
                expr e = dom()->parse_string(move(sexpr));
                PMutCTree tree = new CTree(ctx_mut());
                std::stringstream stream_stree(stree);
                tree->deserialize(stream_stree);
                CHECKF(sexprid.insert(e.hash()).second, "Duplicated expression ({}), loc {}", path,
                       locs.empty() ? str() : locs.at(0));
                //LOG(INFO, "EXPR: ") << e;
                expr tree_e = tree->build_zexpr(CTree::FreeMix);
                CHECKF(count_models(e != tree_e, 1) == 0, "Mismatch tree and expr: {}", path);
                bool is_first = true;
                for (const str &s : locs) {
                    CHECKF(!res.contains(s), "Duplicated location ({}): {}", path, s);
                    res.emplace(s, LocData{e, tree, ignored, is_first, cur_id, move(comments)});
                    is_first = false;
                }
                cur_id++, cnt_ignored += ignored;
                ignored = false, read_state = 0, locs.clear(), comments.clear(), sexpr.clear(), stree.clear();
                continue;
            }
            switch (read_state) {
                case 0: {
                    for (char &c : line) if (c == ',') c = ' ';
                    std::stringstream ss(line);
                    str tok;
                    while (ss >> tok) locs.emplace_back(move(tok));
                    break;
                }
                case 1: {
                    sexpr += line, sexpr += '\n';
                    break;
                }
                case 2: {
                    CHECKF(stree.empty(), "Invalid tree data ({})", path);
                    stree = move(line);
                    break;
                }
                default:
                    CHECK(0);
            }
        }
        LOG(INFO, "Read {} locs ({} uniq, {} ignored) from {}", res.size(), cur_id, cnt_ignored, path);
        return res;
    }

    // compare two output
    void run_analyzer_0() {
        auto finp = get_inp();
        CHECK_EQ(finp.size(), 2) << "Need two input files to compare";
        auto ma = read_file(finp.at(0));
        auto mb = read_file(finp.at(1));
        set<unsigned> sdiff, smissing, slocsa, slocsb, sprintdiff;
        int cntdiff = 0, cntmissing = 0;
        double totalcex = 0;
        z3::expr_vector ve_cex(zctx());
        for (const auto &p : ma) {
            const auto &dat = p.second;
            slocsa.insert(dat.id());
            auto it = mb.find(p.first);
            if (it == mb.end()) {
                VLOG(dat.is_first ? 5 : 7, "{} not found in B", p.first);
                smissing.insert(dat.id()), cntmissing++;
                continue;
            }
            expr expr_cex = dat.e != it->second.e;
            ve_cex.push_back(expr_cex);
            double num_cex = count_models(expr_cex, kMAX<int>, &ve_cex);
            if (num_cex == 0) {
                //VLOG(0, "{} ok", p.first);
            } else {
                VLOG(dat.is_first ? 5 : 7, "{} diff (cex = {:G}){} (id={})",
                     p.first, num_cex,
                     dat.ignored || it->second.ignored ? " (IGNORED)" : "",
                     dat.id());
                if (sprintdiff.insert(dat.id()).second) {
                    totalcex += num_cex;
                    GVLOG(10) << "\nA: " << dat.e << "\nB: " << it->second.e;
                }
                sdiff.insert(dat.id()), cntdiff++;
            }
        }
        for (const auto &p : mb) {
            slocsb.insert(p.second.id());
            if (!ma.contains(p.first)) {
                VLOG(p.second.is_first ? 5 : 7, "{} not found in A", p.first);
                smissing.insert(p.second.id()), cntmissing++;
            }
        }
        LOG(INFO, "Counting total cex");
        double total_cex = count_models(z3::mk_or(ve_cex));

        int repeat_id = ctx()->get_option_as<int>("_repeat_id");
        LOG(INFO, "R[{:>2}] {:=^80}",
            repeat_id, "  FINAL RESULT  ");
        LOG(INFO, "R[{:>2}] Total: diff {:^4}, miss {:^4}, locs A {:^4}, B {:^4}",
            repeat_id, cntdiff, cntmissing, ma.size(), mb.size());
        LOG(INFO, "R[{:>2}] Uniq : diff {:^4}, miss {:^4}, locs A {:^4}, B {:^4}, cex {:G}",
            repeat_id,
            sdiff.size(), smissing.size(),
            slocsa.size(), slocsb.size(), totalcex);
        LOG(INFO, "R[{:>2}] Total CEX: {:G}",
            repeat_id, total_cex);
    }

    // try random conf
    void run_analyze_1() {
        boost::timer::cpu_timer timer;
        ctx()->init_runner();

        auto finp = get_inp();
        CHECK_EQ(finp.size(), 1) << "Need 1 input file";
        auto ma = read_file(finp.at(0));

        int n_iter = 0;
        bool run_full = false;
        if (ctx()->has_option("rand")) n_iter = ctx()->get_option_as<int>("rand");
        else if (ctx()->has_option("full")) n_iter = (int) dom()->config_space(), run_full = true;
        CHECK_GT(n_iter, 0) << "Invalid arguments";

        int n_batch = ctx()->get_option_as<int>("batch-size");
        if (n_batch == 0) n_batch = 100;
        CHECK_GT(n_batch, 0);

        tsl::robin_set<hash_t> all_configs;
        all_configs.reserve(n_iter);
        tsl::robin_set<hash_t> wrong_configs, wrong_configs_nig;
        set<str> wrong_locs, wrong_locs_uniq, missing_locs, empty_set;
        set<str> wrong_locs_nig, wrong_locs_uniq_nig, missing_locs_nig;

        vec<PMutConfig> batch_confs(n_batch), generated_full_configs;
        vec<set<str>> batch_locs;
        for (auto &c : batch_confs) c = new Config(ctx_mut());
        int it = n_batch;

        int nvars = dom()->n_vars();
        vec<int> dom_nvals(nvars);
        for (int var = 0; var < nvars; ++var)
            dom_nvals[var] = dom()->n_values(var);

        if (run_full) {
            LOG(INFO, "Generating full {:G} configs", dom()->config_space());
            generated_full_configs = dom()->gen_all_configs();
            CHECK_EQ(n_iter, generated_full_configs.size());
        }
        int it_full_configs = 0;

        for (int iter = 1; iter <= n_iter; ++iter) {
            if (it == n_batch) {
                if (run_full) {
                    for (int i = 0; i < n_batch; ++i) {
                        if (it_full_configs < sz(generated_full_configs)) {
                            batch_confs[i] = move(generated_full_configs[it_full_configs++]);
                            all_configs.insert(batch_confs[i]->hash());
                        } else {
                            batch_confs.erase(batch_confs.begin() + i, batch_confs.end());
                            break;
                        }
                    }
                } else {
                    for (int i = 0; i < n_batch; ++i) {
                        const PMutConfig &c = batch_confs[i];
                        do {
                            for (int var = 0; var < nvars; ++var)
                                c->set(var, Rand.get(dom_nvals[var]));
                        } while (!all_configs.insert(c->hash(true)).second);
                    }
                }
                batch_locs = ctx_mut()->runner()->run(batch_confs);
                CHECK_EQ(batch_confs.size(), batch_locs.size());
                it = 0;
            }

            const PMutConfig &c = batch_confs.at(it);
            CHECK_NE(c, nullptr);
            const set<str> &e = batch_locs.at(it);
            ++it;

            vec<int8_t> eval_cache(ma.size(), -1);
            eval_cache.reserve(ma.size());
            for (const auto &p : ma) {
                const auto &dat = p.second;
                int8_t &cache_val = eval_cache[dat.id()];
                bool eval_res, uniq_loc;
                if (cache_val == -1) {
                    dat.tree->build_interpreter();
                    cache_val = (eval_res = dat.tree->interpret(*c));
                    //CHECK_EQ(eval_res, c->eval(dat.e));
                    uniq_loc = true;
                } else {
                    eval_res = (bool) cache_val;
                    uniq_loc = false;
                }

                bool truth = e.contains(p.first);
                if (eval_res != truth) {
                    wrong_configs.insert(c->hash());
                    wrong_locs.insert(p.first);
                    if (uniq_loc) wrong_locs_uniq.insert(p.first);

                    if (!dat.ignored) {
                        wrong_configs_nig.insert(c->hash());
                        wrong_locs_nig.insert(p.first);
                        if (uniq_loc) {
                            wrong_locs_uniq_nig.insert(p.first);
                            LOG_FIRST_N(WARNING, 1000)
                                << fmt::format("Wrong: loc {}, ", p.first)
                                << *c << "  #    " << c->to_str_raw();
                        }
                    }
                }
            }
            for (const auto &s : e) {
                if (!ma.contains(s)) {
                    LOG_FIRST_N(ERROR, 1000)
                        << fmt::format("Missing: loc {}, ", s) << *c
                        << "  #    " << c->to_str_raw();
                    missing_locs.insert(s);
                }
            }
            if (it == 1) {
                LOG(INFO, "{:>7} {:>4} {:>7} | {:>4} {:>4} {:>4} | {:>4} {:>4} {:>4} | ",
                    all_configs.size(), wrong_configs.size(), ctx()->runner()->n_cache_hit(),
                    wrong_locs.size(), wrong_locs_uniq.size(), missing_locs.size(),
                    wrong_configs_nig.size(), wrong_locs_nig.size(), wrong_locs_uniq_nig.size()
                )
                        << (sz(wrong_locs_uniq) <= 10 ? wrong_locs_uniq : empty_set)
                        << (sz(missing_locs) <= 10 ? missing_locs : empty_set);
            }
        }
        LOG(INFO, "Total        time: {}", timer.format(0));
        LOG(INFO, "Runner       time: {}", boost::timer::format(ctx()->runner()->timer(), 0));
        LOG(INFO, "Multi-runner time: {}", boost::timer::format(ctx()->runner()->total_elapsed(), 0));
    }

    bool is_pure(const expr &e, Z3_decl_kind outer, Z3_decl_kind inner) {
        CHECK_EQ(e.decl().decl_kind(), outer);
        int n_args = (int) e.num_args();
        CHECK_GT(n_args, 0);
        for (int i = 0; i < n_args; ++i) {
            expr g = e.arg(i);
            if (g.is_eq()) continue;
            if (g.decl().decl_kind() == inner) {
                long last_id = -1;
                int n_args_g = (int) g.num_args();
                for (int j = 0; j < n_args_g; ++j) {
                    expr x = g.arg(j);
                    if (!x.is_eq()) return false;
                    CHECK_EQ(x.num_args(), 2);
                    expr lhs = x.arg(0);

                    bool is_var = false;
                    for (const auto &t : *dom()) {
                        if (lhs.id() == t->zvar().id()) {
                            is_var = true;
                            break;
                        }
                    }
                    if (!is_var) return false;

                    if (last_id == -1) last_id = lhs.id();
                    else if (last_id != lhs.id()) return false;
                }
            } else {
                return false;
            }
        }
        return true;
    }

    bool is_igen_compat(const expr &e) {
        if (!e.is_and()) return false;
        int n_args = (int) e.num_args();
        CHECK_GT(n_args, 0);
        bool has_or_clause = false;
        for (int i = 0; i < n_args; ++i) {
            expr g = e.arg(i);
            if (g.is_eq()) continue;
            if (has_or_clause || !g.is_or()) return false;
            has_or_clause = true;
            if (!is_pure(g, Z3_OP_OR, Z3_OP_UNINTERPRETED)) return false;
        }
        return true;
    }

    // read exp result
    map<str, boost::any> run_analyze_2() {
        map<str, boost::any> ret;
        vec<str> gcmt;
        auto finp = get_inp();
        CHECK_EQ(finp.size(), 1) << "Need 1 input file";
        auto ma = read_file(finp.at(0), &gcmt);

        vec<str> params;
        boost::algorithm::split_regex(params, gcmt.at(2), boost::regex("[ |]+"));
//        LOG(INFO, "params {}", fmt::join(params, ","));

//        fmt::print(out, "# {:>8} {:>4} {:>4} | {:>5} {:>5} | {:>2} {:>4} | {:>5} {:>5} {:>6} | {}\n======\n",
//                   cov()->n_configs(), cov()->n_locs(), v_uniq.size(),
//                   ctx()->runner()->n_cache_hit(), ctx()->runner()->n_locs(),
//                   repeat_id_, iter,
//
//                   timer.elapsed().wall / NS, ctx()->runner()->local_timer().wall / NS,
//                   ctx()->runner()->total_elapsed().wall / NS,
//
//                   state_hash().str()
//        );

        int cnt_singular = 0, cnt_and = 0, cnt_or = 0, cnt_mixed = 0;
        int cnt_pure = 0, cnt_mix_ok = 0, cnt_mix_fail = 0;
        vec<int> vec_len;
        for (const auto &p : ma) {
            const auto &dat = p.second;
            if (!dat.is_first) continue;
            const expr &e = dat.e;

            set<str> vars;
            collect_vars(vars, e);
            vec_len.push_back(sz(vars));

            VLOG(10, "{}\n {}", p.first, e.to_string());
            CHECK(e.is_app());
            if (e.is_const() || e.is_eq()) {
                VLOG(10, "=> Singular");
                cnt_singular++, cnt_pure++;
                continue;
            }
            if (e.is_and() && is_pure(e, Z3_OP_AND, Z3_OP_OR)) {
                VLOG(10, "=> And");
                cnt_and++, cnt_pure++;
                continue;
            }
            if (e.is_or() && is_pure(e, Z3_OP_OR, Z3_OP_UNINTERPRETED)) {
                VLOG(10, "=> Or");
                cnt_or++, cnt_pure++;
                continue;
            }
            bool ig_compat = is_igen_compat(e);
            VLOG(10, "=> Mixed{}", ig_compat ? ", igen compat" : "");
            cnt_mixed++;
            if (ig_compat) cnt_mix_ok++; else cnt_mix_fail++;
        }
#define RET_PARAM(name) ret[#name] = name
        RET_PARAM(cnt_singular);
        RET_PARAM(cnt_and);
        RET_PARAM(cnt_or);
        RET_PARAM(cnt_mixed);
        RET_PARAM(cnt_pure);
        RET_PARAM(cnt_mix_ok);
        RET_PARAM(cnt_mix_fail);
#undef RET_PARAM
        ret["cnt_total"] = cnt_singular + cnt_and + cnt_or + cnt_mixed;
        ret["max_len"] = (int) *std::max_element(vec_len.begin(), vec_len.end());
        ret["median_len"] = (int) vec_median(vec_len);

#define RET_PARAM(type, p, pos) ret[p] = boost::lexical_cast<type>(params.at(pos))
        RET_PARAM(int, "n_configs", 1);
        RET_PARAM(int, "n_locs", 2);
        RET_PARAM(int, "n_locs_uniq", 3);

        RET_PARAM(int, "n_cache_hit", 4);
        RET_PARAM(int, "n_locs_total", 5);

        RET_PARAM(int, "repeat_id", 6);
        RET_PARAM(int, "iter", 7);

        RET_PARAM(int, "t_total", 8);
        RET_PARAM(int, "t_runner", 9);
        ret["t_search"] = boost::any_cast<int>(ret["t_total"]) - boost::any_cast<int>(ret["t_runner"]);
        RET_PARAM(int, "t_runner_total", 10);

        RET_PARAM(std::string, "hash", 11);
#undef RET_PARAM
        return ret;
    }

    double calc_mcc(double tp, double tn, double fp, double fn) {
        double nom = tp * tn - fp * fn;
        double denom = sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn));
        if (denom == 0) denom = 1;

        return nom / denom;
    }

    double calc_fscore(double tp, double fp, double fn) {
        double p = (tp + fp == 0) ? 0 : double(tp) / (tp + fp);
        double r = (tp + fn == 0) ? 0 : double(tp) / (tp + fn);
        double f = (r + p == 0) ? 0 : double(2 * r * p) / (r + p);
        return f;
    }

    // calculate MCC
    // https://en.wikipedia.org/wiki/Matthews_correlation_coefficient
    map<str, boost::any> run_analyze_3() {
        auto finp = get_inp();
        CHECK_EQ(finp.size(), 2) << "Need two input files to compare";
        vec<str> gcmt, params;
        auto m_truth = read_file(finp.at(0));
        auto m_predicate = read_file(finp.at(1), &gcmt);
        boost::algorithm::split_regex(params, gcmt.at(2), boost::regex("[ |]+"));

        double cspace = dom()->config_space();
        vec<PMutConfig> all_confs = dom()->gen_all_configs();
        CHECK_EQ(cspace, all_confs.size());

        double total_mcc = 0, total_fscore = 0;
        int cnt_interactions = 0, cnt_exact = 0;
        std::stringstream wrong_locs;
        for (const auto &p : m_truth) {
            if (!p.second.is_first) continue;
            cnt_interactions++;
            double mcc = 0, fscore = 0;
            auto _pred = m_predicate.find(p.first);
            if (_pred == m_predicate.end()) {
                // Loc not found
                wrong_locs << '!' << p.first << ' ';
                mcc = fscore = 0;
            } else {
                const auto &truth = p.second, &pred = _pred->second;
                if (count_models(truth.e != pred.e, 1) == 0) {
                    mcc = fscore = 1;
                    cnt_exact++;
                } else {
                    wrong_locs << p.first << ' ';
                    double tp = 0, tn = 0, fp = 0, fn = 0;
                    truth.tree->build_interpreter(), pred.tree->build_interpreter();
                    for (const auto &c : all_confs) {
                        bool v_truth = truth.tree->interpret(*c), v_pred = pred.tree->interpret(*c);
                        if (v_truth && v_pred) tp++;
                        else if (!v_truth && !v_pred) tn++;
                        else if (!v_truth && v_pred) fp++;
                        else if (v_truth && !v_pred) fn++;
                    }
                    VLOG(1, "tp, tn, fp, fn = {:10G}, {:10G}, {:10G}, {:10G}", tp, tn, fp, fn);
                    CHECK_EQ(tp + tn + fp + fn, cspace);
                    mcc = calc_mcc(tp, tn, fp, fn);
                    fscore = calc_fscore(tp, fp, fn);
                }
            }
            total_mcc += mcc, total_fscore += fscore;
            if (mcc != 1 || fscore != 1) LOG(INFO, "Loc {}: mcc {}, fscore {}", p.first, mcc, fscore);
        }

        map<str, boost::any> ret;
        ret["avg_mcc"] = (total_mcc / cnt_interactions);
        ret["avg_f"] = (total_fscore / cnt_interactions);
        ret["cnt_interactions"] = cnt_interactions;
        ret["cnt_exact"] = cnt_exact;
        ret["cnt_wrong"] = cnt_interactions - cnt_exact;
        ret["delta_locs"] = sz(m_predicate) - sz(m_truth);
        ret["n_configs"] = boost::lexical_cast<int>(params.at(1));
        ret["wrong_locs"] = wrong_locs.str();
        return ret;
    }

    PMutConfig gen_model(expr e) {
        auto &solver = ctx()->zsolver();
        if (solver.check(1, &e) == z3::sat) {
            PMutConfig c = new Config(ctx_mut());
            z3::model m = solver.get_model();
            for (int id = 0; id < dom()->n_vars(); ++id) {
                expr const_interp = m.get_const_interp(dom()->var(id)->zvar().decl());
                if (const_interp)
                    c->set(id, dom()->var(id)->val_id_of(const_interp));
                else
                    c->set(id, Rand.get(dom()->n_values(id)));
            }
            return c;
        }
        return nullptr;
    }

    map<str, boost::any> run_analyze_cmin() {
        auto finp = get_inp();
        CHECK_EQ(finp.size(), 1) << "Need 1 input file";
        auto ma = read_file(finp.at(0));

        vec<LocData> vd;
        for (const auto &p : ma) {
            const auto &d = p.second;
            d.tree->build_interpreter();
            if (d.is_first) vd.emplace_back(d);
        }
//        int beg_sz = sz(vd);
//        for (int i = 0; i < sz(vd); ++i) {
//            for (int j = i + 1; j < sz(vd); ++j) {
//                if (gen_model(!z3::implies(vd[i].e, vd[j].e)) == nullptr) {
//                    vd.erase(vd.begin() + j);
//                    --j;
//                } else if (gen_model(!z3::implies(vd[j].e, vd[i].e)) == nullptr) {
//                    vd.erase(vd.begin() + i);
//                    --i;
//                    break;
//                }
//            }
//        }
//        LOG(INFO, "Removed {} interactions", beg_sz - sz(vd));
        if (ctx()->get_option_as<int>("_repeat_id") > 0) {
            LOG(INFO, "Randomly shuffle the interactions");
            Rand.shuffle(vd);
        }

        vec<PMutConfig> confs;
        set<hash_t> set_hash;
        for (int L = 0; L < sz(vd);) {
            int R = L;
            z3::expr e = ctx()->ztrue();
            while (R < sz(vd)) {
                const auto &d = vd.at(R);
                for (const auto &c : confs) {
                    if (d.tree->interpret(*c)) {
                        R++;
                        goto continue_while;
                    }
                }
                //VLOG(20, "Solve conj of {} interactions ({} -> {})", R - L, L, R);
                if (gen_model(e && d.e)) {
                    e = e && d.e;
                    R++;
                } else {
                    break;
                }
                continue_while:
                continue;
            }
            CHECK_GT(R, L);
            VLOG(20, "Conj of {} interactions ({} -> {})", R - L, L, R);
            PMutConfig c = gen_model(e);
            CHECK_NE(c, nullptr);
            confs.emplace_back(c);
            CHECK(set_hash.insert(c->hash()).second);
            L = R;
        }

        for (const auto &p : ma) {
            bool ok = false;
            for (const auto &c : confs) {
                if (p.second.tree->interpret(*c)) {
                    ok = true;
                    break;
                }
            }
            CHECK(ok);
        }

        map<str, boost::any> ret;
        ret["min_c"] = sz(confs);
        return ret;
    }

    bool is_var(const expr &e) {
        return e.is_const() && e.decl().decl_kind() == Z3_OP_UNINTERPRETED;
    }

    void collect_vars(set<str> &res, const expr &e) {
        if (is_var(e)) {
            res.insert(e.decl().name().str());
        } else {
            int nargs = e.num_args();
            for (int i = 0; i < nargs; ++i) {
                collect_vars(res, e.arg(i));
            }
        }
    }

    map<str, boost::any> run_analyze_complexity() {
        auto finp = get_inp();
        CHECK_EQ(finp.size(), 1) << "Need 1 input file";
        auto ma = read_file(finp.at(0));

        vec<int> samples;
        int sum = 0, count = 0;
        for (const auto &p : ma) {
            const auto &d = p.second;
            if (!d.is_first) continue;

            set<str> s;
            collect_vars(s, d.e);
            sum += sz(s), count++;
            samples.push_back(sz(s));
        }

        double avg = double(sum) / count;
        double dev = 0;
        for (int x : samples) {
            dev += (x - avg) * (x - avg);
        }
        dev = std::sqrt(dev / count);

        map<str, boost::any> ret;
        ret["avg_len"] = avg;
        ret["dev_len"] = dev;
        return ret;
    }

    vec<str> get_inp() {
        str inp = ctx()->get_option_as<str>("input");
        boost::algorithm::replace_all(inp, ",", " ");
        std::stringstream ss(inp);
        str tok;
        vec<str> res;
        while (ss >> tok) res.push_back(tok);
        return res;
    }

    map<str, boost::any> run_alg() {
        switch (ctx()->get_option_as<int>("alg-version")) {
            case 0:
                run_analyzer_0();
                return {};
            case 1:
                run_analyze_1();
                return {};
            case 2:
                return run_analyze_2();
            case 3:
                return run_analyze_3();
            case 4:
                return run_analyze_cmin();
            case 5:
                return run_analyze_complexity();
            default:
                CHECK(0);
        }
    }
};

map<str, boost::any> run_analyzer(const map<str, boost::any> &opts) {
    map<str, boost::any> ret;
    PMutContext ctx = new Context();
    ctx->set_options(opts);
    ctx->init();
    //ctx->init_runner();
    {
        ptr<Analyzer> ite_alg = new Analyzer(ctx);
        ret = ite_alg->run_alg();
    }
    ctx->cleanup();
    return ret;
}

}


#include <boost/algorithm/string.hpp>
#include "ContentHash.h"

namespace PDLS {

    void ContentHash::convertStringToRawByteHash() {
        if(_stringHash != "") {
            if(_rawByteHash != nullptr) delete _rawByteHash;

            _rawByteHash = new RawByteHash;

            //copy string hash c_str no non const buffer
            char *buffer = new char[_stringHash.size()];
            std::copy(_stringHash.c_str(), _stringHash.c_str() + _stringHash.size(), buffer);

            //reinterpret char* to unsigned char* proper for byte representation
            _rawByteHash->_data = reinterpret_cast<unsigned char*>(buffer);

            switch (_type) {
                // x bits/8 <=> conversion bits to byte
                case ContentHashType::MD5:
                    _rawByteHash->_size = 128/8;
                    break;
                case ContentHashType::SHA1:
                    _rawByteHash->_size = 160/8;
                    break;
                case ContentHashType::SHA224:
                    _rawByteHash->_size = 224/8;
                    break;
                case ContentHashType::SHA256:
                    _rawByteHash->_size = 256/8;
                    break;
                case ContentHashType::SHA384:
                    _rawByteHash->_size = 384/8;
                    break;
                case ContentHashType::SHA512:
                    _rawByteHash->_size = 512/8;
                    break;
            }
        }
    }

    void ContentHash::validateStringHash()  {
        //string contains 1 char per hex digit
        //1 byte <=> 8 bits <=> 2 hex digits ==> x bits / 8 / 2 = x bits / 4 <=> number of hex digits
        switch (_type){
            case ContentHashType::MD5:
                if(_stringHash.size() != 128/4)
                    throw InvalidContentHashException(_stringHash, _prefix, _postfix);
                break;
            case ContentHashType::SHA1:
                if(_stringHash.size() != 160/4)
                    throw InvalidContentHashException(_stringHash, _prefix, _postfix);
                break;
            case ContentHashType::SHA224:
                if(_stringHash.size() != 224/4)
                    throw InvalidContentHashException(_stringHash, _prefix, _postfix);
                break;
            case ContentHashType::SHA256:
                if(_stringHash.size() != 256/4)
                    throw InvalidContentHashException(_stringHash, _prefix, _postfix);
                break;
            case ContentHashType::SHA384:
                if(_stringHash.size() != 384/4)
                    throw InvalidContentHashException(_stringHash, _prefix, _postfix);
                break;
            case ContentHashType::SHA512:
                if(_stringHash.size() != 512/4)
                    throw InvalidContentHashException(_stringHash, _prefix, _postfix);
                break;
        }
    }

    void ContentHash::convertRawByteToStringHash()  {
        if(_rawByteHash != nullptr && _rawByteHash->_data != nullptr) {
            std::ostringstream outs;

            for (size_t i = 0; i < _rawByteHash->_size; i++)
                outs << std::hex << std::setw(2) << std::setfill('0') << static_cast<int>(_rawByteHash->_data[i]);

            _stringHash = outs.str();
        }
    }

    void ContentHash::trimPrefixFromStringHashInput() {
        boost::erase_all(_stringHash, _prefix);
    }

    void ContentHash::trimPostfixFromStringHashInput() {
        boost::erase_all(_stringHash, _postfix);
    }

}


/*
 Copyright (C) 2019 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

#include <qle/termstructures/eqcommoptionsurfacestripper.hpp>
#include <boost/make_shared.hpp>
#include <ql/instruments/impliedvolatility.hpp>
#include <ql/math/solver1d.hpp>
#include <ql/pricingengines/vanilla/analyticeuropeanengine.hpp>
#include <ql/termstructures/volatility/equityfx/blackconstantvol.hpp>
#include <qle/pricingengines/baroneadesiwhaleyengine.hpp>
#include <qle/termstructures/blackvariancesurfacesparse.hpp>
#include <qle/termstructures/pricetermstructureadapter.hpp>

using std::function;
using std::map;
using std::pair;
using std::set;
using std::vector;
using namespace QuantLib;

namespace {

// Utility method to create the list of options to be used at an expiry date for stripping.
function<bool(Real,Real)> comp = [](Real a, Real b) { return !close(a, b) && a < b; };

map<Real, Option::Type, decltype(comp)> createStrikes(Real forward, const vector<Real>& cStrikes,
    const vector<Real>& pStrikes, bool preferOutOfTheMoney) {

    // Firstly create the restricted vector of call and put strikes.
    vector<Real> rcStks;
    copy_if(cStrikes.begin(), cStrikes.end(), back_inserter(rcStks),[forward, preferOutOfTheMoney](Real stk) {
        return (preferOutOfTheMoney && stk >= forward) || (!preferOutOfTheMoney && stk <= forward); });
    vector<Real> rpStks;
    copy_if(pStrikes.begin(), pStrikes.end(), back_inserter(rpStks), [forward, preferOutOfTheMoney](Real stk) {
        return (preferOutOfTheMoney && stk <= forward) || (!preferOutOfTheMoney && stk >= forward); });

    // Create the empty map.
    map<Real, Option::Type, decltype(comp)> res(comp);

    // If both restricted vectors are empty, return an empty map
    if (rcStks.empty() && rpStks.empty())
        return res;

    // At least one of the restricted strike vectors are non empty so populate the map.
    if (!rcStks.empty() && !rpStks.empty()) {
        // Most common case hopefully. Use both sets of strikes.
        // Could have the fwd strike in both restricted sets from the logic above. Favour Call here via overwrite.
        for (Real stk : rpStks)
            res[stk] = Option::Put;
        for (Real stk : rcStks)
            res[stk] = Option::Call;
    } else if (rpStks.empty()) {
        // If restricted put strikes are empty, use all the call strikes
        for (Real stk : cStrikes)
            res[stk] = Option::Call;
    } else if (rcStks.empty()) {
        // If restricted call strikes are empty, use all the put strikes
        for (Real stk : pStrikes)
            res[stk] = Option::Put;
    }

    return res;
}

}

namespace QuantExt {

OptionSurfaceStripper::OptionSurfaceStripper(
    const boost::shared_ptr<OptionInterpolatorBase>& callSurface,
    const boost::shared_ptr<OptionInterpolatorBase>& putSurface,
    const Calendar& calendar,
    const DayCounter& dayCounter,
    Exercise::Type type,
    bool lowerStrikeConstExtrap,
    bool upperStrikeConstExtrap,
    bool timeFlatExtrapolation,
    bool preferOutOfTheMoney,
    Solver1DOptions solverOptions)
    : callSurface_(callSurface),
      putSurface_(putSurface),
      calendar_(calendar),
      dayCounter_(dayCounter),
      type_(type),
      lowerStrikeConstExtrap_(lowerStrikeConstExtrap),
      upperStrikeConstExtrap_(upperStrikeConstExtrap),
      timeFlatExtrapolation_(timeFlatExtrapolation),
      preferOutOfTheMoney_(preferOutOfTheMoney),
      solverOptions_(solverOptions),
      havePrices_(boost::dynamic_pointer_cast<OptionPriceSurface>(callSurface_)) {

    QL_REQUIRE(callSurface_->referenceDate() == putSurface_->referenceDate(),
        "Mismatch between Call and Put reference dates in OptionSurfaceStripper");

    registerWith(Settings::instance().evaluationDate());

    // Set up that is only needed if we have price based surfaces and we are stripping volatilities.
    if (havePrices_) {

        // Check that there is also a put price surface
        QL_REQUIRE(boost::dynamic_pointer_cast<OptionPriceSurface>(putSurface_),
            "OptionSurfaceStripper: call price surface provided but no put price surface.");

        setUpSolver();
    }
}

OptionSurfaceStripper::PriceError::PriceError(const VanillaOption& option, SimpleQuote& volatility, Real targetPrice)
    : option_(option), volatility_(volatility), targetPrice_(targetPrice) {}

Real OptionSurfaceStripper::PriceError::PriceError::operator()(Volatility x) const {

    volatility_.setValue(x);

    // Barone Adesi Whaley fails for very small variance, so wrap in a try catch
    Real npv;
    try {
        npv = option_.NPV();
    } catch (...) {
        npv = 0.0;
    }

    return npv - targetPrice_;
}

void OptionSurfaceStripper::performCalculations() const {

    // Create a set of all expiries
    auto tmp = callSurface_->expiries();
    set<Date> allExpiries(tmp.begin(), tmp.end());
    tmp = putSurface_->expiries();
    allExpiries.insert(tmp.begin(), tmp.end());

    boost::shared_ptr<BlackVarianceSurfaceSparse> callVolSurface;
    boost::shared_ptr<BlackVarianceSurfaceSparse> putVolSurface;

    // Switch based on whether surface is direct volatilities or prices to be stripped.
    boost::shared_ptr<PricingEngine> engine;
    boost::shared_ptr<SimpleQuote> volQuote = boost::make_shared<SimpleQuote>(0.1);
    if (havePrices_) {

        // a black scholes process
        boost::shared_ptr<GeneralizedBlackScholesProcess> gbsp = process(volQuote);

        // hard code the engines here
        if (type_ == Exercise::American) {
            engine = boost::make_shared<QuantExt::BaroneAdesiWhaleyApproximationEngine>(gbsp);
        } else if (type_ == Exercise::European) {
            engine = boost::make_shared<QuantExt::AnalyticEuropeanEngine>(gbsp);
        } else {
            QL_FAIL("Unsupported exercise type for option stripping");
        }

    } else {
        // we have variance surfaces, explicitly cast so we can look up vol later
        callVolSurface = boost::dynamic_pointer_cast<BlackVarianceSurfaceSparse>(callSurface_);
        putVolSurface = boost::dynamic_pointer_cast<BlackVarianceSurfaceSparse>(putSurface_);
    }

    // Need to populate these below to feed to BlackVarianceSurfaceSparse
    vector<Real> volStrikes;
    vector<Real> volData;
    vector<Date> volExpiries;

    // Loop over each expiry
    for (const Date& expiry : allExpiries) {

        // Get the forward price at expiry
        Real fwd = forward(expiry);

        // Get the call and put strikes at the expiry date. Each may be empty.
        vector<Real> callStrikes = strikes(expiry, true);
        vector<Real> putStrikes = strikes(expiry, false);

        // We want a set of prices both sides of ATM forward
        // If preferOutOfTheMoney_ is false, we take calls where strike < atm and puts where strike > atm.
        // If preferOutOfTheMoney_ is true, we take calls where strike > atm and puts where strike < atm.
        auto relevantStrikes = createStrikes(fwd, callStrikes, putStrikes, preferOutOfTheMoney_);
        for (const auto& kv : relevantStrikes) {
            if (havePrices_) {
                // Only use the volatility if the root finding was successful
                Real v = implyVol(expiry, kv.first, kv.second, engine, *volQuote);
                if (v != Null<Real>()) {
                    volExpiries.push_back(expiry);
                    volStrikes.push_back(kv.first);
                    volData.push_back(v);
                }
            } else {
                volExpiries.push_back(expiry);
                volStrikes.push_back(kv.first);
                Real v = kv.second == Option::Call ? callVolSurface->blackVol(expiry, kv.first) :
                    putVolSurface->blackVol(expiry, kv.first);
                volData.push_back(v);
            }
        }
    }

    // Populate the variance surface.
    volSurface_ = boost::make_shared<BlackVarianceSurfaceSparse>(
        callSurface_->referenceDate(), calendar_, volExpiries, volStrikes, volData, dayCounter_,
        lowerStrikeConstExtrap_, upperStrikeConstExtrap_, timeFlatExtrapolation_);
}

vector<Real> OptionSurfaceStripper::strikes(const Date& expiry, bool isCall) const {

    const boost::shared_ptr<OptionInterpolatorBase>& surface = isCall ? callSurface_ : putSurface_;
    auto expiries = surface->expiries();
    auto it = find(expiries.begin(), expiries.end(), expiry);

    if (it != expiries.end()) {
        return surface->strikes().at(distance(expiries.begin(), it));
    } else {
        return {};
    }

}

Real OptionSurfaceStripper::implyVol(Date expiry, Real strike, Option::Type type,
    boost::shared_ptr<PricingEngine> engine, SimpleQuote& volQuote) const {

    // Create the option instrument used in the solver.
    boost::shared_ptr<StrikedTypePayoff> payoff = boost::make_shared<PlainVanillaPayoff>(type, strike);
    boost::shared_ptr<Exercise> exercise;
    if (type_ == Exercise::American) {
        exercise = boost::make_shared<AmericanExercise>(expiry);
    } else if (type_ == Exercise::European) {
        exercise = boost::make_shared<EuropeanExercise>(expiry);
    } else {
        QL_FAIL("OptionSurfaceStripper: unsupported exercise type for option stripping.");
    }
    VanillaOption option(payoff, exercise);
    option.setPricingEngine(engine);

    // Get the target price from the surface.
    Real targetPrice = type == Option::Call ? callSurface_->getValue(expiry, strike)
        : putSurface_->getValue(expiry, strike);

    // Attempt to calculate the implied volatility.
    Real vol = Null<Real>();
    try {
        PriceError f(option, volQuote, targetPrice);
        vol = solver_(f);
    } catch (const Error&) {
    }

    return vol;
}

void OptionSurfaceStripper::setUpSolver() {

    // Check that enough solver options have been provided.
    const Real& guess = solverOptions_.initialGuess;
    QL_REQUIRE(guess != Null<Real>(), "OptionSurfaceStripper: need a valid initial " <<
        "guess for a price based surface.");

    const Real& accuracy = solverOptions_.accuracy;
    QL_REQUIRE(accuracy != Null<Real>(), "OptionSurfaceStripper: need a valid accuracy " <<
        "for a price based surface.");

    // Set maximum evaluations if provided.
    if (solverOptions_.maxEvaluations != Null<Size>())
        brent_.setMaxEvaluations(solverOptions_.maxEvaluations);

    // Check and set the lower bound and upper bound
    if (solverOptions_.lowerBound != Null<Real>() && solverOptions_.upperBound != Null<Real>()) {
        QL_REQUIRE(solverOptions_.lowerBound < solverOptions_.upperBound, "OptionSurfaceStripper: lowerBound (" <<
            solverOptions_.lowerBound << ") should be less than upperBound (" << solverOptions_.upperBound << ")");
    }

    if (solverOptions_.lowerBound != Null<Real>())
        brent_.setLowerBound(solverOptions_.lowerBound);
    if (solverOptions_.upperBound != Null<Real>())
        brent_.setUpperBound(solverOptions_.upperBound);

    // Choose a min/max or step solver based on parameters provided, favouring the min/max based version.
    const Real& min = solverOptions_.minMax.first;
    const Real& max = solverOptions_.minMax.second;
    const Real& step = solverOptions_.step;
    using std::placeholders::_1;
    if (min != Null<Real>() && max != Null<Real>()) {
        typedef Real (Brent::* MinMaxSolver)(const PriceError&, Real, Real, Real, Real) const;
        solver_ = std::bind(static_cast<MinMaxSolver>(&Brent::solve), &brent_, _1, accuracy, guess, min, max);
    } else if (step != Null<Real>()) {
        typedef Real(Brent::* StepSolver)(const PriceError&, Real, Real, Real) const;
        solver_ = std::bind(static_cast<StepSolver>(&Brent::solve), &brent_, _1, accuracy, guess, step);
    } else {
        QL_FAIL("OptionSurfaceStripper: need a valid step size or (min, max) pair for a price based surface.");
    }

}

boost::shared_ptr<BlackVolTermStructure> OptionSurfaceStripper::volSurface() {
    calculate();
    return volSurface_;
}

EquityOptionSurfaceStripper::EquityOptionSurfaceStripper(
    const Handle<EquityIndex>& equityIndex,
    const boost::shared_ptr<OptionInterpolatorBase>& callSurface,
    const boost::shared_ptr<OptionInterpolatorBase>& putSurface,
    const Calendar& calendar,
    const DayCounter& dayCounter,
    Exercise::Type type,
    bool lowerStrikeConstExtrap,
    bool upperStrikeConstExtrap,
    bool timeFlatExtrapolation,
    bool preferOutOfTheMoney,
    Solver1DOptions solverOptions)
    : OptionSurfaceStripper(callSurface, putSurface, calendar, dayCounter, type, lowerStrikeConstExtrap,
        upperStrikeConstExtrap, timeFlatExtrapolation, preferOutOfTheMoney, solverOptions), equityIndex_(equityIndex) {
    registerWith(equityIndex_);
}

boost::shared_ptr<GeneralizedBlackScholesProcess> EquityOptionSurfaceStripper::process(
    const boost::shared_ptr<QuantLib::SimpleQuote>& volatilityQuote) const {

    Handle<BlackVolTermStructure> vts(boost::make_shared<BlackConstantVol>(
        callSurface_->referenceDate(), calendar_, Handle<Quote>(volatilityQuote), dayCounter_));

    return boost::make_shared<BlackScholesMertonProcess>(equityIndex_->equitySpot(),
        equityIndex_->equityDividendCurve(), equityIndex_->equityForecastCurve(), vts);
}

Real EquityOptionSurfaceStripper::forward(const Date& date) const {
    return equityIndex_->forecastFixing(date);
}

CommodityOptionSurfaceStripper::CommodityOptionSurfaceStripper(
    const Handle<PriceTermStructure>& priceCurve,
    const Handle<YieldTermStructure>& discountCurve,
    const boost::shared_ptr<OptionInterpolatorBase>& callSurface,
    const boost::shared_ptr<OptionInterpolatorBase>& putSurface,
    const Calendar& calendar,
    const DayCounter& dayCounter,
    Exercise::Type type,
    bool lowerStrikeConstExtrap,
    bool upperStrikeConstExtrap,
    bool timeFlatExtrapolation,
    bool preferOutOfTheMoney,
    Solver1DOptions solverOptions)
    : OptionSurfaceStripper(callSurface, putSurface, calendar, dayCounter, type, lowerStrikeConstExtrap,
        upperStrikeConstExtrap, timeFlatExtrapolation, preferOutOfTheMoney, solverOptions),
        priceCurve_(priceCurve), discountCurve_(discountCurve) {
    registerWith(priceCurve_);
    registerWith(discountCurve_);
}

boost::shared_ptr<GeneralizedBlackScholesProcess> CommodityOptionSurfaceStripper::process(
    const boost::shared_ptr<QuantLib::SimpleQuote>& volatilityQuote) const {

    QL_REQUIRE(!priceCurve_.empty(), "CommodityOptionSurfaceStripper: price curve is empty");
    QL_REQUIRE(!discountCurve_.empty(), "CommodityOptionSurfaceStripper: discount curve is empty");

    // Volatility term structure for the process
    Handle<BlackVolTermStructure> vts(boost::make_shared<BlackConstantVol>(
        callSurface_->referenceDate(), calendar_, Handle<Quote>(volatilityQuote), dayCounter_));

    // Generate "spot" and "yield" curve for the process.
    Handle<Quote> spot(boost::make_shared<DerivedPriceQuote>(priceCurve_));
    Handle<YieldTermStructure> yield(boost::make_shared<PriceTermStructureAdapter>(*priceCurve_, *discountCurve_));
    yield->enableExtrapolation();

    return boost::make_shared<QuantLib::GeneralizedBlackScholesProcess>(spot, yield, discountCurve_, vts);
}

Real CommodityOptionSurfaceStripper::forward(const Date& date) const {
    QL_REQUIRE(!priceCurve_.empty(), "CommodityOptionSurfaceStripper: price curve is empty");
    return priceCurve_->price(date);
}

}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_IS_SIMD_LOGICAL_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_IS_SIMD_LOGICAL_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-reduction
    Function object implementing is_simd_logical capabilities

    returns @ref True if each vector elements have all bits or no bit set.

    @par Semantic:

    For every parameter of type T

    @code
    logical<scalar_of_t<T>> r = is_simd_logical(x);
    @endcode

    is similar to:

    @code
    T r = all(x == Allbits || x == Zero);
    @endcode

    @see mask2logical

  **/
  const boost::dispatch::functor<tag::is_simd_logical_> is_simd_logical = {};
} }
#endif

#include <boost/simd/function/scalar/is_simd_logical.hpp>
#include <boost/simd/function/simd/is_simd_logical.hpp>

#endif


#ifndef BASISBOOTSTRAPDRIVER_HPP
#define BASISBOOTSTRAPDRIVER_HPP

#include "BALPSolverInterface.hpp"
#include "BAVector.hpp"
#include <boost/scoped_ptr.hpp>

class subsetWrapper;

// serial-only code for now
template<typename BALPSolver, typename RecourseSolver>
class basisBootstrapDriver {
public:

	basisBootstrapDriver(stochasticInput &input, BAContext& ctx) : input(input), ctx(ctx) {
		masterIsValid = false;
		feasible = false;
		int nscen = input.nScenarios();
		rowStates2.resize(nscen);
		colStates2.resize(nscen);
	}

	// start from scratch
	void initializeMaster(std::vector<int> const& scenariosInMaster) {
		int nscen = input.nScenarios();
		this->scenariosInMaster = scenariosInMaster;
		scenariosNotInMaster.clear();
		std::vector<bool> in(nscen,false);
		for (unsigned i = 0; i < scenariosInMaster.size(); i++) {
			in[scenariosInMaster[i]] = true;
		}
		for (int i = 0; i < nscen; i++) {
			if (in[i] == false) scenariosNotInMaster.push_back(i);
		}
		firstSolve = true;
	}
	void solveMaster() {
		subsetWrapper wrapper(input, scenariosInMaster);
		ctx.initializeAssignment(scenariosInMaster.size());
		solver.reset(new BALPSolver(wrapper, ctx, feasible ? BALPSolver::usePrimal : BALPSolver::useDual));

		int nvar1 = input.nFirstStageVars();
		int ncons1 = input.nFirstStageCons();
		int nscen = input.nScenarios();

		if (!firstSolve) { // copy states
			for (int i = 0; i < nvar1; i++) {
				solver->setFirstStageColState(i,colStates1[i]);
			}
			for (int i = 0; i < ncons1; i++) {
				solver->setFirstStageRowState(i,rowStates1[i]);
			}
			for (unsigned k = 0; k < scenariosInMaster.size(); k++) {
				int scen = scenariosInMaster[k];
				int nvar2 = input.nSecondStageVars(scen);
				int ncons2 = input.nSecondStageCons(scen);
				for (int i = 0; i < nvar2; i++) {
					solver->setSecondStageColState(k,i,colStates2[scen][i]);
				}
				for (int i = 0; i < ncons2; i++) {
					solver->setSecondStageRowState(k,i,rowStates2[scen][i]);
				}
			}
			solver->commitStates();
		}
		solver->go();

		colStates1.resize(nvar1); // won't cause reallocation if already done
		rowStates1.resize(ncons1);
		for (int i = 0; i < nvar1; i++) {
			colStates1[i] = solver->getFirstStageColState(i);
		}
		for (int i = 0; i < ncons1; i++) {
			rowStates1[i] = solver->getFirstStageRowState(i);
		}
		for (int k = 0; k < scenariosInMaster.size(); k++) {
			int scen = scenariosInMaster[k];
			int nvar2 = input.nSecondStageVars(scen);
			int ncons2 = input.nSecondStageCons(scen);
			colStates2[scen].resize(nvar2);
			rowStates2[scen].resize(ncons2);
			for (int i = 0; i < nvar2; i++) {
				colStates2[scen][i] = solver->getSecondStageColState(k,i);
			}
			for (int i = 0; i < ncons2; i++) {
				rowStates2[scen][i] = solver->getSecondStageRowState(k,i);
			}
		}

		firstSolve = false;
		masterIsValid = true;
		feasible = true;
	}
	std::vector<double> getFirstStageSolution() const { assert(solver.get()); return solver->getFirstStagePrimalColSolution(); }

	void addScenarioToMaster(int scen) {
		RecourseSolver rsolver(input, scen, getFirstStageSolution());

		// may need to change this
		rsolver.setDualObjectiveLimit(1e10);

		int nvar2 = input.nSecondStageVars(scen);
		int ncons2 = input.nSecondStageCons(scen);
		if (goodRowStates.size() && input.onlyBoundsVary()) {
			for (int i = 0; i < nvar2; i++) {
				rsolver.setSecondStageColState(i,goodColStates[i]);
			}
			for (int i = 0; i < ncons2; i++) {
				rsolver.setSecondStageRowState(i,goodRowStates[i]);
			}
			rsolver.commitStates();
		}
		rsolver.go();

		std::vector<variableState> cols(nvar2), rows(ncons2);
		for (int i = 0; i < nvar2; i++) {
			cols[i] = rsolver.getSecondStageColState(i);
		}
		for (int i = 0; i < ncons2; i++) {
			rows[i] = rsolver.getSecondStageRowState(i);
		}

		addScenarioToMaster(scen, cols, rows, rsolver.getStatus() == Optimal);

	}

	void addScenarioToMaster(int scen, std::vector<variableState> const& colstates,
				 std::vector<variableState> const& rowstates, bool isFeasible) {
		feasible = feasible && isFeasible;
		assert(feasible);
		if (feasible) {
			if (input.onlyBoundsVary()) {
				goodRowStates = rowstates;
				goodColStates = colstates;
			}
			rowStates2[scen] = rowstates;
			colStates2[scen] = colstates;
		} else {
			if (input.onlyBoundsVary()) {
				// reuse a good basis
				rowStates2[scen] = goodRowStates;
				colStates2[scen] = goodColStates;
			} else {
				// slack basis
				colStates2[scen].clear();
				colStates2[scen].resize(input.nSecondStageVars(scen),AtLower);
				rowStates2[scen].clear();
				rowStates2[scen].resize(input.nSecondStageCons(scen),Basic);
			}
		}
		scenariosInMaster.push_back(scen);
		scenariosNotInMaster.erase(std::find(scenariosNotInMaster.begin(),scenariosNotInMaster.end(),scen));
		assert(scenariosInMaster.size() + scenariosNotInMaster.size() == input.nScenarios());
		masterIsValid = false;

	}

	void writeStatus(std::string const &filebase);
	void loadStatus(std::string const &filebase);

	double getObjective() {
		assert(masterIsValid && scenariosInMaster.size() == input.nScenarios());
		return solver->getObjective();
	}

	int countSquare() {
		assert(masterIsValid);
		int count = 0;
		for (unsigned k = 0; k < scenariosInMaster.size(); k++) {
			int scen = scenariosInMaster[k];
			int nvar2 = input.nSecondStageVars(scen);
			int ncons2 = input.nSecondStageCons(scen);
			int nbasic = 0;
			for (int i = 0; i < nvar2; i++) {
				if (colStates2[scen][i] == Basic) nbasic++;
			}
			for (int i = 0; i < ncons2; i++) {
				if (rowStates2[scen][i] == Basic) nbasic++;
			}
			if (nbasic == ncons2) count++;
		}
		printf("%d of %d are square\n",count,(int)scenariosInMaster.size());
		return count;
	}

protected:
	std::vector<variableState> rowStates1, colStates1;
	std::vector<variableState> goodRowStates, goodColStates; // basis from a feasible recourse problem, for hot starting
	std::vector<std::vector<variableState> > rowStates2, colStates2;
	std::vector<int> scenariosInMaster, scenariosNotInMaster;
	BAContext &ctx;
	stochasticInput &input;
	bool masterIsValid; // false if we've added scenarios and not reformulated/resolved master
	bool feasible; // if all scenarios are feasible (so we can use primal simplex)
	bool firstSolve; // we haven't solved a master problem before

	boost::scoped_ptr<BALPSolver> solver;

};

// wrap stochasticInput so that only a subset of scenarios are presented
class subsetWrapper : public stochasticInput {
public:
        subsetWrapper(stochasticInput &inner, std::vector<int> const& includedScenarios) :
		realScenarios(includedScenarios), probabilitySum(0.), inner(inner) {
			for (unsigned i = 0; i < realScenarios.size(); i++) {
				probabilitySum += inner.scenarioProbability(realScenarios[i]);
			}
	}
	virtual int nScenarios() { return realScenarios.size(); }
        virtual int nFirstStageVars() { return inner.nFirstStageVars(); }
        virtual int nFirstStageCons() { return inner.nFirstStageCons(); }
        virtual int nSecondStageVars(int scen) { return inner.nSecondStageVars(realScenarios[scen]); }
        virtual int nSecondStageCons(int scen) { return inner.nSecondStageCons(realScenarios[scen]); }

        virtual std::vector<double> getFirstStageColLB() { return inner.getFirstStageColLB(); }
        virtual std::vector<double> getFirstStageColUB() { return inner.getFirstStageColUB(); }
        virtual std::vector<double> getFirstStageObj() { return inner.getFirstStageObj(); }
        virtual std::vector<std::string> getFirstStageColNames() { return inner.getFirstStageColNames(); }
        virtual std::vector<double> getFirstStageRowLB() { return inner.getFirstStageRowLB(); }
        virtual std::vector<double> getFirstStageRowUB() { return inner.getFirstStageRowUB(); }
        virtual std::vector<std::string> getFirstStageRowNames() { return inner.getFirstStageRowNames(); }
        virtual bool isFirstStageColInteger(int col) { return inner.isFirstStageColInteger(col); }

        virtual std::vector<double> getSecondStageColLB(int scen) { return inner.getSecondStageColLB(realScenarios[scen]); }
        virtual std::vector<double> getSecondStageColUB(int scen) { return inner.getSecondStageColUB(realScenarios[scen]); }
        virtual std::vector<double> getSecondStageObj(int scen) {
		std::vector<double> v = inner.getSecondStageObj(realScenarios[scen]);
		for (unsigned i = 0; i < v.size(); i++) v[i] /= probabilitySum;
		return v;
	}
        virtual std::vector<std::string> getSecondStageColNames(int scen) { return inner.getSecondStageColNames(realScenarios[scen]); }
        virtual std::vector<double> getSecondStageRowUB(int scen) { return inner.getSecondStageRowUB(realScenarios[scen]); }
        virtual std::vector<double> getSecondStageRowLB(int scen) { return inner.getSecondStageRowLB(realScenarios[scen]); }
        virtual std::vector<std::string> getSecondStageRowNames(int scen) { return inner.getSecondStageRowNames(realScenarios[scen]); }
        virtual double scenarioProbability(int scen) { return inner.scenarioProbability(realScenarios[scen])/probabilitySum; }
        virtual bool isSecondStageColInteger(int scen, int col) { return inner.isSecondStageColInteger(realScenarios[scen],col); }

        virtual CoinPackedMatrix getFirstStageConstraints() { return inner.getFirstStageConstraints(); }
        virtual CoinPackedMatrix getSecondStageConstraints(int scen) { return inner.getSecondStageConstraints(realScenarios[scen]); }
        virtual CoinPackedMatrix getLinkingConstraints(int scen) { return inner.getLinkingConstraints(realScenarios[scen]); }

        virtual bool scenarioDimensionsEqual() { return inner.scenarioDimensionsEqual(); }
        virtual bool onlyBoundsVary() { return inner.onlyBoundsVary(); }
        virtual bool allProbabilitiesEqual() { return inner.allProbabilitiesEqual(); }
        virtual bool continuousRecourse() { return inner.continuousRecourse(); }

private:
	std::vector<int> const realScenarios;
	double probabilitySum;
	stochasticInput &inner;
};

// select scenarios by number only
template<typename BALPSolver, typename RecourseSolver>
class simpleBasisBootstrapDriver : public basisBootstrapDriver<BALPSolver,RecourseSolver> {
public:
	simpleBasisBootstrapDriver(stochasticInput &input, BAContext& ctx) : basisBootstrapDriver<BALPSolver,RecourseSolver>(input,ctx) {}

	void goFromScratch(int startingScenarios, int addedPer) {
		std::vector<int> s;
		for (int i = 0; i < startingScenarios; i++) s.push_back(i);
		this->initializeMaster(s);
		this->solveMaster();

		while(this->scenariosNotInMaster.size()) {
			for (int i = 0; i < addedPer && this->scenariosNotInMaster.size(); i++) {
				this->addScenarioToMaster(this->scenariosNotInMaster[0]);
			}
			this->solveMaster();
		}

	}

};
#endif