// Copyright 2021 National Technology & Engineering Solutions of Sandia, LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights in this software. See LICENSE.txt for details.
//
// Author: Denis Bueno

#if defined(__linux__) && !defined(_GNU_SOURCE)
#define _GNU_SOURCE 1
#endif

#include "supp/boolector_supp.h"

#include <bitset>
#include <boost/iterator/zip_iterator.hpp>
#include <gmpxx.h>
#include <iostream>
#include <llvm/Support/Casting.h>

#include "supp/std_supp.h"
#include "supp/wrapstream.h"

using namespace std;
using namespace llvm;

static inline int log2pow2(int size) {
  switch (size) {
    case 1: return 0;
    case 2: return 1;
    case 4: return 2;
    case 8: return 3;
    case 16: return 4;
    case 32: return 5;
    case 64: return 6;
    case 128: return 7;
  }
  EUFORIA_FATAL("size must be power of two < 256: got " + to_string(size));
}

/*-----------------------------------------------------------------------------------*/

namespace boolector {

std::size_t hash_value(const BoolectorNodeWrapper& n) {
  return n.hash();
}

std::ostream& operator<<(std::ostream& os, const BoolectorNodeWrapper& n) {
  auto *file = wrapostream(os);
  boolector_dump_smt2_node(boolector_get_btor(n.c_node_), file, n.c_node_);
  fclose(file);
  return os;
}

BoolectorNodeWrapper BoolectorNodeWrapper::operator()(const BoolectorNodeWrapper& a) {
  vector<BoolectorNodeWrapper> args({a});
  return boolector::apply(args, *this);
}

BoolectorNodeWrapper BoolectorNodeWrapper::operator()(const std::vector<BoolectorNodeWrapper>& args) {
  return boolector::apply(args, *this);
}

/*-----------------------------------------------------------------------------------*/
/* btor methods */

BtorWrapper::BtorWrapper() : c_btor_(boolector_new()), trfp(nullptr) {}

BtorWrapper::~BtorWrapper() {
  boolector_release_all(c_btor_);
  boolector_delete(c_btor_);
  if (trfp) fclose(trfp);
}

BtorWrapper::BtorWrapper(const BtorWrapper& other)
  : c_btor_(boolector_clone(other.c_btor_)),
    trfp(boolector_get_trapi(other.c_btor_)) {
  for (auto& [name, node] : other.bitvecs) {
    bitvecs.insert({name, matchNode(node)});
  }
  for (auto& [name, node] : other.arrays) {
    arrays.insert({name, matchNode(node)});
  }
  for (auto& a : other.assumptions) {
    assumptions.push_back(matchNode(a));
  }
  for (auto& [name, node] : other.params) {
    params.insert({name, node});
  }
  for (auto& f : other.funs) {
    funs.push_back(f);
  }
}

void BtorWrapper::setOption(BtorOption o, uint32_t val) {
  boolector_set_opt(c_btor_, o, val);
}

uint32_t BtorWrapper::getOption(BtorOption o) const {
  return boolector_get_opt(c_btor_, o);
}

void BtorWrapper::PrintStats() const {
  boolector_print_stats(c_btor_);
}

void BtorWrapper::ResetStats() {
  boolector_reset_stats(c_btor_);
}

 BoolectorNodeWrapper BtorWrapper::zero(int width) {
  return Wrap(boolector_zero(c_btor_, boolector_bitvec_sort(c_btor_, width)));
}

 BoolectorNodeWrapper BtorWrapper::bconst(const char *bits) {
  return Wrap(boolector_const(c_btor_, bits));
}

BoolectorNodeWrapper BtorWrapper::bconst(const std::string& bits) {
  return Wrap(boolector_const(c_btor_, bits.c_str()));
}

BoolectorNodeWrapper BtorWrapper::boolVal(bool b) {
  auto ret = Wrap(b ? boolector_true(c_btor_) : boolector_false(c_btor_));
  return ret;
}

BoolectorNodeWrapper BtorWrapper::cint(int i, int width) {
  auto sort = width == 1 ? boolector_bool_sort(c_btor_) : boolector_bitvec_sort(c_btor_, width);
  return Wrap(boolector_int(c_btor_, i, sort));
}

BoolectorNodeWrapper BtorWrapper::var(int width, const std::string& name) {
  if (width < 1) EUFORIA_FATAL("width is not at least 1");
  if (auto loc = bitvecs.find(name); loc != bitvecs.end()) {
    assert(name == loc->second.symbol());
    assert(static_cast<unsigned>(width) == loc->second.width());
    return loc->second;
  }

  auto sort = boolector_bitvec_sort(c_btor_, width);
  auto ret = Wrap(boolector_var(c_btor_, sort, name.c_str()));
  bitvecs.insert({name, ret});
  return ret;
}

BoolectorNodeWrapper BtorWrapper::boolvar(const std::string& name) {
  if (auto search = bools_.find(name); search != bools_.end()) {
    assert(name == search->second.symbol());
    return search->second;
  }
  auto sort = boolector_bool_sort(c_btor_);
  auto ret = Wrap(boolector_var(c_btor_, sort, name.c_str()));
  bools_.insert({name, ret});
  return ret;
}

BoolectorNodeWrapper BtorWrapper::uf(BoolectorSort s, const string& name) {
  return Wrap(boolector_uf(c_btor_, s, name.c_str()));
}

BoolectorNodeWrapper BtorWrapper::param(BoolectorSort sort, const std::string& name) {
  if (auto loc = params.find(name); loc != params.end()) {
    assert(boolector_is_param(c_btor_, loc->second.c_node_));
    assert(name == loc->second.symbol());
    return loc->second;
  }

  auto ret = Wrap(boolector_param(c_btor_, sort, name.c_str()));
  params.insert({name, ret});
  return ret;
}

BoolectorNodeWrapper write(const BoolectorNodeWrapper& arr, const BoolectorNodeWrapper& c_node_, const BoolectorNodeWrapper& val) {
  assert(arr.btor() == c_node_.btor() && c_node_.btor() == val.btor());
//    if (boolector_is_fun_sort(arr.btor(), arr.sort())) {
//      // write is a new fun
//      auto name = "ptr_in" + to_string(++paramCtr);
//      auto ptr_in = BoolectorNodeWrapper(boolector_param(arr.btor(), boolector_fun_get_domain_sort(arr.btor(), arr.c_node_), name.c_str()));
//      vector<BoolectorNodeWrapper> recArg({ptr_in});
//      auto body = ite(ptr_in.eq(c_node_), val, boolector::apply(recArg, arr));
//      BoolectorNode **param = new BoolectorNode*[1];
//      param[0] = ptr_in.c_node_;
//      auto ret = BoolectorNodeWrapper(boolector_fun(arr.btor(), param, 1, body.c_node_));
//      delete[] param;
//      return ret;
//    } else {
    // use boolector_write
    return BoolectorNodeWrapper(arr.btor(), boolector_write(arr.btor(), arr.c_node_, c_node_.c_node_, val.c_node_));
//    }
}

BoolectorNodeWrapper read(const BoolectorNodeWrapper& arr, const BoolectorNodeWrapper& c_node_) {
  assert(arr.btor() == c_node_.btor());
//    if (boolector_is_fun_sort(arr.btor(), arr.sort())) {
//      // read is an apply
//      BoolectorNode **args = new BoolectorNode*[1];
//      args[0] = c_node_.c_node_;
//      auto ret = BoolectorNodeWrapper(boolector_apply(arr.btor(), args, 1, arr.c_node_));
//      delete[] args;
//      return ret;
//    } else {
    return BoolectorNodeWrapper(arr.btor(), boolector_read(arr.btor(), arr.c_node_, c_node_.c_node_));
//    }
}

BoolectorNodeWrapper BtorWrapper::array(const std::string& name, int ptrSize, int valWidth) {
  if (auto loc = arrays.find(name); loc != arrays.end()) {
    assert(name == loc->second.symbol());
    assert(boolector_is_array(loc->second.btor(), loc->second.c_node_));
    return loc->second;
  }

  auto sort = boolector_array_sort(c_btor_, boolector_bitvec_sort(c_btor_, ptrSize), boolector_bitvec_sort(c_btor_, valWidth));
  auto ret = Wrap(boolector_array(c_btor_, sort, name.c_str()));
  arrays.insert({name, ret});
  return ret;
}

BoolectorNodeWrapper BtorWrapper::const_array(
    int ptr_size, const BoolectorNodeWrapper& value) {
  int val_size = static_cast<int>(value.width());
  auto sort = boolector_array_sort(c_btor_,
                                   boolector_bitvec_sort(c_btor_, ptr_size),
                                   boolector_bitvec_sort(c_btor_, val_size));
  auto ret = Wrap(boolector_const_array(c_btor_, sort, value.c_node_));
  return ret;
}

BoolectorNodeWrapper BtorWrapper::fun(const std::vector<BoolectorNodeWrapper>& funParams, BoolectorNodeWrapper body) {
  int paramc = static_cast<int>(funParams.size());
  BoolectorNode **paramnodes = new BoolectorNode*[funParams.size()];
  for (int i = 0; i < paramc; i++) {
    paramnodes[i] = funParams[i].c_node_;
  }
  auto ret = Wrap(boolector_fun(c_btor_, paramnodes, paramc, body.c_node_));
  delete[] paramnodes;
  funs.push_back(ret);
  return ret;
}

BoolectorSort BtorWrapper::fun_sort(BoolectorSort *args, uint32_t arity,
                                    BoolectorSort ret) {
  return boolector_fun_sort(c_btor_, args, arity, ret);
}

BoolectorSort BtorWrapper::array_sort(BoolectorSort index,
                                      BoolectorSort element) {
  return boolector_array_sort(c_btor_, index, element);
}

BoolectorSort BtorWrapper::bitvec_sort(uint32_t width) {
  return boolector_bitvec_sort(c_btor_, width);
}

BoolectorSort BtorWrapper::bool_sort() {
  return boolector_bool_sort(c_btor_);
}

BoolectorNodeWrapper apply(const std::vector<BoolectorNodeWrapper>& args, BoolectorNodeWrapper f) {
  int argc = static_cast<int>(args.size());
  BoolectorNode **argnodes = new BoolectorNode*[args.size()];
  for (int i = 0; i < argc; i++) {
    argnodes[i] = args[i].c_node_;
  }
  auto ret = BoolectorNodeWrapper(f.btor(), boolector_apply(f.btor(), argnodes, argc, f.c_node_));
  delete[] argnodes;
  return ret;
}

void BtorWrapper::fixateAssumptions() {
  boolector_fixate_assumptions(c_btor_);
}

void BtorWrapper::resetAssumptions() {
  boolector_reset_assumptions(c_btor_);
}

BtorWrapper::result BtorWrapper::simplify() {
  ScopedTimeKeeper t(&total_simplify_time_);
  auto result = boolector_simplify(c_btor_);
  switch (result) {
    case BOOLECTOR_SAT: return result::kSat;
    case BOOLECTOR_UNSAT: return result::kUnsat;
    case BOOLECTOR_UNKNOWN: return result::kUnknown;
    default:
      EUFORIA_FATAL("bad result from boolector");
  }
}

BtorWrapper::result BtorWrapper::check() {
  ScopedTimeKeeper t(&total_check_time_);
  auto result = boolector_sat(c_btor_);
  assumptions.clear();
  switch (result) {
    case BOOLECTOR_SAT: return result::kSat;
    case BOOLECTOR_UNSAT: return result::kUnsat;
    default: EUFORIA_FATAL("unhandled switch");
  }
}

shared_ptr<BtorAssignment> BtorWrapper::assignment(const BoolectorNodeWrapper& n) {
  assert(n.btor() == c_btor_);
  shared_ptr<BtorAssignment> ret;
  if (n.isArray()) {
    ret = make_shared<BtorArrayAssignment>(*this, n);
  } else {
    assert(n.isBV());
    ret = make_shared<BtorBvAssignment>(*this, n);
  }
  return ret;
}

BoolectorNodeWrapper BtorWrapper::assignmentNode(const BtorAssignment& a) {
  switch (a.kind) {
    case BtorAssignment::kind::BV: {
      auto& val = static_cast<const BtorBvAssignment&>(a);
      return bconst(val.value());
    }

    case BtorAssignment::kind::Array: {
      auto& arrayAss = static_cast<const BtorArrayAssignment&>(a);
      auto arr = array(arrayAss.symbol, arrayAss.ptrSize, arrayAss.valSize);
      for (int i = 0; i < arrayAss.size(); i++) {
        auto& c_node_ = arrayAss.indices()[i];
        auto& val = arrayAss.values()[i];
        auto nptr = bconst(c_node_);
        auto nval = bconst(val);
        arr = write(arr, nptr, nval);
      }
      return arr;
    }

    default:
      EUFORIA_FATAL("more imples");
  }

}

BoolectorNodeWrapper BtorWrapper::matchNode(const BoolectorNodeWrapper& n) const {
  return Wrap(boolector_match_node(c_btor_, n.c_node_));
}

BtorWrapper::result BtorWrapper::checkSMT2(const std::string& filename) {
  auto *fp = fopen(filename.c_str(), "r");
  assert(nullptr != fp);
  char *error_msg = nullptr;
  int status = -1;
  auto *outfile = fopen("/tmp/btor_checkSMT2_outfile", "w");
  assert(nullptr != outfile);
  switch (boolector_parse_smt2(c_btor_, fp, filename.c_str(), outfile, &error_msg, &status)) {
    case BOOLECTOR_SAT:
      return result::kSat;
    case BOOLECTOR_UNSAT:
      return result::kUnsat;
    case BOOLECTOR_UNKNOWN:
      return result::kUnknown;
    case BOOLECTOR_PARSE_ERROR:
      EUFORIA_FATAL("checkSMT2: boolector parse error");
  }
  EUFORIA_FATAL("unexpected");
}

void BtorWrapper::setTrace(const std::string& filename) {
  trfp = fopen(filename.c_str(), "w");
  assert(trfp);
  boolector_set_trapi(c_btor_, trfp);
}

void BtorWrapper::printModel() const {
  char fmt[5] = "smt2"; // why is this not const char*
  boolector_print_model(c_btor_, fmt, stderr);
}

void BtorWrapper::DumpBenchmark(std::ostream& os) const {
  auto *file = wrapostream(os);
  boolector_dump_smt2(c_btor_, file);
  fclose(file);
}

void BtorWrapper::Print(std::ostream& os) const {
  auto *file = wrapostream(os);
  boolector_dump_smt2(c_btor_, file);
  fclose(file);
}

void BtorWrapper::printAssertions() const {
  //    boolector_dump_smt2(c_btor_, wrapostream(cerr));
  //    boolector_dump_smt2(c_btor_, wraplogger(log));
  stringstream ss;
  boolector_dump_smt2(c_btor_, wrapstringstream(&ss));
  logger.Log(4, "boolector assertions:\n{}", ss.str());
  //log->debug("boolector has {} assertions:", assertions.size());
  //size_t num = 0;
  //for (auto assertion : assertions)
  //  log->debug("assertion {}: {}\n", ++num, assertion);
  //log->debug("boolector has {} assumptions:", assumptions.size());
  //num = 0;
  //for (auto assump : assumptions) {
  //  log->debug("assumption {}: {}\n", ++num, assump);
  //}
}

void BtorWrapper::add(BoolectorNode *n) {
  boolector_assert(c_btor_, n);
}

void BtorWrapper::assume(BoolectorNode *n) {
  assumptions.push_back(Wrap(n));
  boolector_assume(c_btor_, n);
}

/*-----------------------------------------------------------------------------------*/
// assignments

BtorArrayAssignment::BtorArrayAssignment(BtorWrapper& b, const BoolectorNodeWrapper& n)
  : BtorAssignment(BtorAssignment::kind::Array, b,
                   boolector_get_symbol(n.btor(), n.c_node_)),
    ptrSize(boolector_get_index_width(n.btor(), n.c_node_)),
    valSize(boolector_get_width(n.btor(), n.c_node_)) {
  assert(b.c_btor_ == n.btor());
  char **i;
  char **v;
  uint32_t size;
  boolector_array_assignment(b.c_btor_, n.c_node_, &i, &v, &size);
  if (size == 1 && i[0][0] == '*') {
    // constant array
    string val(v[0]);
    assert(val.size() == valSize);
    const_value_ = val;
    boolector_free_array_assignment(b.c_btor_, i, v, size);
    return;
  }

  indices_.reserve(size);
  values_.reserve(size);
  for (uint32_t j = 0; j < size; j++) {
    string idx(i[j]);
    string val(v[j]);
    assert(idx.size() == ptrSize);
    assert(val.size() == valSize);
    indices_.emplace_back(idx);
    values_.emplace_back(val);
  }
  // Fills the vector idx_val_pairs_ with the corresponding pairs from indices_
  // and values.
  idx_val_pairs_.reserve(size);
  using IdxValTuple = boost::tuple<string&,string&>;
  std::for_each(
      boost::make_zip_iterator(
          boost::make_tuple(indices_.begin(), values_.begin())),
      boost::make_zip_iterator(
          boost::make_tuple(indices_.end(), values_.end())),
      [&](const IdxValTuple& t) {
          idx_val_pairs_.emplace_back(std::make_pair(boost::get<0>(t),
                                                     boost::get<1>(t)));
      });
  assert(idx_val_pairs_.size() == size);
  struct {
    bool operator()(const pair<string,string>& iv0,
                    const pair<string,string>& iv1) {
      return std::get<0>(iv0) < std::get<0>(iv1);
    }
  } compare_on_indices;
  // Sorts the array of pairs so that lambdas are equivalent if they're
  // identical (see AsExpr).
  std::sort(idx_val_pairs_.begin(), idx_val_pairs_.end(), compare_on_indices);
  if (size)
    boolector_free_array_assignment(b.c_btor_, i, v, size);
}

BtorArrayAssignment::~BtorArrayAssignment() = default;

static z3::expr bvValExpr(z3::context& ctx, const std::string& ass) {
  int width = static_cast<int>(ass.size());
  bool bits[width];
  for (int i = 0, j = width-1; i < width; i++, j--) {
    bool bit = bool(ass[i] - '0');
    bits[j] = bit;
  }
  return ctx.bv_val(width, bits);
}

z3::expr BtorArrayAssignment::AsExpr(const z3::expr& arr_in) const {
  if (const_value_) {
    return z3::const_array(arr_in.ctx().bv_sort(ptrSize),
                           bvValExpr(arr_in.ctx(), *const_value_));
  }

  assert(indices_.size() == values_.size());
  // assert(bool(arr_in));
  z3::expr arr(arr_in.ctx());
  unsigned i = 0;
  for (auto&& iv_pair : idx_val_pairs_) {
    auto idx = bvValExpr(arr.ctx(), iv_pair.first);
    auto val = bvValExpr(arr.ctx(), iv_pair.second);
    auto var = arr.ctx().constant(("x" + to_string(i++)).c_str(), idx.get_sort());
    if (!bool(arr)) {
      arr = z3::lambda(var,
                       z3::ite(var == idx, val,
                               arr.ctx().bv_val(0, val.get_sort().bv_size())));
    } else {
      arr = z3::lambda(var, z3::ite(var == idx, val, z3::select(arr, idx)));
    }
  }
  return arr;
}

z3::expr BtorArrayAssignment::AsConstraint(const z3::expr& var) const {
  if (const_value_) {
    return var == AsExpr(var);
  }
  assert(indices_.size() == values_.size());
  z3::expr_vector v(var.ctx());
  ExprVectorInserter out(v);
  for (unsigned i = 0; i < indices_.size(); i++) {
    auto idx = bvValExpr(var.ctx(), indices_[i]);
    auto val = bvValExpr(var.ctx(), values_[i]);
    *out++ = z3::select(var, idx) == val;
  }
  return expr_mk_and(v);
}

static bool bitStringToVal(const std::string& index, uint64_t &val) {
  if (index.size() <= 64) {
    val = static_cast<uint64_t>(strtoll(index.c_str(), nullptr, 2));
    return true;
  }
  return false;
}

void BtorArrayAssignment::print(std::ostream& os) const {
  assert(indices_.size() == values_.size());
  if (symbol) {
    os << symbol;
  }
  os << " btor array: [index] = value" << endl;
  if (const_value_) {
    os << "[*] = " << *const_value_ << ";" << std::endl;
  } else {
    for (unsigned i = 0; i < indices_.size(); i++) {
      auto& idx = indices_[i];
      auto& val = values_[i];
      os << "[" << idx << "] = " << val << ";" << std::endl;
    }
    os << "[ ... else ... ] = 0";
  }
}

z3::expr BtorBvAssignment::AsExpr(const z3::expr& e) const {
  if (e.is_bool()) {
    return e.ctx().bool_val(ass[0] == '1');
  }
  return bvValExpr(e.ctx(), ass);
}

z3::expr BtorBvAssignment::AsConstraint(const z3::expr& var) const {
  auto expr = AsExpr(var);
  return expr == var;
}

void BtorBvAssignment::print(ostream& os) const {
  if (symbol && width == 1) {
    if (strcmp(ass, "0") == 0) {
      os << "!";
    }
    os << symbol;
  } else {
    if (symbol)
      os << "<" << width << ">" << symbol << " = " << ass;
    else
      os << "<" << width << "> = " << ass;
  }
}

std::ostream& operator<<(std::ostream& os, const BtorAssignment& a) {
  a.print(os);
  return os;
}

/*-----------------------------------------------------------------------------------*/
// boolector rewriter

Z3ToBtorRewriter::Z3ToBtorRewriter(std::shared_ptr<BtorWrapper> b) : B(b) {}

BoolectorNodeWrapper Z3ToBtorRewriter::visitExpr(const z3::expr& n) {
  Z3_decl_kind k;
  // function that does the shift
  function<BoolectorNodeWrapper(const BoolectorNodeWrapper&,
                                const BoolectorNodeWrapper&)> shift =
      [&](const BoolectorNodeWrapper& n,
          const BoolectorNodeWrapper& shift_amount) {
        if (k == Z3_OP_BSHL) { return n.shl(shift_amount); }
        else if (k == Z3_OP_BASHR) { return n.ashr(shift_amount); }
        else if (k == Z3_OP_BLSHR) { return n.lshr(shift_amount); }
        else { EUFORIA_FATAL("unhandled kind: {}", k); }
      };
  // function that does sign/zero extension
  function<BoolectorNodeWrapper(const BoolectorNodeWrapper&, int)> extend =
      [&](const BoolectorNodeWrapper& n, int extend_amount) {
        if (k == Z3_OP_BSHL) { return n.uext(extend_amount); }
        else if (k == Z3_OP_BASHR) { return n.sext(extend_amount); }
        else if (k == Z3_OP_BLSHR) { return n.uext(extend_amount); }
        else { EUFORIA_FATAL("unhandled kind: {}", k); }
      };
  switch ((k = n.decl().decl_kind())) {
    case Z3_OP_BSHL:
    case Z3_OP_BASHR:
    case Z3_OP_BLSHR: {
      BoolectorNodeWrapper ret;
      assert(n.num_args() == 2);
      const auto& val = Arg(n,0);
      const auto& arg1 = Arg(n,1);
      const unsigned val_width = val.width();
      bitset<sizeof(val_width)*8> bits(val_width);
      // Boolector only supports shifting bit vectors of power-of-two size, but
      // Z3 allows you to shift any bit width vector. To express non
      // powers-of-two in Boolector, we extend the BV to the nearest
      // power-of-two size, shift it, and truncate the result.
      if (bits.count() != 1) {
        auto exp = ceil(log2(val_width));
        // Get the next highest power of two size
        auto new_width = pow(2, exp);
        logger.Log(7, "exp {}", exp);
        logger.Log(7, "bits {}", bits);
        logger.Log(7, "val {}", val);
        logger.Log(7, "val_width {}", val_width);
        logger.Log(7, "arg1 {}", arg1);
        logger.Log(7, "new_width {}", new_width);
        // Extends the value, shifts it, then chops off the top bits
        auto val_extended = extend(val, new_width - val_width);
        auto shift_amount = arg1.slice(exp-1, 0);
        ret = shift(val_extended, shift_amount);
        ret = ret.slice(val_width-1, 0);
        logger.Log(7, "ret.width() {}", ret.width());
      } else {
        auto hi = log2pow2(val_width)-1;
        const auto& shift_amount = arg1.slice(hi, 0);
        ret = shift(val, shift_amount);
        assert(ret.width() == n.get_sort().bv_size());
      }
      assert(ret.width() == n.get_sort().bv_size());
      return ret;
      break;
    }

    default:
      break;
  }
  std::ostringstream ss;
  ss << n;
  EUFORIA_FATAL(fmt::format("no generic handling: {}", ss.str()));
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitUNINTERPRETED(const z3::expr& e) {
  const auto decl = e.decl();
  const auto name = decl.name().str();
  if (e.is_bool()) {
    // Boolean variable
    if (e.num_args() == 0) {
      return B->boolvar(e.decl().name().str().c_str());
    }
    EUFORIA_FATAL("unhandled Boolean function");
  } else if (e.is_bv()) {
    // Function returning a bit vector
    if (e.num_args() == 2) {
      // GROSS stuff
      // So Z3 turns perfectly good (bvsrem ...) ops into uninterpreted function calls.
      // This is because division by zero is undefined.
      // By the way it ALSO turns (ite (not (= x 0)) (/ y x) 0) into uninterpreted calls.
      // Anyhow, here we use the real boolector functions.
      if (starts_with(name, "bvsrem")) {
        return Arg(e,0).srem(Arg(e,1));
      } else if (starts_with(name, "bvsmod")) {
        return Arg(e,0).srem(Arg(e,1));
      } else if (starts_with(name, "bvsdiv")) {
        return Arg(e,0) / Arg(e,1);
      } else if (starts_with(name, "bvurem")) {
        return Arg(e,0).urem(Arg(e,1));
      } else if (starts_with(name, "bvudiv")) {
        return Arg(e,0).udiv(Arg(e,1));
      }
    }
    // Bit vector variable
    if (decl.is_const()) {
      return B->var(e.get_sort().bv_size(), name);
    } else {
      // Uninterpreted function returning a BV
      vector<BoolectorSort> param_sorts;
      for (unsigned i = 0; i < decl.arity(); i++) {
        param_sorts.push_back(TranslateSort(decl.domain(i)));
      }
      BoolectorSort ret_sort = TranslateSort(decl.range());
      auto uf_sort = B->fun_sort(param_sorts.data(), param_sorts.size(),
                                 ret_sort);
      return B->uf(uf_sort, name);
    }
  } else if (e.is_array()) {
    return B->array(e.decl().name().str(),
                   e.get_sort().array_domain().bv_size(),
                   e.get_sort().array_range().bv_size());
  }
  EUFORIA_FATAL("unimplemented");
}

BoolectorSort Z3ToBtorRewriter::TranslateSort(const z3::sort& s) const {
  switch (s.sort_kind()) {
    case Z3_BOOL_SORT:
      return B->bool_sort();
    case Z3_BV_SORT:
      return B->bitvec_sort(s.bv_size());
    case Z3_ARRAY_SORT: {
      auto index_sort = TranslateSort(s.array_domain());
      auto value_sort = TranslateSort(s.array_range());
      return B->array_sort(index_sort, value_sort);
    }
    default: {
      EUFORIA_FATAL(fmt::format("Boolector doesn't support this sort: {}",
                                s));
    }
  }
}

//  int uniqParam = 0;
BoolectorNodeWrapper Z3ToBtorRewriter::visitCONST_ARRAY(const z3::expr& e) {
  return B->const_array(e.get_sort().array_domain().bv_size(),
                        Arg(e, 0));

  // use a lambda \x. 0 to send all reads to zero
  // this is currently unsupported by boolector :(
//    auto ptrSort = boolector_bitvec_sort(B->c_btor_, sort.array_domain().bv_size());
//    auto x = B->param(ptrSort, "ptr_in" + to_string(++uniqParam));
//    vector<BoolectorNodeWrapper> params({x});
//    auto zero = B->cint(0, sort.array_range().bv_size());
//    return B->fun(params, zero);
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitDISTINCT(const z3::expr& n) {
  BoolectorNodeWrapper ret;
  for (unsigned i = 0; i < n.num_args(); i++) {
    for (unsigned j = i+1; j < n.num_args(); j++) {
      auto ne = Arg(n, i).ne(Arg(n, j));
      ret = !ret.is_null() ? ret && ne : ne;
    }
  }
  assert(!ret.is_null());
  return ret;
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitAND(const z3::expr& n) {
  BoolectorNodeWrapper ret = Arg(n,0);
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret && Arg(n,i);
  }
  return ret;
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitOR(const z3::expr& n) {
  BoolectorNodeWrapper ret = Arg(n,0);
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret || Arg(n,i);
  }
  return ret;
}

#define BTOR_BINARY_HANDLER(KIND, EXPR) \
BoolectorNodeWrapper Z3ToBtorRewriter::visit##KIND(const z3::expr& n) { \
  assert(n.num_args() == 2); \
  BoolectorNodeWrapper ret(EXPR); \
  assert(ret.width() > 0); \
  assert(!n.is_bv() || (static_cast<unsigned>(ret.width()) == n.get_sort().bv_size())); \
  assert(!n.is_bool() || (ret.width() == 1)); \
  return ret; \
}

#define BTOR_TERNARY_HANDLER(KIND, EXPR) \
BoolectorNodeWrapper Z3ToBtorRewriter::visit##KIND(const z3::expr& n) { \
  assert(n.num_args() == 3); \
  BoolectorNodeWrapper ret(EXPR); \
  assert(ret.width() > 0); \
  assert(!n.is_bv() || (static_cast<unsigned>(ret.width()) == n.get_sort().bv_size())); \
  assert(!n.is_bool() || (ret.width() == 1)); \
  return ret; \
}

#define BTOR_RASSOC_HANDLER(KIND, COP) \
BoolectorNodeWrapper Z3ToBtorRewriter::visit##KIND(const z3::expr& n) { \
  BoolectorNodeWrapper ret(Arg(n,0)); \
  for (unsigned i = 1; i < n.num_args(); i++) { \
    ret = COP(ret, Arg(n,i)); \
  } \
  return ret; \
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitTRUE(const z3::expr&) {
  return B->boolVal(true);
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitFALSE(const z3::expr&) {
  return B->boolVal(false);
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitNOT(const z3::expr& n) {
  assert(n.num_args() == 1);
  return !Arg(n,0);
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitBIT2BOOL(const z3::expr& e) {
  assert(Z3_get_decl_num_parameters(Z3_context(e.ctx()), Z3_func_decl(e.decl())) == 1);
  auto bit_idx = Z3_get_decl_int_parameter(Z3_context(e.ctx()), Z3_func_decl(e.decl()), 0);
  auto tgt = Arg(e, 0);
  auto ret = tgt.slice(bit_idx, bit_idx);
  assert(ret.width() == 1);
  return ret;
}

//  BTOR_BINARY_HANDLER(EQ, Arg(n, 0).eq(Arg(n, 1)));
BoolectorNodeWrapper Z3ToBtorRewriter::visitEQ(const z3::expr& n) {
  assert(n.num_args() == 2);

  BoolectorNodeWrapper ret(Arg(n, 0) == (Arg(n, 1)));

  assert(!n.is_bv() || (ret.width() == n.get_sort().bv_size()));
  assert(!n.is_bool() || (ret.width() == 1));
  return ret;
}

BTOR_BINARY_HANDLER(IFF, Arg(n, 0).iff(Arg(n, 1)));
BTOR_TERNARY_HANDLER(ITE, ite(Arg(n, 0), Arg(n, 1), Arg(n, 2)));
BTOR_BINARY_HANDLER(IMPLIES, Arg(n, 0).implies(Arg(n, 1)));
BTOR_BINARY_HANDLER(XOR, Arg(n, 0) ^ Arg(n, 1));

// *** bit vector

BoolectorNodeWrapper Z3ToBtorRewriter::visitBNUM(const z3::expr& n) {
  string numstr;
  auto b = n.is_numeral(numstr);
  assert(b);
  _unused(b);
  mpz_class num(numstr);
  string binstr = num.get_str(2);
  assert(binstr.size() <= n.get_sort().bv_size());
  string bits(n.get_sort().bv_size(), '0');
  copy(binstr.begin(), binstr.end(),
       bits.begin() + (bits.size() - binstr.size()));
  assert(bits.size() == n.get_sort().bv_size());

  //unsigned long long i = 0;
  //b = Z3_get_numeral_uint64(Z3_context(n.ctx()), Z3_ast(n), &i);
  //assert(b == Z3_TRUE);
  //string allBits;
  //for (unsigned j = 0; j < n.get_sort().bv_size(); j++) {
  //  allBits.push_back((i>>j) & 1 ? '1' : '0');
  //}
  //std::reverse(begin(allBits), end(allBits));
  //logger.Log(3, "allbits is {}", allBits);
  //assert(allBits == bits);
  //assert(allBits.size() == n.get_sort().bv_size());
  return B->bconst(bits.c_str());
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitCONCAT(const z3::expr& n) {
  BoolectorNodeWrapper ret = Arg(n,0);
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret.concat(Arg(n,i));
  }
  assert(ret.width() == n.get_sort().bv_size());
  return ret;
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitEXTRACT(const z3::expr& e) {
  // extract in z3 uses "parameters" not "arguments" to specify the hi and lo for extracting
  assert(Z3_get_decl_num_parameters(Z3_context(e.ctx()), Z3_func_decl(e.decl())) == 2);
  auto hi = Z3_get_decl_int_parameter(Z3_context(e.ctx()), Z3_func_decl(e.decl()), 0);
  auto lo = Z3_get_decl_int_parameter(Z3_context(e.ctx()), Z3_func_decl(e.decl()), 1);
  auto tgt = Arg(e, 0);
  auto ret = tgt.slice(hi, lo);
  assert(ret.width() == e.get_sort().bv_size());
  return ret;
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitSIGN_EXT(const z3::expr& e) {
  assert(Z3_get_decl_num_parameters(e.ctx(), e.decl()) == 1);
  auto addedBits = Z3_get_decl_int_parameter((e.ctx()), (e.decl()), 0);
  auto tgt = Arg(e, 0);
  auto ret = tgt.sext(addedBits);
  assert(ret.width() == e.get_sort().bv_size());
  return ret;
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitZERO_EXT(const z3::expr& e) {
  assert(Z3_get_decl_num_parameters(e.ctx(), e.decl()) == 1);
  auto added_bits = Z3_get_decl_int_parameter((e.ctx()), (e.decl()), 0);
  auto tgt = Arg(e, 0);
  auto ret = tgt.uext(added_bits);
  assert(ret.width() == e.get_sort().bv_size());
  return ret;
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitBNOT(const z3::expr& n) {
  assert(n.num_args() == 1);
  BoolectorNodeWrapper ret(Arg(n,0));
  ret = ~ret;
  assert(ret.width() == n.get_sort().bv_size());
  return ret;
}

BoolectorNodeWrapper Z3ToBtorRewriter::visitBNEG(const z3::expr& n) {
  assert(n.num_args() == 1);
  BoolectorNodeWrapper ret(Arg(n,0));
  ret = ret.negate();
  assert(ret.width() == n.get_sort().bv_size());
  return ret;
}

BTOR_BINARY_HANDLER(BSUB, Arg(n, 0) - Arg(n, 1));
BTOR_BINARY_HANDLER(BMUL, Arg(n, 0) * Arg(n, 1));
BTOR_BINARY_HANDLER(BUDIV, Arg(n, 0).udiv(Arg(n, 1)));
BTOR_BINARY_HANDLER(BUDIV_I, Arg(n, 0).udiv(Arg(n, 1)));
BTOR_BINARY_HANDLER(BSDIV, Arg(n, 0) / Arg(n, 1));
BTOR_BINARY_HANDLER(BSDIV_I, Arg(n, 0) / Arg(n, 1));
BTOR_BINARY_HANDLER(BSMOD, Arg(n, 0) % Arg(n, 1));
BTOR_BINARY_HANDLER(BSMOD_I, Arg(n, 0) % Arg(n, 1));
BTOR_BINARY_HANDLER(BSREM, Arg(n, 0).srem(Arg(n, 1)));
BTOR_BINARY_HANDLER(BSREM_I, Arg(n, 0).srem(Arg(n, 1)));
BTOR_BINARY_HANDLER(BUREM, Arg(n, 0).urem(Arg(n, 1)));
BTOR_BINARY_HANDLER(BUREM_I, Arg(n, 0).urem(Arg(n, 1)));
BTOR_BINARY_HANDLER(SLT, Arg(n, 0) < Arg(n, 1));
BTOR_BINARY_HANDLER(SLEQ, Arg(n, 0) <= Arg(n, 1));
BTOR_BINARY_HANDLER(SGT, Arg(n, 0) > Arg(n, 1));
BTOR_BINARY_HANDLER(SGEQ, Arg(n, 0) >= Arg(n, 1));
BTOR_BINARY_HANDLER(ULT, Arg(n, 0).ult(Arg(n, 1)));
BTOR_BINARY_HANDLER(ULEQ, Arg(n, 0).ulte(Arg(n, 1)));
BTOR_BINARY_HANDLER(UGT, Arg(n, 0).ugt(Arg(n, 1)));
BTOR_BINARY_HANDLER(UGEQ, Arg(n, 0).ugte(Arg(n, 1)));
BoolectorNodeWrapper Z3ToBtorRewriter::visitBADD(const z3::expr& n) {
  BoolectorNodeWrapper ret(Arg(n,0));
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret + Arg(n,i);
  }
  assert(!n.is_bv() || (static_cast<unsigned>(ret.width()) == n.get_sort().bv_size()));
  return ret;
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitBXOR(const z3::expr& n) {
  BoolectorNodeWrapper ret(Arg(n,0));
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret ^ Arg(n,i);
  }
  assert(ret.width() == n.get_sort().bv_size());
  return ret;
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitBAND(const z3::expr& n) {
  BoolectorNodeWrapper ret(Arg(n,0));
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret & Arg(n,i);
  }
  assert(ret.width() == n.get_sort().bv_size());
  return ret;
}
BoolectorNodeWrapper Z3ToBtorRewriter::visitBOR(const z3::expr& n) {
  BoolectorNodeWrapper ret(Arg(n,0));
  for (unsigned i = 1; i < n.num_args(); i++) {
    ret = ret | Arg(n,i);
  }
  assert(ret.width() == n.get_sort().bv_size());
  return ret;
}

BTOR_BINARY_HANDLER(SELECT, read(Arg(n,0), Arg(n,1)));
BTOR_TERNARY_HANDLER(STORE, write(Arg(n,0), Arg(n,1), Arg(n,2)));

#undef BTOR_BINARY_HANDLER
#undef BTOR_TERNARY_HANDLER
#undef BTOR_RASSOC_HANDLER

}

void mylog(const boolector::BoolectorNodeWrapper& n) {
  cerr << n;
}


#pragma once

#include <boost/graph/bipartite.hpp>
#include <boost/graph/detail/set_adaptor.hpp>
#include <boost/graph/filtered_graph.hpp>
#include <boost/graph/connected_components.hpp>
#include <boost/graph/biconnected_components.hpp>
#include <boost/property_map/property_map.hpp>
#include <vector>

#include "matroid_graph.hpp"
#include "graph_utils.hpp"

#include <boost/graph/adjacency_list_io.hpp>

namespace tu
{

  /**
   * Finds a matroid element which is parallel to a row element in minor of a given matroid or
   * a unit vector in a column element in the minor.
   *
   * @param matroid The given matroid
   * @param matrix Representation matrix of the given matroid
   * @param minor_height Number of rows of the minor
   * @param minor_width Number of columns of the minor
   * @param row The element must be parallel to this row (along the minor)
   * @return The matroid element that was found to be parallel / unit vector
   */

  template <typename MatroidType, typename MatrixType>
  int find_parallel_to_row(const MatroidType& matroid, const MatrixType& matrix, const size_t minor_height,
      const size_t minor_width, const size_t row)
  {
    /// Look for unit vector
    size_t last = 0;
    size_t count = 0;
    for (size_t c = 0; c < minor_width; ++c)
    {
      if (matrix(row, c) != 0)
      {
        last = c;
        ++count;
      }
    }
    if (count == 1)
      return matroid.name2(last);

    /// Check for parallel
    for (size_t r = 0; r < minor_height; ++r)
    {
      bool same = true;
      for (size_t c = 0; c < minor_width; ++c)
      {
        if (matrix(r, c) != matrix(row, c))
        {
          same = false;
          break;
        }
      }
      if (same)
      {
        return matroid.name1(r);
      }
    }

    throw std::logic_error("find_parallel_to_row did not find any parallel / unit vector!");
  }

  /**
   * A filter for a graph which excludes some edges and a given vertex
   */

  template <typename Graph>
  struct articulation_edge_filter
  {
    typedef typename boost::graph_traits<Graph>::vertex_descriptor vertex_descriptor;
    typedef typename boost::graph_traits<Graph>::edge_descriptor edge_descriptor;
    typedef std::set<edge_descriptor> edge_set;

    /**
     * Constructs the filter.
     */

    articulation_edge_filter() :
      graph_(NULL), articulation_vertex_(NULL), evil_edges_(NULL)
    {
    }

    /**
     * Constructs the filter
     *
     * @param graph The original graph
     * @param articulation_vertex The excluded vertex
     * @param evil_edges The set of excluded edges
     */

    articulation_edge_filter(const Graph* graph, const vertex_descriptor* articulation_vertex,
        const edge_set* evil_edges) :
      graph_(graph), articulation_vertex_(articulation_vertex), evil_edges_(evil_edges)
    {
    }

    /**
     * Filter operator
     *
     * @param e The edge to be considered
     * @return true if and only if this edge is included in the filtered graph
     */

    template <typename Edge>
    bool operator()(const Edge& e) const
    {
      if (evil_edges_->find(e) != evil_edges_->end())
        return false;

      if (boost::source(e, *graph_) == *articulation_vertex_)
        return false;

      if (boost::target(e, *graph_) == *articulation_vertex_)
        return false;

      return true;
    }

  private:
    const Graph* graph_;
    const vertex_descriptor* articulation_vertex_;
    const edge_set* evil_edges_;
  };

  /**
   * Constructs an articulation edge filter for a given graph.
   *
   * @param graph The given graph
   * @param articulation_vertex Vertex to be excluded
   * @param evil_edges Edges to be excluded
   * @return The constructed filter
   */

  template <typename Graph, typename Vertex, typename EdgeSet>
  inline struct articulation_edge_filter<Graph> make_articulation_edge_filter(const Graph* graph,
      const Vertex* articulation_vertex, const EdgeSet* evil_edges)
  {
    return articulation_edge_filter<Graph> (graph, articulation_vertex, evil_edges);
  }

  /**
   * Tries to extend a 3-connected graph of a graphic minor or detects that
   * the extended minor is not graphic.
   *
   * @param graph The graph of the graphic minor
   * @param matroid The matroid
   * @param matrix Representation matrix of the matroid
   * @param minor_height Number of row elements of the minor
   * @param minor_width Number of column elements of the minor
   * @param extension_type Type of the extension
   * @return true if and only if the extension was possible.
   */

  template <typename MatroidType, typename MatrixType>
  bool extend_graph(matroid_graph& graph, const MatroidType& matroid, const MatrixType& matrix,
      const size_t minor_height, const size_t minor_width, const nested_minor_sequence::extension_type extension_type)
  {
    typedef boost::graph_traits<matroid_graph> traits;

    matroid_element_map element_map = boost::get(edge_matroid_element, graph);
    std::map<int, traits::edge_descriptor> reverse_element_map;

    typename traits::edge_iterator edge_iter, edge_end;
    for (boost::tie(edge_iter, edge_end) = boost::edges(graph); edge_iter != edge_end; ++edge_iter)
    {
      reverse_element_map[element_map[*edge_iter]] = *edge_iter;
    }

    /// Distinct extension cases
    switch (extension_type)
    {
    case nested_minor_sequence::TWO_ROWS_ONE_COLUMN:
    {
      int first_edge_element = find_parallel_to_row(matroid, matrix, minor_height, minor_width, minor_height);
      int second_edge_element = find_parallel_to_row(matroid, matrix, minor_height, minor_width, minor_height + 1);

      /// Find vertices of corresponding edges

      traits::vertex_descriptor first_vertex1 = boost::source(reverse_element_map[first_edge_element], graph);
      traits::vertex_descriptor first_vertex2 = boost::target(reverse_element_map[first_edge_element], graph);
      traits::vertex_descriptor second_vertex1 = boost::source(reverse_element_map[second_edge_element], graph);
      traits::vertex_descriptor second_vertex2 = boost::target(reverse_element_map[second_edge_element], graph);

      if (first_vertex1 == second_vertex2)
        std::swap(second_vertex1, second_vertex2);
      if (first_vertex2 == second_vertex1)
        std::swap(first_vertex1, first_vertex2);
      if (first_vertex2 == second_vertex2)
      {
        std::swap(first_vertex1, first_vertex2);
        std::swap(second_vertex1, second_vertex2);
      }
      if (first_vertex1 != second_vertex1)
        return false;

      /// Remove old edges

      boost::remove_edge(reverse_element_map[first_edge_element], graph);
      boost::remove_edge(reverse_element_map[second_edge_element], graph);

      /// Create new vertices

      traits::vertex_descriptor first_breaker = boost::add_vertex(graph);
      traits::vertex_descriptor second_breaker = boost::add_vertex(graph);

      /// Create new edges

      boost::add_edge(first_vertex2, first_breaker, first_edge_element, graph);
      boost::add_edge(first_breaker, first_vertex1, matroid.name1(minor_height), graph);
      boost::add_edge(second_vertex2, second_breaker, second_edge_element, graph);
      boost::add_edge(second_breaker, second_vertex1, matroid.name1(minor_height + 1), graph);
      boost::add_edge(first_breaker, second_breaker, matroid.name2(minor_width), graph);

      return true;
    }
    break;
    case nested_minor_sequence::ONE_ROW_TWO_COLUMNS:
    {
      int first_edge_element = find_parallel_to_row(make_transposed_matroid(matroid), make_transposed_matrix(matrix),
          minor_width, minor_height, minor_width);
      int second_edge_element = find_parallel_to_row(make_transposed_matroid(matroid), make_transposed_matrix(matrix),
          minor_width, minor_height, minor_width + 1);

      /// Find vertices of corresponding edges

      traits::vertex_descriptor first_vertex1 = boost::source(reverse_element_map[first_edge_element], graph);
      traits::vertex_descriptor first_vertex2 = boost::target(reverse_element_map[first_edge_element], graph);
      traits::vertex_descriptor second_vertex1 = boost::source(reverse_element_map[second_edge_element], graph);
      traits::vertex_descriptor second_vertex2 = boost::target(reverse_element_map[second_edge_element], graph);

      if (first_vertex1 == second_vertex2)
        std::swap(second_vertex1, second_vertex2);
      if (first_vertex2 == second_vertex1)
        std::swap(first_vertex1, first_vertex2);
      if (first_vertex2 == second_vertex2)
      {
        std::swap(first_vertex1, first_vertex2);
        std::swap(second_vertex1, second_vertex2);
      }
      if (first_vertex1 != second_vertex1)
        return false;

      /// Create new vertex

      traits::vertex_descriptor additional_vertex = boost::add_vertex(graph);

      /// Create new edges

      boost::add_edge(first_vertex2, additional_vertex, matroid.name2(minor_width), graph);
      boost::add_edge(second_vertex2, additional_vertex, matroid.name2(minor_width + 1), graph);
      boost::add_edge(first_vertex1, additional_vertex, matroid.name1(minor_height), graph);

      return true;
    }
    break;
    case nested_minor_sequence::ONE_ROW_ONE_COLUMN:
    {
      int row_edge_element = find_parallel_to_row(matroid, matrix, minor_height, minor_width, minor_height);
      int column_edge_element = find_parallel_to_row(make_transposed_matroid(matroid), make_transposed_matrix(matrix),
          minor_width, minor_height, minor_width);

      /// Find vertices of corresponding edges

      traits::vertex_descriptor row_vertex1 = boost::source(reverse_element_map[row_edge_element], graph);
      traits::vertex_descriptor row_vertex2 = boost::target(reverse_element_map[row_edge_element], graph);
      traits::vertex_descriptor column_vertex1 = boost::source(reverse_element_map[column_edge_element], graph);
      traits::vertex_descriptor column_vertex2 = boost::target(reverse_element_map[column_edge_element], graph);

      if (row_vertex1 == column_vertex2)
        std::swap(column_vertex1, column_vertex2);
      if (row_vertex2 == column_vertex1)
        std::swap(row_vertex1, row_vertex2);
      if (row_vertex2 == column_vertex2)
      {
        std::swap(row_vertex1, row_vertex2);
        std::swap(column_vertex1, column_vertex2);
      }
      if (row_vertex1 != column_vertex1)
        return false;

      /// Remove old edges

      boost::remove_edge(reverse_element_map[row_edge_element], graph);

      /// Create new vertex

      traits::vertex_descriptor row_breaker = boost::add_vertex(graph);

      /// Create new edges

      boost::add_edge(row_vertex2, row_breaker, row_edge_element, graph);
      boost::add_edge(row_breaker, row_vertex1, matroid.name1(minor_height), graph);
      boost::add_edge(row_breaker, column_vertex2, matroid.name2(minor_width), graph);

      return true;
    }
    break;
    case nested_minor_sequence::ONE_COLUMN:
    {
      typedef std::set<traits::edge_descriptor> edge_set_t;
      typedef boost::filtered_graph<matroid_graph, boost::is_in_subset<edge_set_t> > edge_subset_graph_t;

      edge_set_t edge_set;

      for (size_t r = 0; r < minor_height; ++r)
      {
        if (matrix(r, minor_width) == 1)
          edge_set.insert(reverse_element_map[matroid.name1(r)]);
      }

      /// Check if edges form a path
      std::vector<traits::vertex_descriptor> path;

      edge_subset_graph_t edge_subset_graph(graph, boost::is_in_subset<edge_set_t>(edge_set));
      if (!tu::util::is_path(edge_subset_graph, boost::get(boost::vertex_index, edge_subset_graph), path))
      {
        return false;
      }

      boost::add_edge(path[0], path[path.size() - 1], matroid.name2(minor_width), graph);
      return true;
    }
    break;
    case nested_minor_sequence::ONE_ROW:
    {
      typedef traits::edge_descriptor edge_t;
      typedef traits::vertex_descriptor vertex_t;
      typedef std::set<vertex_t> vertex_set;
      typedef std::set<edge_t> edge_set;
      typedef std::vector<vertex_t> vertex_vector_t;
      typedef std::vector<edge_t> edge_vector_t;

      typedef boost::vec_adj_list_vertex_id_map<boost::no_property, traits::vertices_size_type> IndexMap;
      IndexMap index_map = boost::get(boost::vertex_index, graph);

      /// Collect all 1-edges

      edge_set one_edges;
      for (size_t c = 0; c < minor_width; ++c)
      {
        if (matrix(minor_height, c) != 0)
          one_edges.insert(reverse_element_map[matroid.name2(c)]);
      }

      traits::vertex_descriptor the_vertex = traits::null_vertex();
      if (tu::util::find_star_vertex(boost::make_filtered_graph(graph, boost::is_in_subset<edge_set>(one_edges)),
          the_vertex))
      {
        /// Create the new vertex and connect it with the_vertex.

        traits::vertex_descriptor new_vertex = boost::add_vertex(graph);
        boost::add_edge(the_vertex, new_vertex, matroid.name1(minor_height), graph);

        /// Reconnect the 1-edges

        for (edge_set::const_iterator edge_iter = one_edges.begin(); edge_iter != one_edges.end(); ++edge_iter)
        {
          traits::edge_descriptor edge = *edge_iter;
          traits::vertex_descriptor other_vertex = boost::source(edge, graph);
          if (other_vertex == the_vertex)
            other_vertex = boost::target(edge, graph);
          int name = element_map[edge];
          boost::remove_edge(the_vertex, other_vertex, graph);
          boost::add_edge(new_vertex, other_vertex, name, graph);
        }

        return true;
      }

      /// Count for each vertex the number of paths that use it

      std::vector<size_t> common_vertex_count(boost::num_vertices(graph), 0);
      for (size_t c = 0; c < minor_width; ++c)
      {
        if (matrix(minor_height, c) == 0)
          continue;

        edge_set edges;
        for (size_t r = 0; r < minor_height; ++r)
        {
          if (matrix(r, c) == 1)
            edges.insert(reverse_element_map[matroid.name1(r)]);
        }

        vertex_set vertices;
        tu::util::used_vertices(boost::make_filtered_graph(graph, boost::is_in_subset<edge_set>(edges)), vertices);

        for (typename vertex_set::const_iterator iter = vertices.begin(); iter != vertices.end(); ++iter)
        {
          common_vertex_count[boost::get(index_map, *iter)]++;
        }
      }

      /// Find articulation points for the graph without 1-edges

      vertex_vector_t articulation_points;
      boost::articulation_points(boost::make_filtered_graph(graph, boost::is_not_in_subset<edge_set>(one_edges)),
          std::back_inserter(articulation_points));

      std::vector<traits::vertex_descriptor> the_vertex_candidates;
      the_vertex = traits::null_vertex();
      for (vertex_vector_t::const_iterator iter = articulation_points.begin(); iter != articulation_points.end();
          ++iter)
      {
        if (common_vertex_count[boost::get(index_map, *iter)] == one_edges.size())
        {
          the_vertex_candidates.push_back(*iter);
        }
      }

      if (the_vertex_candidates.empty() || the_vertex_candidates.size() > 2)
        return false;

      for (vertex_vector_t::const_iterator iter = the_vertex_candidates.begin(); iter != the_vertex_candidates.end();
          ++iter)
      {
        the_vertex = *iter;

        /// Filter the chosen articulation point and the one-edges and check the remaining graph

        vertex_set articulation_set;
        articulation_set.insert(the_vertex);
        std::vector<size_t> component_vector(boost::num_vertices(graph));

        size_t num_components = boost::connected_components(
            boost::make_filtered_graph(graph, make_articulation_edge_filter(&graph, &the_vertex, &one_edges),
                boost::is_not_in_subset<vertex_set>(articulation_set)),
            boost::make_iterator_property_map(component_vector.begin(), index_map));

        /// We should really have articulation point + 2 further components
        assert(num_components >= 2);

        typedef boost::adjacency_list<boost::vecS, boost::vecS, boost::undirectedS> component_graph_t;
        component_graph_t component_graph(num_components);

        bool abort = false;
        for (typename edge_set::const_iterator iter = one_edges.begin(); iter != one_edges.end(); ++iter)
        {
          typename boost::graph_traits<component_graph_t>::vertex_descriptor source, target;
          source = boost::source(*iter, graph);
          target = boost::target(*iter, graph);
          if (source == the_vertex || target == the_vertex)
            continue;

          size_t source_component = component_vector[boost::get(index_map, source)];
          size_t target_component = component_vector[boost::get(index_map, target)];

          if (source_component == target_component)
          {
            /// There cannot be a 1-edge inside one component.
            abort = true;
            break;
          }

          boost::add_edge(boost::vertex(source_component, component_graph),
              boost::vertex(target_component, component_graph), component_graph);
        }

        if (abort)
          continue;

        boost::one_bit_color_map<boost::vec_adj_list_vertex_id_map<boost::no_property, traits::vertices_size_type> > bipartition(
            num_components, boost::get(boost::vertex_index, component_graph));

        if (!boost::is_bipartite(component_graph, boost::get(boost::vertex_index, component_graph), bipartition))
        {
          continue;
        }

        for (size_t i = 0; i < component_vector.size(); ++i)
        {
          if (boost::vertex(i, graph) == the_vertex)
            continue;
        }

        vertex_t new_vertex = boost::add_vertex(graph);

        edge_vector_t reconnect_edges;
        typename traits::out_edge_iterator out_edge_iter, out_edge_end;
        for (boost::tie(out_edge_iter, out_edge_end) = boost::incident_edges(the_vertex, graph);
            out_edge_iter != out_edge_end; ++out_edge_iter)
        {
          vertex_t incident_vertex = boost::target(*out_edge_iter, graph);

          bool reconnect = boost::get(bipartition,
              boost::vertex(component_vector[boost::get(index_map, incident_vertex)], component_graph))
              != boost::one_bit_white;

          if (one_edges.find(*out_edge_iter) != one_edges.end())
            reconnect = !reconnect;

          if (reconnect)
            reconnect_edges.push_back(*out_edge_iter);
        }

        for (typename edge_vector_t::const_iterator iter = reconnect_edges.begin(); iter != reconnect_edges.end();
            ++iter)
        {
          util::reconnect_edge(graph, *iter, the_vertex, new_vertex);
        }

        boost::add_edge(the_vertex, new_vertex, matroid.name1(minor_height), graph);

        return true;
      }

      return false;
    }
    break;
    default:
      throw std::logic_error("Unknown extension in graphicness test.");
    }
  }

  /**
   * Either constructs a graph whose forest matroid is the given matroid or detects
   * that the given matroid is non-graphic.
   *
   * @param matroid The given matroid
   * @param matrix Representation matrix of the given matroid
   * @param nested_minors Sequence of nested minors
   * @return The constructed graph or NULL if the matroid is not graphic.
   */

  template <typename MatroidType, typename MatrixType, typename NestedMinorSequenceType>
  matroid_graph* construct_matroid_graph(const MatroidType& matroid, const MatrixType& matrix,
      const NestedMinorSequenceType& nested_minors, size_t& largest_graphic_minor)
  {
    typedef boost::graph_traits<matroid_graph>::vertex_descriptor vertex_descriptor;

    /// Initialize W3

    largest_graphic_minor = 0;
    matroid_graph* graph = new matroid_graph(4);

    vertex_descriptor center_vertex = boost::vertex(0, *graph);
    vertex_descriptor border_vertex1 = boost::vertex(1, *graph);
    vertex_descriptor border_vertex2 = boost::vertex(2, *graph);
    vertex_descriptor border_vertex3 = boost::vertex(3, *graph);

    boost::add_edge(center_vertex, border_vertex1, matroid.name1(0), *graph).first;
    boost::add_edge(center_vertex, border_vertex2, matroid.name1(1), *graph).first;
    boost::add_edge(center_vertex, border_vertex3, matroid.name2(2), *graph).first;
    boost::add_edge(border_vertex1, border_vertex2, matroid.name2(0), *graph).first;
    boost::add_edge(border_vertex1, border_vertex3, matroid.name2(1), *graph).first;
    boost::add_edge(border_vertex2, border_vertex3, matroid.name1(2), *graph).first;

    size_t minor_height = 3;
    size_t minor_width = 3;

    for (size_t i = 0; i < nested_minors.size(); ++i)
    {
      if (!extend_graph(*graph, matroid, matrix, minor_height, minor_width, nested_minors.get_extension(i)))
      {
        delete graph;
        return NULL;
      }

      minor_height += nested_minors.get_extension_height(i);
      minor_width += nested_minors.get_extension_width(i);
      largest_graphic_minor++;
    }

    return graph;
  }

} /* namespace tu */


//=============================================================================================================
/**
* @file     adaptivemp.cpp
* @author   Martin Henfling <martin.henfling@tu-ilmenau.de>
*           Daniel Knobl <daniel.knobl@tu-ilmenau.de>
*
* @version  1.0
* @date     July, 2014
*
* @section  LICENSE
*
* Copyright (C) 2014, Daniel Knobl and Martin Henfling All rights reserved.
*
* ported to mne-cpp by Martin Henfling and Daniel Knobl in May 2014
* original code was implemented in Matlab Code by Maciej Gratkowski
*
* Redistribution and use in source and binary forms, with or without modification, are permitted provided that
* the following conditions are met:
*     * Redistributions of source code must retain the above copyright notice, this list of conditions and the
*       following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and
*       the following disclaimer in the documentation and/or other materials provided with the distribution.
*     * Neither the name of MNE-CPP authors nor the names of its contributors may be used
*       to endorse or promote products derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
* WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
* PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
* INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
* HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
* NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*
*
* @brief    Implemetation of the Matching Pursuit Algorithm introduced by Stephane Mallat and Zhifeng Zhang.
*           Matlabimplemetation of Maciej Gratkowski is used as Source and reference.
*
*/

//*************************************************************************************************************
//=============================================================================================================
// INCLUDES
//=============================================================================================================

#include "adaptivemp.h"

//*************************************************************************************************************
//=============================================================================================================
// STL INCLUDES
//=============================================================================================================

#include <iostream>
#include <vector>
#include <math.h>

//*************************************************************************************************************
//=============================================================================================================
// Eigen INCLUDES
//=============================================================================================================

#include <unsupported/Eigen/FFT>
#include <Eigen/SparseCore>

//*************************************************************************************************************
//=============================================================================================================
// Qt INCLUDES
//=============================================================================================================

#include <QFuture>
#include <QtConcurrent>

//*************************************************************************************************************
//=============================================================================================================
// USED NAMESPACES
//=============================================================================================================

using namespace UTILSLIB;

//*************************************************************************************************************
//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================

AdaptiveMp::AdaptiveMp()
: it(0)
, max_it(0)
, signal_energy(0)
, current_energy(0)
, fix_phase(0)
, epsilon(0)
, max_iterations(0)
{

}

//*************************************************************************************************************

QList<QList<GaborAtom> > AdaptiveMp::matching_pursuit(MatrixXd signal, qint32 max_iterations, qreal epsilon, bool fix_phase = false, qint32 boost = 0,
                                                      qint32 simplex_it = 1E3, qreal simplex_reflection = 1.0, qreal simplex_expansion = 0.2 ,
                                                      qreal simplex_contraction = 0.5, qreal simplex_full_contraction = 0.5, bool trial_separation = false)
{
    //stop_running = false;
    std::cout << "\nAdaptive Matching Pursuit Algorithm started...\n";

    max_it = max_iterations;
    Eigen::FFT<double> fft;
    MatrixXd residuum = signal; //residuum initialised with signal
    qint32 sample_count = signal.rows();
    qint32 channel_count = signal.cols();

    signal_energy = 0;
    qreal residuum_energy = 0;
    qreal energy_threshold = 0;

    //calculate signal_energy
    for(qint32 channel = 0; channel < channel_count; channel++)
    {
        for(qint32 sample = 0; sample < sample_count; sample++)
            signal_energy += (signal(sample, channel) * signal(sample, channel));

        energy_threshold = 0.01 * epsilon * signal_energy;
        residuum_energy = signal_energy;
    }
    std::cout << "absolute energy of signal: " << residuum_energy << "\n";

    while(it < max_iterations && (energy_threshold < residuum_energy) && sample_count > 1)
    {
        channel_count = channel_count * (boost / 100.0); //reducing the number of observed channels in the algorithm to increase speed performance
        if(boost == 0 || channel_count == 0)
            channel_count = 1;

        //variables for dyadic sampling
        qreal s = 1;                             //scale
        qint32 j = 1;
        VectorXd max_scalar_product = VectorXd::Zero(channel_count);            //inner product for choosing the best matching atom
        qreal k = 0;                             //for modulation 2*pi*k/N
        qint32 p = floor(sample_count / 2);      //translation
        GaborAtom *gabor_Atom = new GaborAtom();
        gabor_Atom->sample_count = sample_count;
        gabor_Atom->energy = 0;
        qreal phase = 0;

        while(s < sample_count)
        {
            k = 0;                               //for modulation 2*pi*k/N
            p = floor(sample_count / 2);         //translation
            VectorXd envelope = GaborAtom::gauss_function(sample_count, s, p);
            VectorXcd fft_envelope = RowVectorXcd::Zero(sample_count);
            fft.fwd(fft_envelope, envelope);

            while(k < sample_count/2)
            {
                p = floor(sample_count/2);
                VectorXcd modulation = modulation_function(sample_count, k);
                VectorXcd modulated_resid = VectorXcd::Zero(sample_count);
                VectorXcd fft_modulated_resid = VectorXcd::Zero(sample_count);
                VectorXcd fft_m_e_resid = VectorXcd::Zero(sample_count);
                VectorXd corr_coeffs = VectorXd::Zero(sample_count);

                //iteration for multichannel, depending on boost setting
                for(qint32 chn = 0; chn < channel_count; chn++)
                {
                    qint32 max_index = 0;
                    qreal maximum = 0;
                    phase = 0;
                    p = floor(sample_count/2);//here is difference to dr. gratkowski´s code (he didn´t reset parameter p)

                    //complex correlation of signal and sinus-modulated gaussfunction
                    for(qint32 l = 0; l< sample_count; l++)
                        modulated_resid[l] = residuum(l, chn) * modulation[l];

                    fft.fwd(fft_modulated_resid, modulated_resid);

                    for( qint32 m = 0; m < sample_count; m++)
                        fft_m_e_resid[m] = fft_modulated_resid[m] * conj(fft_envelope[m]);

                    fft.inv(corr_coeffs, fft_m_e_resid);
                    maximum = corr_coeffs[0];

                    //find index of maximum correlation-coefficient to use in translation
                    for(qint32 i = 1; i < corr_coeffs.rows(); i++)
                        if(maximum < corr_coeffs[i])
                        {
                            maximum = corr_coeffs[i];
                            max_index = i;
                        }

                    //adapting translation p to create atomtranslation correctly
                    if(max_index >= p) p = max_index - p + 1;
                    else p = max_index + p;

                    VectorXd atom_parameters = calculate_atom(sample_count, s, p, k, chn, residuum, RETURNPARAMETERS, fix_phase);
                    qreal temp_scalar_product = 0;
                    if(trial_separation) temp_scalar_product = max_scalar_product[chn];
                    else temp_scalar_product = max_scalar_product[0];

                    if(std::fabs(atom_parameters[4]) >= std::fabs(temp_scalar_product))
                    {
                        //set highest scalarproduct, in comparison to best matching atom
                        gabor_Atom->scale              = atom_parameters[0];
                        gabor_Atom->translation        = atom_parameters[1];
                        gabor_Atom->modulation         = atom_parameters[2];
                        gabor_Atom->phase              = atom_parameters[3];
                        gabor_Atom->max_scalar_product = atom_parameters[4];
                        gabor_Atom->bm_channel         = chn;

                        if(trial_separation)
                        {
                            max_scalar_product[chn]    = atom_parameters[4];

                            if(atoms_in_chns.length() < channel_count)
                                atoms_in_chns.append(*gabor_Atom);
                            else
                                atoms_in_chns.replace(chn, *gabor_Atom);
                        }
                        else
                            max_scalar_product[0]      = atom_parameters[4];

                    }

                }
                k += pow(2.0,(-j))*sample_count/2;

            }
            j++;
            s = pow(2.0,j);
        }
        std::cout << "\n" << "===============" << " found parameters " << it + 1 << "===============" << ":\n\n"<<
                     "scale: " << gabor_Atom->scale << " trans: " << gabor_Atom->translation <<
                     " modu: " << gabor_Atom->modulation << " phase: " << gabor_Atom->phase << " sclr_prdct: " << gabor_Atom->max_scalar_product << "\n";

        //replace atoms with s==N and p = floor(N/2) by such atoms that do not have an envelope
        k = 0;
        s = sample_count;
        p = floor(sample_count / 2);
        j = floor(log10(sample_count)/log10(2));//log(sample_count) / log(2));
        phase = 0;

        //iteration for multichannel, depending on boost setting
        for(qint32 chn = 0; chn < channel_count; chn++)
        {
            k = 0;

            while(k < sample_count / 2)
            {
                VectorXd parameters_no_envelope = calculate_atom(sample_count, s, p, k, chn, residuum, RETURNPARAMETERS, fix_phase);

                qreal temp_scalar_product = 0;
                if(trial_separation) temp_scalar_product = max_scalar_product[chn];
                else temp_scalar_product = max_scalar_product[0];
                if(std::fabs(parameters_no_envelope[4]) > std::fabs(temp_scalar_product))
                {
                    //set highest scalarproduct, in comparison to best matching atom

                    gabor_Atom->scale              = parameters_no_envelope[0];
                    gabor_Atom->translation        = parameters_no_envelope[1];
                    gabor_Atom->modulation         = parameters_no_envelope[2];
                    gabor_Atom->phase              = parameters_no_envelope[3];
                    gabor_Atom->max_scalar_product = parameters_no_envelope[4];
                    gabor_Atom->bm_channel         = chn;

                    if(trial_separation)
                    {
                        max_scalar_product[chn]    = parameters_no_envelope[4];
                        atoms_in_chns.replace(chn, *gabor_Atom);
                    }
                    else
                        max_scalar_product[0]      = parameters_no_envelope[4];

                }
                k += pow(2.0,(-j))*sample_count/2;

            }

        }
        std::cout << "      after comparison to NoEnvelope " << ":\n"<< "scale: " << gabor_Atom->scale << " trans: " << gabor_Atom->translation <<
                     " modu: " << gabor_Atom->modulation << " phase: " << gabor_Atom->phase << " sclr_prdct: " << gabor_Atom->max_scalar_product << "\n\n";

        //simplexfunction to find minimum of target among parameters s, p, k

        if(trial_separation && simplex_it != 0)
        {
            for(qint32 chn = 0; chn < atoms_in_chns.length(); chn++)
            {
                *gabor_Atom = atoms_in_chns.at(chn);
                simplex_maximisation(simplex_it, simplex_reflection, simplex_expansion, simplex_contraction, simplex_full_contraction,
                                     gabor_Atom, max_scalar_product, sample_count, fix_phase, residuum, trial_separation, chn);

            }
        }
        else if(simplex_it != 0)
            simplex_maximisation(simplex_it, simplex_reflection, simplex_expansion, simplex_contraction, simplex_full_contraction,
                                 gabor_Atom, max_scalar_product, sample_count, fix_phase, residuum, trial_separation, gabor_Atom->bm_channel);

        //calc multichannel parameters phase and max_scalar_product
        channel_count = signal.cols();
        VectorXd best_match = VectorXd::Zero(sample_count);

        for(qint32 chn = 0; chn < channel_count; chn++)
        {

            if(trial_separation)
            {
                *gabor_Atom = atoms_in_chns.at(chn);
                gabor_Atom->energy = 0;
                best_match = gabor_Atom->create_real(gabor_Atom->sample_count, gabor_Atom->scale, gabor_Atom->translation,
                                                     gabor_Atom->modulation, gabor_Atom->phase);
            }
            else
            {
                VectorXd channel_params = calculate_atom(sample_count, gabor_Atom->scale, gabor_Atom->translation,
                                                         gabor_Atom->modulation, chn, residuum, RETURNPARAMETERS, fix_phase);
                gabor_Atom->phase_list.append(channel_params[3]);

                gabor_Atom->max_scalar_list.append(channel_params[4]);

                best_match = gabor_Atom->create_real(gabor_Atom->sample_count, gabor_Atom->scale, gabor_Atom->translation,
                                                     gabor_Atom->modulation, gabor_Atom->phase_list.at(chn));
            }

            //substract best matching Atom from Residuum in each channel
            for(qint32 j = 0; j < gabor_Atom->sample_count; j++)
            {
                if(!trial_separation)
                {
                    residuum(j,chn) -= gabor_Atom->max_scalar_list.at(chn) * best_match[j];
                    gabor_Atom->energy += pow(gabor_Atom->max_scalar_list.at(chn) * best_match[j], 2);
                }
                else
                {
                    residuum(j,chn) -= gabor_Atom->max_scalar_product * best_match[j];
                    gabor_Atom->energy += pow(gabor_Atom->max_scalar_product * best_match[j], 2);
                }
            }
            if(trial_separation)
            {
                atoms_in_chns.replace(chn, *gabor_Atom);            // change energy
                residuum_energy -= atoms_in_chns.at(chn).energy;// / channel_count;;
                current_energy  += atoms_in_chns.at(chn).energy;// / channel_count;;
            }
        }

        if(!trial_separation)
        {
            residuum_energy -= gabor_Atom->energy;
            current_energy  += gabor_Atom->energy;
        }
        else
            /*for(qint32 i = 0; i < atoms_in_chns.length(); i++)
            {
                residuum_energy -= atoms_in_chns.at(i).energy / channel_count;
                current_energy  += atoms_in_chns.at(i).energy / channel_count;
            }*/

        std::cout << "absolute energy of residue: " << residuum_energy << "\n";

        if(!trial_separation)
            atoms_in_chns.append(*gabor_Atom);

        atom_list.append(atoms_in_chns);

        atoms_in_chns.clear();
        delete gabor_Atom;
        it++;

        emit current_result(it, max_it, current_energy, signal_energy, residuum, atom_list, fix_dict_list);

        if( QThread::currentThread()->isInterruptionRequested())
        {
            send_warning(10);
            break;
        }

    }//end iterations
    std::cout << "\nAdaptive Matching Pursuit Algorithm finished.\n";
    emit finished_calc();
    return atom_list;
}

//*************************************************************************************************************

VectorXcd AdaptiveMp::modulation_function(qint32 N, qreal k)
{
    VectorXcd modulation = VectorXcd::Zero(N);

    for(qint32 n = 0; n < N; n++)
    {
        modulation[n] = std::polar(1 / sqrt(qreal(N)), 2 * PI * k / qreal(N) * qreal(n));
    }
    return modulation;
}

//*************************************************************************************************************

VectorXd AdaptiveMp::calculate_atom(qint32 sample_count, qreal scale, qint32 translation, qreal modulation, qint32 channel, MatrixXd residuum, ReturnValue return_value = RETURNATOM, bool fix_phase = false)
{
    GaborAtom *gabor_Atom = new GaborAtom();
    qreal phase = 0;
    //create complex Gaboratom
    VectorXcd complex_gabor_atom = gabor_Atom->create_complex(sample_count, scale, translation, modulation);

    //calculate Inner Product: preparation to find the parameter phase
    std::complex<double> inner_product(0, 0);

    if(fix_phase == false)
    {
        for(qint32 i = 0; i < sample_count; i++)
            inner_product += residuum(i, channel) * conj(complex_gabor_atom[i]);
    }
    else
    {
        for(qint32 chn = 0; chn < residuum.cols(); chn++)
        {
            for(qint32 i = 0; i < sample_count; i++)
                inner_product += residuum(i, chn) * conj(complex_gabor_atom[i]);
        }
        if(residuum.cols() != 0)
            inner_product /= residuum.cols();
    }

    //calculate phase to create realGaborAtoms
    phase = std::arg(inner_product);
    if (phase < 0) phase = 2 * PI - phase;
    VectorXd real_gabor_atom = gabor_Atom->create_real(sample_count, scale, translation, modulation, phase);

    delete gabor_Atom;

    switch(return_value)
    {
    case RETURNPARAMETERS:
    {

        qreal scalar_product = 0;

        for(qint32 i = 0; i < sample_count; i++)
            scalar_product += real_gabor_atom[i] * residuum(i, channel);

        VectorXd atom_parameters = VectorXd::Zero(5);

        atom_parameters[0] = scale;
        atom_parameters[1] = translation;
        atom_parameters[2] = modulation;
        atom_parameters[3] = phase;
        atom_parameters[4] = scalar_product;

        return atom_parameters;
    }

    case RETURNATOM: {return real_gabor_atom;} //returns normalized realGaborAtom

    default: return real_gabor_atom;// only to avoid compiler warnings
    }
}

//*************************************************************************************************************

void AdaptiveMp::simplex_maximisation(qint32 simplex_it, qreal simplex_reflection, qreal simplex_expansion, qreal simplex_contraction, qreal simplex_full_contraction,
                                      GaborAtom *gabor_Atom, VectorXd max_scalar_product, qint32 sample_count, bool fix_phase, MatrixXd residuum, bool trial_separation, qint32 chn)
{
    //Maximisation Simplex Algorithm implemented by Botao Jia, adapted to the MP Algorithm by Martin Henfling. Copyright (C) 2010 Botao Jia
    //ToDo: change to clean use of EIGEN, @present its mixed with Namespace std and <vector>
    std::vector<double> init;

    init.push_back(gabor_Atom->scale);
    init.push_back(gabor_Atom->translation);
    init.push_back(gabor_Atom->modulation);

    double tol = 1E8 * std::numeric_limits<double>::epsilon();
    std::vector<std::vector<double> > x = std::vector<std::vector<double> >();
    qint32 iterations = simplex_it;
    qint32 N = qint32 (init.size());                //space dimension

    VectorXd atom_fxc_params = VectorXd::Zero(5);   //initialisation for contraction coefficients

    qreal a = simplex_reflection, b = simplex_expansion, g = simplex_contraction, h = simplex_full_contraction;
                                                //coefficients default a = 1, b = 0.2, g = 0.5, h = 0.5
                                                //a: reflection  -> xr step away from worst siplex found
                                                //b: expansion   -> xe if better with a so go in this direction with b
                                                //g: contraction -> xc calc new worst point an bring closer to middle of simplex
                                                //h: full contraction to x1
    std::vector<double> xcentroid_old(N,0);     //simplex center * (N+1)
    std::vector<double> xcentroid_new(N,0);     //simplex center * (N+1)
    std::vector<double> vf(N+1,0);              //f evaluated at simplex vertices
    qint32 x1 = 0, xn = 0, xnp1 = 0;            //x1:   f(x1) = min { f(x1), f(x2)...f(x_{n+1} }
                                                //xnp1: f(xnp1) = max { f(x1), f(x2)...f(x_{n+1} }
                                                //xn:   f(xn)<f(xnp1) && f(xn)> all other f(x_i)
    qint32 cnt = 0; //iteration step number

    if(x.size()== 0) //if no initial simplex is specified
    {
        //construct the trial simplex
        //based upon the initial guess parameters
        std::vector<double> del( init );
        std::transform(del.begin(), del.end(), del.begin(),
                       std::bind2nd( std::divides<double>() , 20) );//'20' is picked
        //assuming initial trail close to true

        for(qint32 i = 0; i < N; ++i)
        {
            std::vector<double> tmp( init );
            tmp[i] +=  del[i];
            x.push_back( tmp );
        }

        x.push_back(init);//x.size()=N+1, x[i].size()=N

        //xcentriod
        std::transform(init.begin(), init.end(), xcentroid_old.begin(), std::bind2nd(std::multiplies<double>(), N+1) );
    }//constructing the simplex finished

    //optimization begins
    for(cnt=0; cnt<iterations; ++cnt)
    {
        for(qint32 i=0; i < N+1; ++i)
        {
            VectorXd atom_fx = VectorXd::Zero(sample_count);

            if(gabor_Atom->scale == sample_count && gabor_Atom->translation == floor(sample_count / 2))
                atom_fx = calculate_atom(sample_count, sample_count, floor(sample_count / 2), x[i][2], chn, residuum, RETURNATOM, fix_phase);

            else
                atom_fx = calculate_atom(sample_count, x[i][0], x[i][1], x[i][2], chn, residuum, RETURNATOM, fix_phase);

            //create targetfunction of realGaborAtom and Residuum
            double target = 0;
            for(qint32 k = 0; k < atom_fx.rows(); k++)
            {
                target -=atom_fx[k]*residuum(k, chn); //ToDo: old residuum(k,0)
            }

            vf[i] = target;
        }

        x1=0; xn=0; xnp1=0;//find index of max, second max, min of vf.

        for(quint32 i=0; i < vf.size(); ++i)
        {
            if(vf[i]<vf[x1])      x1 = i;
            if(vf[i]>vf[xnp1])    xnp1 = i;
        }

        xn = x1;

        for(quint32 i=0; i<vf.size();++i) if(vf[i]<vf[xnp1] && vf[i]>vf[xn])  xn=i;

        //x1, xn, xnp1 are found

        std::vector<double> xg(N, 0);//xg: centroid of the N best vertexes

        for(quint32 i=0; i<x.size(); ++i) if(i!=xnp1) std::transform(xg.begin(), xg.end(), x[i].begin(), xg.begin(), std::plus<double>() );

        std::transform(xg.begin(), xg.end(), x[xnp1].begin(), xcentroid_new.begin(), std::plus<double>());
        std::transform(xg.begin(), xg.end(), xg.begin(), std::bind2nd(std::divides<double>(), N) );
        //xg found, xcentroid_new updated

        //termination condition
        double diff=0;          //calculate the difference of the simplex centers

        //see if the difference is less than the termination criteria
        for(qint32 i=0; i<N; ++i) diff += std::fabs(xcentroid_old[i]-xcentroid_new[i]);

        if (diff/N < tol) break;              //terminate the optimizer
        else xcentroid_old.swap(xcentroid_new); //update simplex center

        //reflection:
        std::vector<double> xr(N,0);

        for( qint32 i=0; i<N; ++i) xr[i]=xg[i]+a*(xg[i]-x[xnp1][i]);
        //reflection, xr found

        VectorXd atom_fxr = VectorXd::Zero(sample_count);

        if(gabor_Atom->scale == sample_count && gabor_Atom->translation == floor(sample_count / 2))
            atom_fxr = calculate_atom(sample_count, sample_count, floor(sample_count / 2), xr[2], chn, residuum, RETURNATOM, fix_phase);

        else
            atom_fxr = calculate_atom(sample_count, xr[0], xr[1], xr[2], chn, residuum, RETURNATOM, fix_phase);

        //create targetfunction of realGaborAtom and Residuum
        double fxr = 0;
        for(qint32 k = 0; k < atom_fxr.rows(); k++) fxr -=atom_fxr[k]*residuum(k,chn);//ToDo: old residuum(k,0)

        //double fxr = target;//record function at xr

        if(vf[x1]<=fxr && fxr<=vf[xn]) std::copy(xr.begin(), xr.end(), x[xnp1].begin());

        //expansion:
        else if(fxr<vf[x1])
        {
            std::vector<double> xe(N,0);

            for( qint32 i=0; i<N; ++i) xe[i]=xr[i]+b*(xr[i]-xg[i]);

            VectorXd atom_fxe = VectorXd::Zero(sample_count);

            if(gabor_Atom->scale == sample_count && gabor_Atom->translation == floor(sample_count / 2))
                atom_fxe = calculate_atom(sample_count, sample_count, floor(sample_count / 2), xe[2], chn, residuum, RETURNATOM, fix_phase);

            else
                atom_fxe = calculate_atom(sample_count, xe[0], xe[1], xe[2], chn, residuum, RETURNATOM, fix_phase);

            //create targetfunction of realGaborAtom and Residuum
            double fxe = 0;
            for(qint32 k = 0; k < atom_fxe.rows(); k++) fxe -=atom_fxe[k]*residuum(k,chn);//ToDo: old residuum(k,0)

            if( fxe < fxr ) std::copy(xe.begin(), xe.end(), x[xnp1].begin() );
            else std::copy(xr.begin(), xr.end(), x[xnp1].begin() );
        }//expansion finished,  xe is not used outside the scope

        //contraction:
        else if( fxr > vf[xn] )
        {
            std::vector<double> xc(N,0);

            for( qint32 i=0; i<N; ++i)
                xc[i]=xg[i]+g*(x[xnp1][i]-xg[i]);

            if(gabor_Atom->scale == sample_count && gabor_Atom->translation == floor(sample_count / 2))
                atom_fxc_params = AdaptiveMp::calculate_atom(sample_count, sample_count, floor(sample_count / 2), xc[2], chn, residuum, RETURNPARAMETERS, fix_phase);

            else
                atom_fxc_params = AdaptiveMp::calculate_atom(sample_count, xc[0], xc[1], xc[2], chn, residuum, RETURNPARAMETERS, fix_phase);

            VectorXd atom_fxc = gabor_Atom->create_real(gabor_Atom->sample_count, atom_fxc_params[0], atom_fxc_params[1], atom_fxc_params[2], atom_fxc_params[3]);

            atom_fxc_params[4] = 0;

            for(qint32 i = 0; i < sample_count; i++)
                atom_fxc_params[4] += atom_fxc[i] * residuum(i, chn);

            //create targetfunction of realGaborAtom and Residuum
            double fxc = 0;

            for(qint32 k = 0; k < atom_fxc.rows(); k++)
                fxc -=atom_fxc[k]*residuum(k,chn);//ToDo: old residuum(k,0)

            if( fxc < vf[xnp1] )
                std::copy(xc.begin(), xc.end(), x[xnp1].begin() );

            else
                for( quint32 i=0; i<x.size(); ++i )
                    if( i!=x1 )
                        for(qint32 j=0; j<N; ++j)
                            x[i][j] = x[x1][j] + h * ( x[i][j]-x[x1][j] );
        }//contraction finished, xc is not used outside the scope
    }//optimization is finished
    //end Maximisation Copyright (C) 2010 Botao Jia

    if(iterations != 0)
    {

        if(gabor_Atom->scale == sample_count && gabor_Atom->translation == floor(sample_count / 2))
            atom_fxc_params = AdaptiveMp::calculate_atom(sample_count, sample_count, floor(sample_count / 2), x[x1][2], chn, residuum, RETURNPARAMETERS, fix_phase);

        else
            atom_fxc_params = AdaptiveMp::calculate_atom(sample_count, x[x1][0], x[x1][1], x[x1][2], chn, residuum, RETURNPARAMETERS, fix_phase);

        qreal temp_scalar_product = 0;
        if(trial_separation) temp_scalar_product = max_scalar_product[chn];
        else temp_scalar_product = max_scalar_product[0];

        if(std::fabs(atom_fxc_params[4]) > std::fabs(temp_scalar_product) && atom_fxc_params[1] < sample_count && atom_fxc_params[1] > 0)//ToDo: find a way to make the simplex not running out of bounds
        {
            //set highest scalarproduct, in comparison to best matching atom

            gabor_Atom->scale              = atom_fxc_params[0];
            gabor_Atom->translation        = atom_fxc_params[1];
            gabor_Atom->modulation         = atom_fxc_params[2];
            gabor_Atom->phase              = atom_fxc_params[3];
            gabor_Atom->max_scalar_product = atom_fxc_params[4];

            if(trial_separation)
                atoms_in_chns.replace(chn, *gabor_Atom);
        }

        if(cnt==iterations)//max number of iteration achieves before tol is satisfied
        {
            send_warning(11);
            std::cout<<"Simplex Iteration limit of "<<iterations<<" achieved, result may not be optimal\n";
        }

        std::cout <<  "      after simplex optimization " << ":\n"<< "scale: " << gabor_Atom->scale << " trans: " << gabor_Atom->translation <<
                      " modu: " << gabor_Atom->modulation << " phase: " << gabor_Atom->phase << " sclr_prdct: " << gabor_Atom->max_scalar_product << "\n\n";
    }
}

//*************************************************************************************************************

void AdaptiveMp::recieve_input(Eigen::MatrixXd signal, qint32 max_iterations, qreal epsilon, bool fix_phase = false, qint32 boost = 0, qint32 simplex_it = 1E3,
                               qreal simplex_reflection = 1.0, qreal simplex_expansion = 0.2, qreal simplex_contraction = 0.5, qreal simplex_full_contraction = 0.5, bool trial_separation = false)
{
    matching_pursuit(signal, max_iterations, epsilon, fix_phase, boost, simplex_it, simplex_reflection, simplex_expansion, simplex_contraction, simplex_full_contraction, trial_separation);
}

//*************************************************************************************************************

AdaptiveMp::~AdaptiveMp()
{
}