#include "tile/lang/type.h"

#include <cmath>
#include <deque>
#include <queue>
#include <set>
#include <stdexcept>
#include <vector>

#include <boost/format.hpp>

#include "base/util/error.h"
#include "tile/lang/builtins.h"
#include "tile/lang/fpconv.h"
#include "tile/lang/gen_special.h"
#include "tile/lang/parser.h"
#include "tile/lang/replace.h"
#include "tile/lang/sym_poly.h"

namespace vertexai {
namespace tile {
namespace lang {

using namespace math;  // NOLINT

namespace {
DataType g_floatx = DataType::FLOAT32;
}  // namespace

void SetFloatX(DataType dtype) { g_floatx = dtype; }

std::string Binding::key() const {
  switch (tag) {
    case Binding::TENSOR:
      return "tensor";
    case Binding::ICONST:
      return std::to_string(iconst);
    case Binding::FCONST: {
      std::string c = DoubleToString(fconst);
      if (c.find_first_of(".e") == std::string::npos) {
        c += ".0";
      }
      return c;
    }
    case Binding::TUPLE:
    default:
      throw std::logic_error{"Invalid binding for key"};
  }
}

bool Binding::operator==(const Binding& rhs) const {
  if (tag != rhs.tag) {
    return false;
  }
  switch (tag) {
    case Binding::TENSOR:
      return shape == rhs.shape;
    case Binding::ICONST:
      return iconst == rhs.iconst;
    case Binding::FCONST:
      return fconst == rhs.fconst;
    case Binding::TUPLE:
      return tuple == rhs.tuple;
    default:
      throw std::logic_error{"Invalid binding"};
  }
}

bool Binding::operator!=(const Binding& rhs) const { return !(*this == rhs); }

static double ConstantPropagate(const std::string& op, const std::vector<double>& x) {
  if (op == "ident") {
    return x[0];
  }
  if (op == "broadcast") {
    if (x[0] != x[1] && x[0] != 1 && x[1] != 1) {
      throw std::runtime_error("Type check failed due to mismatched tensor sizes: " + std::to_string(x[0]) +
                               " != " + std::to_string(x[1]));
    }
    return x[0] == 1 ? x[1] : x[0];
  }
  if (op == "match") {
    if (x[0] != x[1]) {
      throw std::runtime_error("Type check failed due to mismatched tensor sizes: " + std::to_string(x[0]) +
                               " != " + std::to_string(x[1]));
    }
    return x[0];
  }
  if (op == "neg") {
    return -x[0];
  }
  if (op == "recip") {
    return 1.0 / x[0];
  }
  if (op == "add") {
    return x[0] + x[1];
  }
  if (op == "sub") {
    return x[0] - x[1];
  }
  if (op == "mul") {
    return x[0] * x[1];
  }
  if (op == "div") {
    return x[0] / x[1];
  }
  if (op == "mod") {
    return static_cast<double>(int64_t(x[0]) % int64_t(x[1]));
  }
  if (op == "cmp_eq") {
    return x[0] == x[1];
  }
  if (op == "cmp_ne") {
    return x[0] != x[1];
  }
  if (op == "cmp_lt") {
    return x[0] < x[1];
  }
  if (op == "cmp_gt") {
    return x[0] > x[1];
  }
  if (op == "cmp_le") {
    return x[0] <= x[1];
  }
  if (op == "cmp_ge") {
    return x[0] >= x[1];
  }
  if (op == "cond") {
    return (x[0] ? x[1] : x[2]);
  }
  if (op == "max") {
    return (x[0] < x[1] ? x[1] : x[0]);
  }
  if (op == "ceil") {
    return std::ceil(x[0]);
  }
  if (op == "floor") {
    return std::floor(x[0]);
  }
  throw std::runtime_error("Unknown op " + op + " during constant propagation");
}

std::string to_string(const std::vector<TensorDimension>& dims) {
  std::string result = "[";
  bool first = true;
  for (const auto& dim : dims) {
    if (first) {
      first = false;
    } else {
      result += ',';
    }
    result += std::to_string(dim.size);
  }
  result += ']';
  return result;
}

namespace {

bool BroadcastTo(std::vector<TensorDimension>* dims, std::string* dims_source,
                 const std::vector<TensorDimension>& vdims, const std::string& vdims_source) {
  if (vdims.size() != 0) {
    if (dims->size() == 0) {
      *dims = vdims;
      *dims_source = vdims_source;
      return false;
    } else if (*dims != vdims) {
      // Check for broadcast compatibility.
      IVLOG(4, "Checking compatibility between " << to_string(*dims) << " and " << to_string(vdims));
      auto di = dims->rbegin();
      auto vi = vdims.rbegin();
      for (;; ++di, ++vi) {
        IVLOG(4, "Top of check loop");
        if (vi == vdims.rend()) {
          // vdims can be broadcast to dims.
          IVLOG(4, "vdims broadcasts to dims");
          break;
        }
        if (di == dims->rend()) {
          // Anything that was used to produce dims can be broadcast
          // to vdims; we just need to augment dims with the remaining
          // elements of vdims.
          dims->insert(dims->begin(), vdims.begin(), vdims.begin() + (vdims.size() - (vi - vdims.rbegin())));
          IVLOG(4, "dims broadcasts to vdims; dims = " << to_string(*dims));
          break;
        }
        IVLOG(4, "Considering " << di->size << " vs. " << vi->size);
        if (vi->size == di->size) {
          IVLOG(4, "No broadcasting needed (here)");
          continue;
        }
        if (vi->size == 1) {
          // This dimension of vi can be broadcast to whatever is in dims.
          IVLOG(4, "di broadcasts to vi");
          continue;
        }
        if (di->size == 1) {
          // This dimension of di can be broadcast to whatever is in di.
          di->size = vi->size;
          IVLOG(4, "vi broadcasts to di");
          continue;
        }
        // Otherwise, broadcasting cannot be done.
        throw std::runtime_error("Mismatched tensor shapes in elementwise operation: " + *dims_source +
                                 to_string(*dims) + " can't match " + vdims_source + to_string(vdims));
      }
      IVLOG(4, "Broadcast possible; LCM dims=" << to_string(*dims));
      return true;
    }
  }
  return false;
}

int64_t ExtractInteger(const Bindings& vars, const std::string name) {
  // Handle raw constants
  if (name[0] >= '0' && name[0] <= '9') {
    return std::stoi(name);
  }
  auto it = vars.find(name);
  if (it == vars.end()) {
    throw std::runtime_error("Unknown variable " + name + " in expression");
  }
  if (it->second.tag != Binding::ICONST) {
    throw std::runtime_error("Variable " + name + " used in a context which requires it to be a constant integer");
  }
  return it->second.iconst;
}

}  // namespace

void TypeCheck(Program* prog, Bindings* vars) {
  IVLOG(4, "Before ApplyDefines: " << prog->ops);
  ApplyDefines(prog, InlineDefines);
  IVLOG(4, "After ApplyDefines: " << prog->ops);
  // First, process inputs, validating types and extracting constants
  for (Input& in : prog->inputs) {
    IVLOG(3, "Type checking input " << in.name);
    // Pull the variable by id from the bindings
    auto it = vars->find(in.name);
    // If it doesn't exist, fail
    if (it == vars->end()) {
      throw std::runtime_error("Input " + in.name + " not bound");
    }
    // If input isn't fixed, it's always cool
    if (in.tag == Input::VARIABLE) {
      IVLOG(3, "  Looks like a variable");
      continue;
    }
    // Otherwise, verify it's a tensor
    if (it->second.tag != Binding::TENSOR && in.dims.size() > 0) {
      throw std::runtime_error("Non tensor binding passed to tensor input " + in.name);
    }
    // Extract it's shape and make sure dimensions match
    const TensorShape& shape = it->second.shape;
    if (in.dims.size() != shape.dims.size()) {
      throw std::runtime_error(str(boost::format("Mismatch number of dimensions, %lu vs %lu for %s") %  //
                                   in.dims.size() % shape.dims.size() % in.name));
    }
    // Now, go over each dimension and extract constants
    for (size_t i = 0; i < in.dims.size(); i++) {
      // Verify that the size polynomial is valid
      size_t size = shape.dims[i].size;
      std::string var = in.dims[i];
      // Check for var in our bindings
      auto it2 = vars->find(var);
      if (it2 == vars->end()) {
        // If it's not bound, bind it
        vars->emplace(var, Binding(static_cast<int64_t>(size)));
      } else {
        // Otherwise, make sure it's a match to the binding
        if (it2->second.tag != Binding::ICONST) {
          throw std::runtime_error(var + " used as tensor size and is not an integer");
        }
        size_t othersize = static_cast<size_t>(it2->second.iconst);
        if (othersize != size) {
          throw std::runtime_error(str(boost::format("Mismatched sizes for %s, %lu vs %lu") % var % othersize % size));
        }
      }
    }
  }

  // Now, walk over all the ops and process them...
  for (Op& op : prog->ops) {
    IVLOG(3, "Processing op " << op);
    // Error on already computed values
    if (vars->count(op.output)) {
      throw std::runtime_error("Reassignment to " + op.output);
    }
    // Handle constants
    if (op.tag == Op::CONSTANT && op.f.fn == "fconst") {
      IVLOG(4, "Found fconst op " << to_string(op));
      vars->emplace(op.output, Binding(std::stof(op.inputs[0].c_str()), g_floatx));
      continue;
    }
    if (op.tag == Op::CONSTANT && op.f.fn == "iconst") {
      IVLOG(4, "Found iconst op " << to_string(op));
      vars->emplace(op.output, Binding(int64_t(std::stoll(op.inputs[0].c_str()))));
      continue;
    }

    // Handle simple_reduce function before processing output types etc
    if (op.tag == Op::FUNCTION && op.f.fn == "simple_reduce" && op.inputs.size() >= 1) {
      IVLOG(4, "Running simple_reduce");
      Binding tensorBinding = vars->at(op.inputs[0]);
      if (tensorBinding.tag != Binding::TENSOR) {
        throw std::runtime_error("Internal error in simple_reduce: invalid output dim type");
      }
      TensorShape shape = tensorBinding.shape;
      std::vector<size_t> new_shape;
      for (size_t i = 1; i < op.inputs.size(); i++) {
        Binding sizeBinding = vars->at(op.inputs[i]);
        if (sizeBinding.tag != Binding::ICONST) {
          throw std::runtime_error("Internal error in simple_reduce: invalid input dim type");
        }
        new_shape.push_back(sizeBinding.iconst);
      }
      std::vector<SymbolicPolynomialPtr> src_indices;
      for (size_t idx = 0; idx < shape.dims.size(); ++idx) {
        src_indices.push_back(SymbolicPolynomial::MakeIndex(std::string("i") + std::to_string(idx)));
      }
      std::vector<SymbolicPolynomialPtr> out_indices;
      size_t i = 0;
      for (auto it = src_indices.begin() + (shape.dims.size() - new_shape.size()); it != src_indices.end(); ++it, ++i) {
        if (new_shape[i] == 1) {
          out_indices.push_back(SymbolicPolynomial::MakeLiteral(0));
        } else {
          out_indices.push_back(*it);
        }
      }
      std::vector<std::string> dims;
      for (auto& dim : new_shape) {
        dims.emplace_back(std::to_string(dim));
      }

      bool nontrivial_contraction = false;
      if (new_shape.size() == shape.dims.size()) {
        for (size_t i = 0; i < shape.dims.size(); ++i) {
          if (new_shape[i] != shape.dims[i].size) {
            nontrivial_contraction = true;
          }
        }
      } else {
        nontrivial_contraction = true;
      }
      if (nontrivial_contraction) {
        IVLOG(4, "Rewriting comb and agg ops with simple_reduce.");
        op.c.comb_op = CombinationOp::PLUS;
        op.c.agg_op = AggregationOp::SUM;
        op.c.output_size = dims;
        op.c.specs.resize(2);
        op.c.specs[0].id = op.output;
        op.c.specs[1].id = op.inputs[0];
        op.c.specs[0].sspec = out_indices;
        op.c.specs[1].sspec = src_indices;
        op.tag = Op::CONTRACTION;  // simple_reduce created a contraction
        op.f = Function();         // Clear function
      } else {
        IVLOG(4, "Using identity for simple_reduce");
        op.f.fn = "ident";
        op.inputs.resize(1);
      }
    }

    // Validate all inputs have a type already computed
    for (const auto& s : op.inputs) {
      if (vars->count(s) == 0) {
        throw std::runtime_error("Op is using undefined input: " + s);
      }
    }

    // Compute result type by 'upcasting' to the highest type in the hierarchy
    DataType out_type = DataType::INVALID;
    if (op.tag == Op::FUNCTION && op.f.fn == "prng_step") {
      out_type = DataType::PRNG;
    } else if (op.tag == Op::FUNCTION && op.f.fn == "prng_state") {
      out_type = DataType::UINT32;
    } else if (op.tag == Op::FUNCTION && op.f.fn == "prng_value") {
      out_type = DataType::FLOAT32;
    } else if (op.tag == Op::FUNCTION && op.f.fn == "quantize") {
      out_type = DataType::INT8;
    } else {
      for (size_t i = 0; i < op.inputs.size(); i++) {
        const auto& s = op.inputs[i];
        // Skip condition booleans
        if (op.tag == Op::FUNCTION && op.f.fn == "cond" && i == 0) {
          continue;
        }
        // Move types up the hierarchy
        DataType cur = vars->at(s).shape.type;
        IVLOG(4, "  Adding type " << to_string(cur));
        out_type = CommonSupertype(out_type, cur);
      }
    }
    if (out_type == DataType::INVALID) {
      throw error::FailedPrecondition("Output result type is ill-defined");
    }
    // Set output type
    IVLOG(4, "Derived type " << to_string(out_type));
    if (op.tag == Op::CONTRACTION) {
      // Handle the contraction case
      if (op.c.comb_op == CombinationOp::EQ) {
        // == is always type BOOLEAN
        out_type = DataType::BOOLEAN;
      }
      // Check we have proper output sizes
      if (op.c.output_size.size() != op.c.specs[0].sspec.size()) {
        throw std::runtime_error(
            str(boost::format("Mismatched output indices and output size. Indicies: %1%, Sizes: %2%") %
                op.c.specs[0].sspec.size() % op.c.output_size.size()));
      }
      // Get each size
      std::vector<size_t> dims;
      for (const auto& s : op.c.output_size) {
        dims.push_back(static_cast<size_t>(ExtractInteger(*vars, s)));
      }
      // Construct a simple shape + add it
      vars->emplace(op.output, Binding(SimpleShape(out_type, dims)));
      // Finalize the polynominals
      for (auto& ts : op.c.specs) {
        for (auto& ss : ts.sspec) {
          ts.spec.push_back(ss->Evaluate(*vars));
        }
        ts.sspec.clear();
      }
      // Finalize constraint sizes
      for (auto& cc : op.c.constraints) {
        // Bind the range
        int64_t range = ExtractInteger(*vars, cc.range);
        // Bind the polynomial
        Polynomial<Rational> poly = cc.poly->Evaluate(*vars);
        // Update the concrete range
        cc.bound = RangeConstraint(poly, range);
        cc.poly = SymbolicPolynomialPtr();
        cc.range = "";
      }
      IVLOG(4, "ContractionOp " << to_string(op) << " produces dims=" << to_string(dims));
    } else {
      // Handle the function case
      if (op.f.fn.substr(0, 7) == "assert_") {
        const Binding& b = vars->at(op.inputs[0]);
        if (b.tag != Binding::ICONST) {
          throw std::runtime_error("Assert must be constant at bind time");
        }
        if (b.iconst == 0) {
          throw std::runtime_error("Assertion failure: " + op.f.fn);
        }
      }
      if (op.f.fn.substr(0, 3) == "cmp") {
        // All cmp_ ops make booleans
        out_type = DataType::BOOLEAN;
      }
      if (op.f.fn == "gather") {
        Binding ds = vars->at(op.inputs[0]);
        Binding is = vars->at(op.inputs[1]);
        if (ds.tag != Binding::TENSOR || is.tag != Binding::TENSOR) {
          throw std::runtime_error("Both inputs to gather must be tensors (not constants)");
        }
        if (ds.shape.dims.size() == 0) {
          throw std::runtime_error("Data input to gather must have at least 1 dimension");
        }
        if (is.shape.type != DataType::INT32) {
          // TODO: Handle other integer types?  Floor floats?
          throw std::runtime_error("Datatype for index input to gather must be INT32");
        }
        out_type = ds.shape.type;
        std::vector<size_t> out_shape;
        for (size_t i = 0; i < is.shape.dims.size(); i++) {
          out_shape.push_back(is.shape.dims[i].size);
        }
        for (size_t i = 1; i < ds.shape.dims.size(); i++) {
          out_shape.push_back(ds.shape.dims[i].size);
        }
        // Add the output type
        vars->emplace(op.output, Binding(SimpleShape(out_type, out_shape)));
        IVLOG(4, "FunctionOp " << to_string(op) << " produces dims=" << to_string(out_shape));
        continue;
      }
      if (op.f.fn == "scatter") {
        Binding es = vars->at(op.inputs[0]);
        Binding is = vars->at(op.inputs[1]);
        Binding vs = vars->at(op.inputs[2]);
        if (es.tag != Binding::TENSOR || is.tag != Binding::TENSOR || vs.tag != Binding::TENSOR) {
          throw std::runtime_error("All inputs to scatter must be tensors (not constants)");
        }
        if (es.shape.dims.size() == 0) {
          throw std::runtime_error("Data input to scatter must have at least 1 dimension");
        }
        if (is.shape.type != DataType::INT32) {
          // TODO: Handle other integer types?  Floor floats?
          throw std::runtime_error("Datatype for index input to scatter must be INT32");
        }
        out_type = es.shape.type;
        std::vector<size_t> out_shape;
        out_shape.push_back(vs.shape.dims[0].size);
        for (size_t i = is.shape.dims.size(); i < es.shape.dims.size(); i++) {
          out_shape.push_back(es.shape.dims[i].size);
        }
        // Add the output type
        vars->emplace(op.output, Binding(SimpleShape(out_type, out_shape)));
        IVLOG(4, "FunctionOp " << to_string(op) << " produces dims=" << to_string(out_shape));
        continue;
      }
      if (op.f.fn == "tuple") {
        std::vector<Binding> tuple;
        for (const auto& i : op.inputs) {
          tuple.push_back(vars->at(i));
        }
        vars->emplace(op.output, Binding(tuple));
        continue;
      }
      if (op.f.fn == "element") {
        if (op.inputs.size() != 2) {
          throw std::runtime_error("Element requires exactly two inputs.");
        }
        const Binding& it = vars->at(op.inputs[0]);
        if (it.tag != Binding::TUPLE) {
          throw std::runtime_error("Element requires it's first input to be a tuple");
        }
        if (vars->at(op.inputs[1]).tag != Binding::ICONST) {
          throw std::runtime_error("Element requires it's second input to be an integer");
        }
        int64_t elem = vars->at(op.inputs[1]).iconst;
        if (elem < 0 || elem >= it.tuple.size()) {
          throw std::runtime_error("Element requires it's tuple position to be in bound");
        }
        vars->emplace(op.output, it.tuple[elem]);
        continue;
      }
      if (op.f.fn == "shape") {
        if (op.inputs.size() != 1) {
          throw std::runtime_error("Shape requires exactly one input.");
        }
        Binding it = vars->at(op.inputs[0]);
        if (it.tag != Binding::TENSOR) {
          throw std::runtime_error("Shape requires one input that is a tensor");
        }
        out_type = DataType::INT32;
        std::vector<size_t> out_shape;
        out_shape.push_back(it.shape.dims.size());  // one dim, size of the rank of the input tensor
        vars->emplace(op.output, Binding(SimpleShape(out_type, out_shape)));
        continue;
      }
      if (op.f.fn == "reshape") {
        if (op.inputs.size() < 1) {
          throw std::runtime_error("Reshape requires at least one input.");
        }
        const Binding& it = vars->at(op.inputs[0]);
        if (it.tag != Binding::TENSOR) {
          throw std::runtime_error("Reshape requires one input that is a tensor");
        }
        std::vector<size_t> sizes;
        for (size_t i = 1; i < op.inputs.size(); i++) {
          if (vars->at(op.inputs[i]).tag != Binding::ICONST) {
            throw std::runtime_error("Additional parameters to reshape call must be constant integers");
          }
          sizes.push_back(vars->at(op.inputs[i]).iconst);
        }
        vars->emplace(op.output, Binding(SimpleShape(it.shape.type, sizes)));
        continue;
      }
      if (op.f.fn == "prng_step") {
        if (op.inputs.size() < 1) {
          throw std::runtime_error("prng_step must have at least one parameter");
        }
        // Valididate PRNG state size
        TensorShape prng_shape = vars->at(op.inputs[0]).shape;
        if (!(prng_shape == SimpleShape(DataType::UINT32, {3, k_rng_size}))) {
          throw std::runtime_error("Invalid PRNG state tensor");
        }
        // Get the output shape sizes
        std::vector<size_t> sizes;
        for (size_t i = 1; i < op.inputs.size(); i++) {
          if (vars->at(op.inputs[i]).tag != Binding::ICONST) {
            throw std::runtime_error("Additional parameters to PRNG call must be constant integers");
          }
          sizes.push_back(vars->at(op.inputs[i]).iconst);
        }
        vars->emplace(op.output, Binding(SimpleShape(DataType::PRNG, sizes)));
        continue;
      }
      if (op.f.fn == "prng_state" || op.f.fn == "prng_value") {
        if (op.inputs.size() != 1) {
          throw std::runtime_error("prng_state must have exactly one parameter");
        }
        TensorShape prng_shape = vars->at(op.inputs[0]).shape;
        if (prng_shape.type != DataType::PRNG) {
          throw std::runtime_error("prng_state must have a prng_step output as input");
        }
        if (op.f.fn == "prng_state") {
          vars->emplace(op.output, Binding(SimpleShape(DataType::UINT32, {3, k_rng_size})));
        } else {
          vars->emplace(op.output, Binding(TensorShape(DataType::FLOAT32, prng_shape.dims)));
        }
        continue;
      }

      if (op.f.fn.substr(0, 3) == "as_") {
        Binding val = vars->at(op.inputs[0]);
        Binding bitbind = vars->at(op.inputs[1]);
        // Bits must be bound to an integer constant
        if (bitbind.tag != Binding::ICONST) {
          throw std::runtime_error("Cast width must be an integer constant");
        }
        int64_t bits = bitbind.iconst;
        const std::string typefamily = op.f.fn.substr(3);
        if ("float" == typefamily) {
          switch (bits) {
            case 16:
              out_type = DataType::FLOAT16;
              break;
            case 32:
              out_type = DataType::FLOAT32;
              break;
            case 64:
              out_type = DataType::FLOAT64;
              break;
            default:
              throw std::runtime_error("Float width must be 16, 32, or 64");
          }
        } else if ("int" == typefamily) {
          switch (bits) {
            case 8:
              out_type = DataType::INT8;
              break;
            case 16:
              out_type = DataType::INT16;
              break;
            case 32:
              out_type = DataType::INT32;
              break;
            case 64:
              out_type = DataType::INT64;
              break;
            default:
              throw std::runtime_error("Int width must be 8, 16, 32, or 64");
          }
        } else if ("uint" == typefamily) {
          switch (bits) {
            case 8:
              out_type = DataType::UINT8;
              break;
            case 16:
              out_type = DataType::UINT16;
              break;
            case 32:
              out_type = DataType::UINT32;
              break;
            case 64:
              out_type = DataType::UINT64;
              break;
            default:
              throw std::runtime_error("UInt width must be 8, 16, 32, or 64");
          }
        }
        // compute out_type from the function name and possibly inputs[1]
        std::vector<size_t> out_shape;
        for (size_t i = 0; i < val.shape.dims.size(); ++i) {
          out_shape.push_back(val.shape.dims[i].size);
        }
        vars->emplace(op.output, Binding(SimpleShape(out_type, out_shape)));
        continue;
      }

      if (op.f.fn == "index") {
        Binding val = vars->at(op.inputs[0]);
        Binding idx_num_var = vars->at(op.inputs[1]);
        // Index must be bound to an integer constant
        if (idx_num_var.tag != Binding::ICONST) {
          throw std::runtime_error("Index number must be an integer constant");
        }
        int64_t idx_num = idx_num_var.iconst;
        // compute out_type from the function name and possibly inputs[1]
        std::vector<size_t> out_shape;
        for (size_t i = 0; i < val.shape.dims.size(); ++i) {
          out_shape.push_back(val.shape.dims[i].size);
        }
        vars->emplace(op.output, Binding(SimpleShape(DataType::INT32, out_shape)));
        continue;
      }

      // Check if all inputs are constants, and also extract as doubles
      bool all_const = true;
      std::vector<double> dins;
      for (const auto& s : op.inputs) {
        switch (vars->at(s).tag) {
          case Binding::TENSOR:
          case Binding::TUPLE:
            all_const = false;
            break;
          case Binding::ICONST:
            dins.push_back(static_cast<double>(vars->at(s).iconst));
            break;
          case Binding::FCONST:
            dins.push_back(vars->at(s).fconst);
            break;
        }
      }

      // If it's a constant, do the propagation
      if (all_const) {
        double r = ConstantPropagate(op.f.fn, dins);
        op.inputs.resize(1);
        if (is_float(out_type)) {
          vars->emplace(op.output, Binding(r, g_floatx));
          op.f.fn = "fconst";
          op.inputs[0] = DoubleToString(r);
        } else {
          vars->emplace(op.output, Binding(static_cast<int64_t>(r)));
          op.f.fn = "iconst";
          op.inputs[0] = std::to_string(static_cast<int64_t>(r));
        }
        op.tag = Op::CONSTANT;
        continue;
      }
      if (!all_const && op.f.fn == "mod") {
        throw std::runtime_error("Modulus only allowed in bind time code");
      }

      // Validate all sizes either match, are broadcast-compatible, or are 0-dim
      std::vector<TensorDimension> dims;
      std::string dims_source;
      bool did_broadcast = false;
      for (const std::string& in : op.inputs) {
        did_broadcast = BroadcastTo(&dims, &dims_source, vars->at(in).shape.dims, in) || did_broadcast;
      }
      if (did_broadcast) {
        // Recompute strides in dims.
        size_t stride = 1;
        for (auto it = dims.rbegin(); it != dims.rend(); ++it) {
          it->stride = stride;
          stride *= it->size;
        }
      }
      // Add the output type
      vars->emplace(op.output, Binding(TensorShape(out_type, dims)));
      IVLOG(4, "FunctionOp " << to_string(op) << " produces dims=" << to_string(dims));
    }
  }
}

void OptimizeProgram(Program* p, const std::set<std::string>& inputs, const std::set<std::string>& outputs,
                     const Bindings& vars) {
  // Figure out where variables are defined, and also setup identity mappings
  // IVLOG(1, "Pre optimize:\n"  << to_string(*p));
  std::map<std::string, size_t> defs;
  std::map<std::string, std::string> first_def;
  for (size_t i = 0; i < p->ops.size(); i++) {
    defs[p->ops[i].output] = i;
    if (p->ops[i].tag == Op::FUNCTION && p->ops[i].f.fn == "ident" && !outputs.count(p->ops[i].output)) {
      std::string first = p->ops[i].inputs[0];
      if (first_def.count(first)) {
        first = first_def.at(first);
      }
      first_def[p->ops[i].output] = first;
    }
    // Pull 'elements' through from original tuple
    if (p->ops[i].tag == Op::FUNCTION && p->ops[i].f.fn == "element") {
      const std::string& tup_name = p->ops[i].inputs[0];
      const std::string& elem_name = p->ops[i].inputs[1];
      size_t elem = vars.at(elem_name).iconst;
      const Op& tup_op = p->ops[defs.at(tup_name)];
      first_def[p->ops[i].output] = tup_op.inputs[elem];
    }
  }
  // IVLOG(1, "Identity backrefs" << first_def);
  // Backtrack from outputs till we hit inputs or constants
  std::queue<std::string> to_proc;
  std::set<std::string> keep;
  for (const auto& s : outputs) {
    if (vars.at(s).tag == Binding::TENSOR && !vars.at(s).shape.elem_size()) {
      // This is a zero-sized tensor; we don't need to generate it.
      continue;
    }
    keep.insert(s);
    to_proc.push(s);
  }
  auto deident = [&first_def](std::string& s) {
    if (first_def.count(s)) {
      s = first_def.at(s);
    }
  };
  while (!to_proc.empty()) {
    std::string s = to_proc.front();
    to_proc.pop();
    Op& op = p->ops[defs[s]];
    if (op.tag == Op::CONSTANT) {
      continue;
    }
    for (std::string& i : op.inputs) {
      deident(i);
      if (keep.count(i) || inputs.count(i)) {
        continue;
      }
      if (vars.at(i).tag == Binding::TENSOR && !vars.at(i).shape.elem_size()) {
        // This is a zero-sized tensor; we don't need to generate it.
        continue;
      }
      keep.insert(i);
      to_proc.push(i);
    }
    if (op.tag != Op::CONTRACTION) {
      continue;
    }
    if (op.c.use_default != "") {
      std::string& i = op.c.use_default;
      deident(i);
      if (!keep.count(i) && !inputs.count(i)) {
        keep.insert(i);
        to_proc.push(i);
      }
    }
    for (auto& s : op.c.output_size) {
      deident(s);
    }
    for (auto& s : op.c.specs) {
      deident(s.id);
    }
  }
  // IVLOG(1, "Replaced program:\n" << first_def);
  // Remove needless ops (TODO: do common subexpr elimination)
  std::vector<Op> new_ops;
  for (const Op& op : p->ops) {
    if (keep.count(op.output)) {
      new_ops.push_back(op);
    }
  }
  p->ops = new_ops;
}

Bindings BindProgram(Program* p, const ShapeMap& inputs, const ShapeMap& outputs) {
  // Copy input shapes into vars, also track names for OptimizeProgram
  Bindings vars;
  std::set<std::string> input_vars;
  std::set<std::string> output_vars;
  for (const auto& kvp : inputs) {
    // TODO: do we need to check for float and use floatx here?
    vars.emplace(kvp.first, Binding(kvp.second));
    input_vars.insert(kvp.first);
  }
  for (const auto& kvp : outputs) {
    if (!kvp.second.elem_size()) {
      continue;
    }
    output_vars.insert(kvp.first);
  }
  // Do typing
  TypeCheck(p, &vars);
  IVLOG(3, "After typecheck: " << p->ops);
  IVLOG(3, "Types:: " << vars);
  // Verify outputs match
  for (const auto& kvp : outputs) {
    if (!kvp.second.elem_size()) {
      continue;
    }
    auto it = vars.find(kvp.first);
    if (it == vars.end()) {
      throw std::runtime_error("No type deduced for output " + kvp.first);
    }
    std::vector<TensorDimension> dims{it->second.shape.dims};
    std::string src = "program variable";
    BroadcastTo(&dims, &src, kvp.second.dims, "program input");
    // if (!(kvp.second == it->second.shape)) {
    //   IVLOG(1, "Shape mismatch: " << kvp.first << " " << kvp.second << " " << it->second);
    //   throw std::runtime_error("Mismatched output type");
    // }
  }
  // Finally, run program 'optimization' pass
  OptimizeProgram(p, input_vars, output_vars, vars);
  IVLOG(3, "After optimize: " << p->ops);

  return vars;
}

DataType CommonSupertype(DataType left, DataType right) {
  DataType out = left;
  if (is_float(right) != is_float(left)) {
    if (is_float(right)) {
      out = right;
    }
  } else {
    // TODO: This is a bit primitive; for example, it will pick
    // the first of "int32" or "float32".  We may want to make it
    // a bit more sophisticated.
    if (bit_width(right) > bit_width(left)) {
      out = right;
    }
  }
  return out;
}

}  // namespace lang
}  // namespace tile
}  // namespace vertexai


//
//  main.cpp
//  PhyloAcc
//
//  Created by hzr on 3/8/16.
//  Copyright © 2016 hzr. All rights reserved.
//
#include <dirent.h>
#include <stdio.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <algorithm>
#include <iomanip>
//#include <omp.h> //libiomp/
#include <string>
#include <armadillo>
#include "profile.h"
#include "newick.h"
#include "newick2.h"
#include "bpp.hpp"
#include "bpp_c.hpp"
#include "utils.h"
#include <math.h>
#include <gsl/gsl_errno.h>

using namespace std;
using namespace arma;

// parameters, input and output files paths
string params_path;
string phytree_path;
string align_path;
string output_path="";
string output_path2="";
string segment_path;
string id_path="";
string result_prefix="test";
string tree_coal_unit="";

//string refspecies;
string outgroup;
string targetspecies;
string conservegroup; // can't be missing in more than 50%
string deepcoal_species = ""; //species in which deep coal can happen
double conserve_prop = 0.8;

int num_thread = 1;

// running parameters
int num_burn = 200;         // num of burn-in updates is num_burn * num_thin
int num_mcmc = 800;         // num of MCMC updates, num_burn * num_mcmc
int num_thin = 1;    // num of updates between two samples, adaptive frequency = 50 * num_thin
int num_chain; // outer loop of updates Q matrix and hyperparameter of substitution rates

double prep_lrate = 0.5;
double prep_lrate2 = 0.0; //0.1
double prep_grate = 0.8; // initalization

double prior_lrate2_a = 0.0,prior_lrate2_b = 1 ; // beta prior for lrate2, 0.5
double prior_lrate_a = 1 ,prior_lrate_b = 1 ; // beta prior for lrate, 1,9
double prior_grate_a = 1,prior_grate_b = 1; // beta prior for grate, 3,1

double ratio0 = 0.5; //initial conserved rate,0.5
double ratio1 = 1; // initial accelerated rate
double missing_thres = 0.8;

double nprior_a = 10, nprior_b = 0.2;  //around 1
double cprior_a = 5, cprior_b = 0.04;  //around ratio
int ropt = 1;
double cub = 1;
double nlb = 0.6;

int batch = -1 ;
int seed = 1;
int seed2 = 1;
double indel = 0;
double indel2;
bool sample_indel = 0;
bool sample_hyper = false;
char gapchar = '-';
bool verbose = 0;
double consToMis = 0.5; //for simulation, 0.01;
int block = 15; // 25; //90;
bool prune=0;
double revgap=0.9;  // 1
int min_length = 50;
bool WL = true;
bool simulate=false;
bool verboseGT = 1; //Han*: output trace_genetrees
double br_sample_cutoff = 10.0;
double theta_cutoff = 1.0;

//Han*: add Dirichlet prior for stationary distribution pi:
//vector<double> prior_dir_par(4,10);
vector<double> prior_dir_par(2,10);  //Han: it's beta prior for pi. Keep the old name.

// load the program parameters
void LoadParams(int argc, char* argv[])
{
    cout << "Loading input data and running parameters......" << endl;

    if (argc > 1)
        params_path = string(argv[1]);
    else
        params_path = "/Users/zhiruihu/Phylogenetics/test_mammalData/param_chr2_Long2.txt"; //params2

    cout << "Loading program configurations from " << params_path << "......" <<endl;

    const int BUFF_SIZE = 1024;
    char line_buff[BUFF_SIZE];

    ifstream in_params(params_path.c_str());
    if (!in_params)
    {
        cerr << "Cannot open the parameters file: " << params_path.c_str() << endl;
        exit(1);
    }
    while(!in_params.eof())
    {
        in_params.getline(line_buff, BUFF_SIZE);
        istringstream line_stream(line_buff);
        string tmp; line_stream >> tmp;

        // simulate or inference
        if (tmp=="SIMULATE")
            line_stream >> simulate;
        // input and output file paths
        else if (tmp=="PHYTREE_FILE")
            line_stream >> phytree_path;
        else if (tmp=="TREE_IN_COALESCENT_UNIT")
            line_stream >> tree_coal_unit;
        else if (tmp=="ALIGN_FILE")
            line_stream >> align_path;
        else if (tmp=="SEG_FILE")
            line_stream >> segment_path;
        else if (tmp=="ID_FILE")
            line_stream >> id_path;
        else if (tmp=="BATCH")
            line_stream >> batch;
        else if (tmp=="RESULT_FOLDER")
            line_stream >> output_path;
        else if (tmp=="PREFIX")
            line_stream >> result_prefix;
        else if (tmp=="SEED")
            line_stream >> seed;
        else if (tmp=="SEEDS")
            line_stream >> seed2;
        else if (tmp=="INIT_CONSERVE_RATE")
            line_stream >> ratio0;
        else if (tmp=="INIT_ACCE_RATE")
            line_stream >> ratio1;
        else if (tmp=="CONSERVE_PRIOR_A")
            line_stream >> cprior_a;
        else if (tmp=="CONSERVE_PRIOR_B")
            line_stream >> cprior_b;
        else if (tmp=="ACCE_PRIOR_A")
            line_stream >> nprior_a;
        else if (tmp=="ACCE_PRIOR_B")
            line_stream >> nprior_b;
        else if (tmp=="ROPT")
            line_stream >> ropt;
        else if (tmp=="CUB")
            line_stream >> cub;
        else if (tmp=="NLB")
            line_stream >> nlb;

        // running parameters
        else if (tmp=="BURNIN")
            line_stream >> num_burn;
        else if (tmp=="MCMC")
            line_stream >> num_mcmc;
        else if (tmp=="THIN") //ADAPT_FREQ
            line_stream >> num_thin;
        else if (tmp=="INIT_LRATE")
            line_stream >> prep_lrate;
        else if (tmp=="INIT_GRATE")
            line_stream >> prep_grate;
        else if (tmp=="HYPER_LRATE_A")
            line_stream >> prior_lrate_a;
        else if (tmp=="HYPER_LRATE_B")
            line_stream >> prior_lrate_b;
        else if (tmp=="HYPER_GRATE_A")
            line_stream >> prior_grate_a;
        else if (tmp=="HYPER_GRATE_B")
            line_stream >> prior_grate_b;
        else if (tmp=="CHAIN")
            line_stream >> num_chain;
        else if (tmp=="WL")
            line_stream >> WL;
        else if (tmp=="BLK_WL")
            line_stream >> block;

        // constraint
        else if (tmp == "OUTGROUP")
            line_stream >> outgroup;
        else if (tmp == "TARGETSPECIES")
            line_stream >> targetspecies;
        else if (tmp == "CONSERVE")
            line_stream >> conservegroup;
        else if (tmp == "CONSERVE_PROP")
            line_stream >> conserve_prop;
        else if (tmp == "GAP_PROP")
            line_stream >> missing_thres;
        //else if (tmp == "REF")
         //   line_stream >> refspecies;
        else if (tmp == "CONSTOMIS")
            line_stream >> consToMis;
        else if (tmp=="BR_SAMPLE_THRESHOLD")
            line_stream >> br_sample_cutoff;
        else if (tmp=="THETA_CUTOFF")
            line_stream >> theta_cutoff;
        else if( tmp=="DEEP_COAL_BRANCH")
            line_stream >> deepcoal_species;

        // treat indel as additional character
        else if (tmp == "GAPCHAR")
            line_stream >> gapchar;
        else if (tmp == "PRUNE_TREE")
            line_stream >> prune;
        else if (tmp == "TRIM_GAP_PERCENT")
            line_stream >> revgap;
        else if (tmp == "MIN_LEN")
            line_stream >> min_length;
        else if (tmp == "INDEL") // not used
            line_stream >> indel;
        else if (tmp == "INDEL2") // not used
            line_stream >> indel2;
        else if(tmp == "SAMPLE_INDEL")  // not used
            line_stream >> sample_indel;
        else if(tmp == "SAMPLE_HYPER")
            line_stream >> sample_hyper;
        else if(tmp == "VERBOSE")
            line_stream >> verbose;
        else if(tmp == "VERBOSE_GENETREE") //Han*: output sampled gene trees
            line_stream>>verboseGT;
        else if(tmp == "NUM_THREAD")
            line_stream >> num_thread;
        else if(tmp != "")
            cout << "Unknown parameter: " << tmp <<endl;

    }

    // trimming file names
    phytree_path = strutils::trim(phytree_path, " \"\t\n");
    align_path = strutils::trim(align_path, " \"\t\n");
    output_path  = strutils::trim(output_path,  " \"\t\n");
    segment_path = strutils::trim(segment_path, " \"\t\n");
    tree_coal_unit = strutils::trim(tree_coal_unit, " \"\t\n");
}

bool DirectoryExists( string pzPath )
{
    if ( pzPath == "") return false;

    DIR *pDir;
    bool bExists = false;

    pDir = opendir (pzPath.c_str());

    if (pDir != NULL)
    {
        bExists = true;
        (void) closedir (pDir);
    }

    return bExists;
}

void DispParams(PhyloProf profile, int seed)
{
    double mean_seg_size = 0;
    for(unsigned int c=0; c<profile.C; c++)
        mean_seg_size += (double)(profile.element_pos[c][1] - profile.element_pos[c][0]) / profile.C;
    cout << "  # total length = " << profile.G << " (" << profile.C << ")" << ". # Species = " << profile.S << ". # elements = " << profile.C << ". Mean gene set size = " << mean_seg_size << "." << endl;
    cout << "# Burn-ins = " << num_burn*num_thin << ". # MCMC Updates = " << num_mcmc*num_thin << ". # thin = " << num_thin << ".  RND SEED = " << seed << "." << endl ; //
    cout << "# Threads = " << num_thread << endl << endl;
}

int main(int argc, char* argv[])
{
    time_t start = time(NULL);

    cout << std::fixed << setprecision(4);
    srand(time(NULL));

    // load the program parameters
    LoadParams(argc, argv);

    // check output path
    if(! DirectoryExists(output_path))
    {
    	cout << "output path doesn't exist or empty!" << endl;
    	return 1;
    }

    // load the phylogenetic profile
    PhyloProf profile = LoadPhyloProfiles(align_path,segment_path);
    // init and display the running parameters
    DispParams(profile, seed);

    PhyloTree_theta tree2;
    if(tree_coal_unit !=""){
        tree2 = LoadPhyloTree_theta(tree_coal_unit);
    }else{
        cout<<"please input a phylogengy with branch length in coalescent unit."<<endl;
        return 1;
    }

    // load the phylogenetic tree
    PhyloTree phytree = LoadPhyloTree(phytree_path); //Han: .subs_rate contains Q

    //get thetas
    double theta_cum=0;
    int count_cum=0;
    int N = phytree.nodes_names.size();
    vector<int> pos_cum=vector<int> (N,0);
    for(int i=0; i<(N-1); i++){
        if(phytree.nodes_names[i]!=tree2.nodes_names[i]){
            cout<<"i="<<i<<". tree1_name="<<phytree.nodes_names[i]<<", tree2_name="<<tree2.nodes_names[i]<<endl;
            cerr<<"two trees do not have the same topology"<<endl;
            exit(1);
        }else{
            if(i<phytree.S){
                phytree.thetas[i]=0;
            }else{
                //cout<<"\ni="<<i<<": ";
                if((tree2.distances[i]!=1.0) && (tree2.distances[i]!= 7.0) && (tree2.distances[i]!= 0)){
                    phytree.thetas[i]=2*phytree.distances[i]/tree2.distances[i];
                    if(phytree.thetas[i]>= theta_cutoff){
                        pos_cum[i]=1;
                    }else{
                        theta_cum+=phytree.thetas[i];
                        count_cum+=1;
                    }
                }else{
                    pos_cum[i]=1;
                }
            }
        }
    }
    double theta_aver=theta_cum/count_cum;
    pos_cum[N-1]=1;
    for(int i=phytree.S; i<N; i++){
        if(pos_cum[i]==1) phytree.thetas[i]=theta_aver;
        //cout<<"node "<<phytree.nodes_names[i]<<" theta="<<phytree.thetas[i]<<endl;
    }

    // create and init the BPP object
    //BPP bpp(0, profile, phytree, output_path, targetspecies, outgroup, conserve_prop, conservegroup, ratio0, ratio1, ropt, cub, nlb, nprior_a, nprior_b, cprior_a, cprior_b, seed, seed2, prep_grate, prep_lrate, prep_lrate2, prior_grate_a, prior_grate_b,prior_lrate_a, prior_lrate_b,prior_lrate2_a, prior_lrate2_b,  indel, indel2, missing_thres, sample_indel,prior_dir_par, br_sample_cutoff);
    BPP bpp(0, profile, phytree, output_path, targetspecies, outgroup, conserve_prop, conservegroup, ratio0, ratio1, ropt, cub, nlb, nprior_a, nprior_b, cprior_a, cprior_b, seed, seed2, prep_grate, prep_lrate, prep_lrate2, prior_grate_a, prior_grate_b,prior_lrate_a, prior_lrate_b,prior_lrate2_a, prior_lrate2_b,  indel, indel2, missing_thres, sample_indel,prior_dir_par, br_sample_cutoff, deepcoal_species);

    //initialize the MCMC sampling
    bpp.InitMCMC(num_burn, num_mcmc, num_thin);

    output_path = output_path + "/" + result_prefix ;
    output_path2 = output_path;
    string outpath_Z0 = output_path + "_rate_postZ_M" +to_string(0) +".txt";
    string outpath_Z1 = output_path + "_rate_postZ_M" +to_string(1) +".txt";
    string outpath_Z2 = output_path + "_rate_postZ_M" +to_string(2) +".txt";

    string outpath_hyper = output_path+"_hyper.txt";
    ofstream out_hyper(outpath_hyper.c_str());
    out_hyper << "iter\tnprior_a\tnprior_b\tcprior_a\tcprior_b\tprior_l_a\tprior_l_b\tprior_g_a\tprior_g_b\n";
    out_hyper << 0 << "\t"<< nprior_a<< "\t"<< nprior_b <<"\t"<< cprior_a << "\t"<< cprior_b << "\t"<< prior_lrate_a << "\t"<< prior_lrate_b << "\t"<< prior_grate_a << "\t"<< prior_grate_b <<endl;

    ofstream out_lik;
    if(sample_hyper) {
        string outpath_elem = output_path+ "_elem_lik.txt";
        out_lik.open(outpath_elem.c_str());
        out_lik.precision(8);
        out_lik << "No.\tID\tloglik_Full\tloglik_Max"<<endl;
    }

    ofstream out_Z0(outpath_Z0.c_str());
    ofstream out_Z1(outpath_Z1.c_str());
    ofstream out_Z2(outpath_Z2.c_str());

    // output species name
    string species_name = output_path+"_species_names.txt";
    ofstream out_species(species_name.c_str());

    out_Z0 << "No.\tn_rate\tc_rate\tg_rate\tl_rate\tl2_rate"; out_Z1 << "No.\tn_rate\tc_rate\tg_rate\tl_rate\tl2_rate"; out_Z2 << "No.\tn_rate\tc_rate\tg_rate\tl_rate\tl2_rate";
    for(int s=0; s<bpp.N;s++){
         for(int k=0;k<4;k++){
            out_Z0 <<"\t"<<bpp.nodes_names[s]<<"_"<<k;
            out_Z1 <<"\t"<<bpp.nodes_names[s]<<"_"<<k;
            out_Z2 <<"\t"<<bpp.nodes_names[s]<<"_"<<k;
         }
         out_species << bpp.nodes_names[s] << endl;
    }
	out_Z0 <<endl; out_Z1 <<endl; out_Z2 <<endl;

    out_species.close();

    // output tree
    outpath_Z0 = output_path + "_tree_M" +to_string(0) +".txt";
    outpath_Z1 = output_path + "_tree_M" +to_string(1) +".txt";
    outpath_Z2 = output_path + "_tree_M" +to_string(2) +".txt";

    ofstream out_tree0(outpath_Z0.c_str());
    ofstream out_tree1(outpath_Z1.c_str());
    ofstream out_tree2(outpath_Z2.c_str());

    out_tree0 << "No.\tprop\tgenetree\n";
    out_tree1 << "No.\tprop\tgenetree\n";
    out_tree2 << "No.\tprop\tgenetree\n";

    double lrate_prop = 0.5, grate_prop = 0.5;

    vector<int> ids;
    if(id_path=="")
    {
        if(batch==-1)
        {
          for(int c =0;c<500;c++) //bpp.C,
          {
            ids.push_back(c);
          }
        }else{
          int temp = ceil(bpp.C/3);
          for(int c =batch*temp ;c< (batch+1)*temp;c++)
          {
            if(c >= bpp.C) break;
            ids.push_back(c);
          }

        }

    }else{
        ifstream in_params(id_path.c_str());
        if (!in_params)
        {
            cerr << "Cannot open the id file: " << id_path.c_str() << endl;
            exit(1);
        }
        while(!in_params.eof())
        {
            const int BUFF_SIZE = 1024;
            char line_buff[BUFF_SIZE];
            in_params.getline(line_buff, BUFF_SIZE);
            istringstream line_stream(line_buff);
            string tmp; line_stream >> tmp;
            tmp = strutils::trim(tmp);
            if(tmp=="") continue;
            ids.push_back(atoi(tmp.c_str()));

        }

    }

    cout << ids.size() << " elements to be computed" << endl;

    if(sample_hyper)
    {
        for(int iter  =0; iter<num_chain; iter++)
        {
            cout << "Running MCMC chain " << iter +1 << " ..." << endl;
            // Gibbs sampling
            #pragma omp parallel for schedule (guided) num_threads(num_thread)
            for(std::size_t i = 0; i < ids.size(); i++ )
            {
                int c = ids[i];
                bool filter = false;

                try{
                    BPP_C bppc(c, profile, bpp, gapchar, missing_thres, filter, verbose, consToMis, block, prune, revgap, min_length);  // for individual element
                    if(filter) {
                        if(verbose) cerr << "filter: "<< c <<endl;
                        continue;
                    }

                    bppc.initMCMC(0,5,bpp,1,prune);
                    bppc.Gibbs(0, 4, bpp,out_Z2,output_path,output_path2,1, true, sample_hyper, lrate_prop, grate_prop, false);  // Gibbs run to get Z for each element

                    if(bppc.verbose || bppc.failure) bppc.Output_sampling(iter, output_path2, bpp, 2);
                    bppc.Output_init(output_path,output_path2,bpp,out_Z2, out_tree2, bppc.verbose); //sort rates!!

                }catch (exception& e){
                    cout << c << " Standard exception: " << e.what() << endl;
                }
            }
            bpp.sample_hyperparam(iter, ids, out_hyper);
            bpp.Output_init0(profile,out_lik, ids);

        }
    }else if(simulate)
    {
        for(std::size_t i = 0; i < ids.size(); i++ )
        {
            int c = ids[i];
            bool filter = false;
            try{
                // accelerate in target species
                BPP_C bppc(c, profile, bpp, gapchar, missing_thres, filter, verbose, consToMis, block, prune, revgap, min_length);  // for individual element
                bppc.simulate(bpp, profile, gapchar,prune);
            }catch (exception& e){
                cout << c << " Standard exception: " << e.what() << endl;
            }
        }

        //write out simulate sequence
        bpp.Output_simu(profile, output_path, ids.size());
    }else{
        // Gibbs sampling
        #pragma omp parallel for schedule (guided) num_threads(num_thread)
        for(std::size_t i = 0; i < ids.size(); i++ ) //
        {
            int c = ids[i];
            bool filter = false;

            try{

                BPP_C bppc(c, profile, bpp, gapchar, missing_thres, filter, verbose, verboseGT, consToMis, block, prune, revgap, min_length);  // for individual element
                cout<<"element "<<to_string(c)<<", number of base pair="<<to_string(bppc.GG)<<endl;
                if(filter) {
                  if(verbose) cerr << "filter: "<< c <<endl;
                  continue;
                }

                int tot = 0;
                if(bppc.idblk_count==0){
                    double nblk=(double)(bppc.GG - 15)/block;
                    int nblk2=(bppc.GG - 15)/block;
                    if( (nblk- nblk2)< ((double) 1.0/3.0)){
                        tot=nblk2;
                    }else{
                        tot=nblk2+1;
                    }
                    //tot=ceil((double)(bppc.GG - 15)/block);  //first bp length is always 15: len={0,15}. If change 15, change in Gibbs as well.
                    //have to make sure 15<=block. So better change 15 to block
                }else{
                    double nblk = (double)(bppc.GG - bppc.idblk_count-15)/block;
                    int nblk2 = (bppc.GG - bppc.idblk_count-15)/block;
                    if((nblk- nblk2)< ((double) 1.0/3.0)){
                        tot = 1+nblk2;
                    }else{
                        tot = 2+nblk2;
                    }
                    //tot=1+ceil((double)(bppc.GG - bppc.idblk_count-15)/block);
                }
                // int tot=ceil((double)(bppc.GG - 15)/block);
                //cout<<"tot ="<<tot<<endl;

                //null model
                cout<<"start null model\n";
                for (int iter = 0; iter <= tot; iter++)
                {
                    // cout<<"start initMCMC"<<endl;
                    bppc.initMCMC(iter, tot, bpp, 0, prune, false);

                    // cout<<"start Gibbs"<<endl;
                    bppc.Gibbs(iter, tot, bpp, out_Z0, output_path, output_path2, 0, true, sample_hyper, lrate_prop, grate_prop, WL); // Gibbs run to get Z for each element
                    if (bppc.verbose || bppc.failure)
                    {
                        bppc.Output_sampling(iter, output_path2, bpp, 0);
                        bppc.Output_tree(iter, output_path2, bpp, 0);
                    }
                }
                bppc.Output_init(output_path,output_path2,bpp,out_Z0, out_tree0, bppc.verbose); //sort rates!!, posterior median of nrate and crate; posterior mean of Z

                // //res model
                cout<<"start restricted model\n";
                for (int iter = 0; iter <= tot; iter++)
                {
                    bppc.initMCMC(iter, tot, bpp, 2, prune, false);
                    // cout<<"start Gibbs"<<endl;
                    bppc.Gibbs(iter, tot, bpp, out_Z2, output_path, output_path2, 2, true, sample_hyper, lrate_prop, grate_prop, WL);
                    if (bppc.verbose || bppc.failure)
                    {
                        bppc.Output_sampling(iter, output_path2, bpp, 1);
                        bppc.Output_tree(iter, output_path2, bpp, 1);
                    }
                }
                bppc.Output_init(output_path,output_path2,bpp,out_Z1, out_tree1, bppc.verbose);

                // full model
                cout<<"start full model\n";
                for (int iter = 0; iter <= tot; iter++)
                {
                    bppc.initMCMC(iter, tot, bpp, 1, prune, false); // not constrain log_prob_back
                    bppc.Gibbs(iter, tot, bpp, out_Z2, output_path, output_path2, 1, true, sample_hyper, lrate_prop, grate_prop, WL);
                    if (bppc.verbose || bppc.failure)
                    {
                        bppc.Output_sampling(iter, output_path2, bpp, 2);
                        bppc.Output_tree(iter, output_path2, bpp, 2);
                    }
                }
                bppc.Output_init(output_path,output_path2,bpp,out_Z2, out_tree2, bppc.verbose);

                cout << c << "\t" << bpp.log_liks_WL[0][c] <<"\t" <<  bpp.log_liks_WL[2][c] <<"\t" <<  bpp.log_liks_WL[1][c] <<endl;
                cout<<"\t" << bpp.log_liks_Z[0][c] << "\t" << bpp.log_liks_Z[2][c]<<"\t" << bpp.log_liks_Z[1][c] << endl;

            }catch (exception& e){
              cout << c << " Standard exception: " << e.what() << endl;
            }
      }

      bpp.Output_init(profile,output_path, ids);
    }

    out_Z0.close();
    out_Z1.close();
    out_Z2.close();
    out_hyper.close();
    out_lik.close();
    out_tree0.close();
    out_tree1.close();
    out_tree2.close();

    cout << endl << endl << "time used:  " << (time(NULL)-start)/60 << " min." << endl << endl;
    return 0;
}


#include "ompu/plug/game/drawer.hpp"

#include "ompu/ui/keyboard.hpp"
#include "ompu/music/key.hpp"

#include "saya/zed/blackhole.hpp"

#include <boost/format.hpp>
#include <boost/range/irange.hpp>
#include <boost/range/adaptor/indexed.hpp>

#include <boost/core/ignore_unused.hpp>

#include "JuceHeader.h"

namespace ompu { namespace plug { namespace game {

namespace {

template<std::size_t I>
std::nullptr_t sep_step_single(juce::Graphics& g, juce::Rectangle<float> const& rc)
{
    using ompu::ui::Keyboard;
    namespace kbs = ompu::ui::keyboards;

    using rcalc_type = ompu::geo::RatioCalculator<
        kbs::clearance_x_ratio_t<Keyboard, I>,
        std::ratio<0, 1>,
        typename Keyboard::clearance_ratio::w_ratio,
        typename Keyboard::clearance_ratio::h_ratio
    >;

    auto const rc_key = rcalc_type::rect_for_step<juce::Rectangle<float>>(
        rc.getX(), rc.getY(),
        rc.getWidth(), rc.getHeight()
    );
    g.fillRect(rc_key);

    return nullptr;
}

template<std::size_t... Is>
void sep_step(juce::Graphics& g, juce::Rectangle<float> const& rc, std::index_sequence<Is...>)
{
    saya::zed::blackhole{(sep_step_single<Is>(g, rc), 0)...};
}

namespace detail {

template<std::size_t I, class Enabled = void>
struct note_info
{
    static void draw(juce::Graphics&, juce::Rectangle<float> const&) {}
};

template<std::size_t I>
struct note_info<I, std::enable_if_t<midi::key_octave_index_v<I> == 0>>
{
    static void draw(juce::Graphics& g, juce::Rectangle<float> const& rc)
    {
        g.setColour(juce::Colour::fromRGBA(0, 0, 0, 200));
        g.drawText(
            "*",
            rc,
            juce::Justification::centredTop
        );
    }
};

template<std::size_t I>
struct note_info<I, std::enable_if_t<I == midi::Spec::A4>>
{
    static void draw(juce::Graphics& g, juce::Rectangle<float> const& rc)
    {
        g.setColour(juce::Colour::fromRGBA(0, 0, 0, 240));
        g.drawText(
            "A4",
            rc,
            juce::Justification::centredTop
        );
    }
};

} // detail

template<std::size_t I>
std::nullptr_t white_key_step_single(juce::Graphics& g, juce::Rectangle<float> const& rc, ompu::ui::Keyboard const& c)
{
    using ompu::ui::Keyboard;
    namespace kbs = ompu::ui::keyboards;

    using rcalc_type = ompu::geo::RatioCalculator<
        kbs::white_key_x_ratio_t<Keyboard, I>,
        std::ratio<0, 1>,
        typename Keyboard::white_key_ratio::w_ratio,
        typename Keyboard::white_key_ratio::h_ratio
    >;

    auto const rc_key = rcalc_type::rect_for_step<juce::Rectangle<float>>(
        rc.getX(), rc.getY(),
        rc.getWidth(), rc.getHeight()
    );

    auto const& key = c.nth_key<I>();

#if 0
    // (optimized out)
    g.setColour(juce::Colour::fromRGB(240, 240, 240));
    g.fillRect(rc_key);
#endif

    if (key.is_pressed()) {
        g.setColour(juce::Colour::fromRGBA(0, 0, 0, 20));
        g.drawRect(rc_key, rc_key.getWidth() * .1f);
    }

    // -------------------------------------
#if 0
    juce::Font const f{juce::Font::getDefaultSansSerifFontName(), 16, juce::Font::FontStyleFlags::bold};
    g.setFont(f);

    detail::note_info<kbs::to_midi_index_v<Keyboard, I>>::draw(
        g,
        decltype(rc_key){
            rc_key.getX(), rc_key.getY() + rc_key.getHeight() * .9f,
            rc_key.getWidth(), rc_key.getHeight()
        }
    );
#endif
    return nullptr;
}

template<std::size_t... Is>
void white_key_step(
    juce::Graphics& g, juce::Rectangle<float> const& rc, ompu::ui::Keyboard const& c,
    std::index_sequence<Is...>
) {
    saya::zed::blackhole{(white_key_step_single<Is>(g, rc, c), 0)...};
}

template<std::size_t I>
std::nullptr_t black_key_step_single(juce::Graphics& g, juce::Rectangle<float> const& rc, ompu::ui::Keyboard const& c)
{
    using ompu::ui::Keyboard;
    namespace kbs = ompu::ui::keyboards;

    using rcalc_type = ompu::geo::RatioCalculator<
        kbs::black_key_x_ratio_t<Keyboard, I>,
        std::ratio<0, 1>,
        typename Keyboard::black_key_ratio::w_ratio,
        typename Keyboard::black_key_ratio::h_ratio
    >;

    auto const rc_key = rcalc_type::rect_for_step<juce::Rectangle<float>>(
        rc.getX(), rc.getY(),
        rc.getWidth(), rc.getHeight()
    );
    g.setColour(juce::Colour::fromRGB(20, 20, 20));
    g.fillRect(rc_key);

    auto const& key = c.nth_key<I>();

    if (key.is_pressed()) {
        g.setColour(juce::Colour::fromRGBA(255, 255, 255, 200));
        g.drawRect(rc_key, rc_key.getWidth() * .1f);
    }

#if 0
    char const* note_name = nullptr;
    switch (ZeroI % 5) {
    case 0:
        note_name = "C#"; break;
    case 1:
        note_name = "D#"; break;
    case 2:
        note_name = "F#"; break;
    case 3:
        note_name = "G#"; break;
    case 4:
        note_name = "A#"; break;
    }

    auto const text_w = rc_key.getWidth() * 5.f;
    juce::Rectangle<float> const rc_text{
        rc_key.getX(),
        rc_key.getY() + rc_key.getHeight() * .2f,
        text_w,
        rc_key.getHeight() * .2f
    };

    g.setColour(juce::Colour::fromRGB(255, 255, 255));
    g.drawText(
        note_name,
        rc_text,
        juce::Justification::topLeft
    );
#endif

    return nullptr;
}

template<std::size_t... Is>
void black_key_step(
    juce::Graphics& g, juce::Rectangle<float> const& rc, ompu::ui::Keyboard const& c,
    std::index_sequence<Is...>
)
{
    saya::zed::blackhole{(black_key_step_single<Is>(g, rc, c), 0)...};
}

} // anon

void Drawer::partial(ompu::ui::Keyboard const& c) const
{
    using boost::adaptors::indexed;
    static auto const BLACK = juce::Colour::fromRGB(20, 20, 20);
    static auto const WHITE = juce::Colour::fromRGB(240, 240, 240);
    static auto const SEP_BLACK = juce::Colour::fromRGBA(0, 0, 0, 40);

    static float const ROOT_WIDTH = 1.f;
    static float const ROOT_HEIGHT = .15f;

    // -----------------------------------------

    juce::Rectangle<float> const rc_root{
        float(jc_->getPosition().x),
        jc_->getY() + jc_->getHeight() * (1 - ROOT_HEIGHT),
        float(jc_->getWidth()),
        jc_->getHeight() * ROOT_HEIGHT
    };

    // master bg
    g_->setColour(BLACK);
    g_->fillRect(rc_root);

    using ompu::ui::Keyboard;
    namespace kbs = ompu::ui::keyboards;

    auto const rc_key_root = rc_root.reduced(
        rc_root.proportionOfWidth(.05f),
        rc_root.proportionOfHeight(.02f)
    );

    g_->setColour(WHITE);
    g_->fillRect(rc_key_root);

    // white keys
    white_key_step(
        *g_, rc_key_root, c,
        kbs::make_white_key_sequence<Keyboard>()
    );

    // key seps
    g_->setColour(SEP_BLACK);
    sep_step(*g_, rc_key_root, kbs::make_clearance_sequence<Keyboard>{});

    // black keys
    black_key_step(
        *g_, rc_key_root, c,
        kbs::make_black_key_sequence<Keyboard>()
    );
}

}}} // ompu


#ifndef W_POT_EMBEM_H
#define W_POT_EMBEM_H

#include <armadillo>
#include <vector>

namespace willow { namespace pot {

extern void   libint2_init ();
extern void   libint2_end  ();

extern double embem (const arma::mat& pos,
		     const arma::ivec& Zs,
		     const arma::vec& Qs);

extern double  mbem (const arma::mat& pos,
		     const arma::ivec& Zs);

extern void esp_update (const arma::mat& pos,
			const arma::ivec& Zs,
			arma::vec& Qs);

} } // namespace willow::pot

#endif


#include <cctbx/boost_python/flex_fwd.h>
#include "../merge_equivalents.h"

#include <boost/python/class.hpp>
#include <boost/python/return_value_policy.hpp>
#include <boost/python/return_by_value.hpp>

namespace yamtbx { namespace dataproc { namespace crystfel { namespace boost_python {

namespace {

  struct merge_equivalents_crystfel_wrappers
  {
    typedef merge_equivalents_crystfel<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      typedef return_value_policy<return_by_value> rbv;
      class_<w_t>("merge_equivalents_crystfel")
	.def("add_observations", (void(w_t::*)(af::const_ref<cctbx::miller::index<> > const&,
					       af::const_ref<double> const&))&w_t::add_observations,
	     (arg("unmerged_indices"), arg("unmerged_data")))
	.def("add_observations", (void(w_t::*)(af::const_ref<cctbx::miller::index<> > const&,
					       af::const_ref<double> const&,
					       af::const_ref<double> const&))&w_t::add_observations,
	     (arg("unmerged_indices"), arg("unmerged_data"), arg("unmerged_sigmas")))
	.def("merge", &w_t::merge, arg("sigma")="population sigma")
        .add_property("indices", make_getter(&w_t::indices, rbv()))
        .add_property("data", make_getter(&w_t::data, rbv()))
        .add_property("sigmas", make_getter(&w_t::sigmas, rbv()))
        .add_property("redundancies", make_getter(&w_t::redundancies, rbv()))
      ;
    }
  };
} // namespace <anoymous>

  void wrap_merge_equivalents()
  {
    merge_equivalents_crystfel_wrappers::wrap();
  }

}}}} // namespace yamtbx::dataproc::crystfel::boost_python


/*
 * This is part of the fl library, a C++ Bayesian filtering library
 * (https://github.com/filtering-library)
 *
 * Copyright (c) 2015 Max Planck Society,
 * 				 Autonomous Motion Department,
 * 			     Institute for Intelligent Systems
 *
 * This Source Code Form is subject to the terms of the MIT License (MIT).
 * A copy of the license can be found in the LICENSE file distributed with this
 * source code.
 */

/**
 * \file multi_sensor_sigma_point_update_policy.hpp
 * \date August 2015
 * \author Jan Issac (jan.issac@gmail.com)
 */

#pragma once

#include <Eigen/Dense>

#include <fl/util/meta.hpp>
#include <fl/util/types.hpp>
#include <fl/util/traits.hpp>
#include <fl/util/descriptor.hpp>
#include <fl/model/sensor/joint_sensor_iid.hpp>
#include <fl/filter/gaussian/transform/point_set.hpp>
#include <fl/filter/gaussian/quadrature/sigma_point_quadrature.hpp>

namespace fl
{
// Forward declarations
template <typename...>
class MultiSensorSigmaPointUpdatePolizzle;

/**
 * \internal
 */
template <typename SigmaPointQuadrature, typename NonJoinSensor>
class MultiSensorSigmaPointUpdatePolizzle<SigmaPointQuadrature,
                                          NonJoinSensor>
{
    static_assert(
        std::is_base_of<internal::JointSensorIidType,
                        NonJoinSensor>::value,
        "\n\n\n"
        "====================================================================\n"
        "= Static Assert: You are using the wrong observation model type    =\n"
        "====================================================================\n"
        "  Observation model type must be a JointSensor<...>.      \n"
        "  For single observation model, use the regular Gaussian filter     \n"
        "  or the regular SigmaPointUpdatePolicy if you are specifying       \n"
        "  the update policy explicitly fo the GaussianFilter.               \n"
        "===================================================================="
        "\n");
};

template <typename SigmaPointQuadrature, typename MultipleOfLocalSensor>
class MultiSensorSigmaPointUpdatePolizzle<
    SigmaPointQuadrature,
    JointSensor<MultipleOfLocalSensor>>
    : public MultiSensorSigmaPointUpdatePolizzle<
          SigmaPointQuadrature,
          NonAdditive<JointSensor<MultipleOfLocalSensor>>>
{
};

template <typename SigmaPointQuadrature, typename MultipleOfLocalSensor>
class MultiSensorSigmaPointUpdatePolizzle<
    SigmaPointQuadrature,
    NonAdditive<JointSensor<MultipleOfLocalSensor>>>
    : public Descriptor
{
public:
    typedef JointSensor<MultipleOfLocalSensor> JointModel;
    typedef typename MultipleOfLocalSensor::Type FeatureSensor;
    typedef typename FeatureSensor::EmbeddedSensor BodyTailModel;
    typedef typename BodyTailModel::BodySensor BodyModel;
    typedef typename BodyTailModel::TailSensor TailModel;

    typedef typename JointModel::State State;
    typedef typename JointModel::Obsrv Obsrv;
    typedef typename JointModel::Noise Noise;

    typedef typename Traits<JointModel>::LocalObsrv LocalFeature;
    typedef Vector1d LocalObsrvNoise;

    template <typename Belief>
    void operator()(JointModel& obsrv_function,
                    const SigmaPointQuadrature& quadrature,
                    const Belief& prior_belief,
                    const Obsrv& y,
                    Belief& posterior_belief)
    {
        auto& feature_model = obsrv_function.local_sensor();
        auto& body_tail_model = feature_model.embedded_sensor();

        /* ------------------------------------------ */
        /* - Determine the number of quadrature     - */
        /* - points needed for the given quadrature - */
        /* - in conjunction with the joint Gaussian - */
        /* - p(State, LocalObsrvNoise)              - */
        /* ------------------------------------------ */
        enum : signed int
        {
            NumberOfPoints = SigmaPointQuadrature::number_of_points(
                JoinSizes<SizeOf<State>::Value,
                          SizeOf<LocalObsrvNoise>::Value>::Size)
        };

        /* ------------------------------------------ */
        /* - PointSets                              - */
        /* - [p_X, p_Q] ~ p(State, LocalObsrvNoise) - */
        /* ------------------------------------------ */
        PointSet<State, NumberOfPoints> p_X;
        PointSet<LocalObsrvNoise, NumberOfPoints> p_R;

        /* ------------------------------------------ */
        /* - Transform p(State, LocalObsrvNoise) to - */
        /* - point sets [p_X, p_Q]                  - */
        /* ------------------------------------------ */
        quadrature.transform_to_points(
            prior_belief,
            Gaussian<LocalObsrvNoise>(feature_model.noise_dimension()),
            p_X,
            p_R);

        auto mu_x = p_X.mean();
        auto X = p_X.centered_points();

        auto W = p_X.covariance_weights_vector().asDiagonal();
        auto c_xx = (X * W * X.transpose()).eval();
        auto c_xx_inv = c_xx.inverse().eval();

        auto C = c_xx_inv;
        auto D = State();
        D.setZero(mu_x.size());

        const int sensor_count = obsrv_function.count_local_models();
        const int dim_y = y.size() / sensor_count;

        auto h_body = [&](const State& x, const typename BodyModel::Noise& w)
        {
            return feature_model.feature_obsrv(
                body_tail_model.body_model().observation(x, w));
        };
        auto h_tail = [&](const State& x, const typename TailModel::Noise& w)
        {
            return feature_model.feature_obsrv(
                body_tail_model.tail_model().observation(x, w));
        };
        PointSet<LocalFeature, NumberOfPoints> p_Y_body;
        PointSet<LocalFeature, NumberOfPoints> p_Y_tail;

        for (int i = 0; i < sensor_count; ++i)
        {
            // validate sensor value, i.e. make sure it is finite
            if (!is_valid(y, i * dim_y, i * dim_y + dim_y)) continue;

            feature_model.id(i);

            /* ------------------------------------------ */
            /* - Integrate body                         - */
            /* ------------------------------------------ */

            quadrature.propagate_points(h_body, p_X, p_R, p_Y_body);
            auto mu_y_body = p_Y_body.mean();

            // validate sensor value, i.e. make sure it is finite
            if (!is_valid(mu_y_body, 0, dim_y)) continue;

            auto Y_body = p_Y_body.centered_points();
            auto c_yy_body = (Y_body * W * Y_body.transpose()).eval();
            auto c_xy_body = (X * W * Y_body.transpose()).eval();

            /* ------------------------------------------ */
            /* - Integrate tail                         - */
            /* ------------------------------------------ */
            quadrature.propagate_points(h_tail, p_X, p_R, p_Y_tail);
            auto mu_y_tail = p_Y_tail.mean();
            auto Y_tail = p_Y_tail.centered_points();
            auto c_yy_tail = (Y_tail * W * Y_tail.transpose()).eval();
            auto c_xy_tail = (X * W * Y_tail.transpose()).eval();

            /* ------------------------------------------ */
            /* - Fuse and center                        - */
            /* ------------------------------------------ */
            auto w = body_tail_model.tail_weight();
            auto mu_y = ((1.0 - w) * mu_y_body + w * mu_y_tail).eval();

            // non centered moments
            auto m_yy_body =
                (c_yy_body + mu_y_body * mu_y_body.transpose()).eval();
            auto m_yy_tail =
                (c_yy_tail + mu_y_tail * mu_y_tail.transpose()).eval();
            auto m_yy = ((1.0 - w) * m_yy_body + w * m_yy_tail).eval();

            // center
            auto c_yy = (m_yy - mu_y * mu_y.transpose()).eval();
            auto c_xy = ((1.0 - w) * c_xy_body + w * c_xy_tail).eval();

            auto c_yx = c_xy.transpose().eval();
            auto A_i = (c_yx * c_xx_inv).eval();
            auto c_yy_given_x = (c_yy - c_yx * c_xx_inv * c_xy).eval();
            auto innovation = (y.middleRows(i * dim_y, dim_y) - mu_y).eval();

            C += A_i.transpose() * solve(c_yy_given_x, A_i);
            D += A_i.transpose() * solve(c_yy_given_x, innovation);
        }

        /* ------------------------------------------ */
        /* - Update belief according to PAPER REF   - */
        /* ------------------------------------------ */
        posterior_belief.dimension(prior_belief.dimension());
        posterior_belief.covariance(C.inverse());
        posterior_belief.mean(mu_x + posterior_belief.covariance() * D);
    }

    virtual std::string name() const
    {
        return "MultiSensorSigmaPointUpdatePolizzle<" +
               this->list_arguments("SigmaPointQuadrature",
                                    "NonAdditive<SensorFunction>") +
               ">";
    }

    virtual std::string description() const
    {
        return "Multi-Sensor Sigma Point based filter update policy "
               "for joint observation model of multiple local observation "
               "models with non-additive noise.";
    }

private:
    /**
     * \brief Checks whether all vector components within the range (start, end)
     *        are finiate, i.e. not NAN nor Inf.
     */
    template <typename Vector>
    bool is_valid(Vector&& vector, int start, int end) const
    {
        for (int k = start; k < end; ++k)
        {
            if (!std::isfinite(vector(k))) return false;
        }

        return true;
    }
};
}


#pragma once

#include <polyfem/ElementAssemblyValues.hpp>
#include <polyfem/ElementBases.hpp>
#include <polyfem/AutodiffTypes.hpp>
#include <polyfem/Types.hpp>

#include <Eigen/Dense>
#include <vector>
#include <array>
#include <functional>

namespace polyfem
{
	constexpr int SMALL_N = POLYFEM_SMALL_N;
	constexpr int BIG_N = POLYFEM_BIG_N;

	Eigen::VectorXd gradient_from_energy(const int size, const int n_bases, const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da,
		const std::function<DScalar1<double, Eigen::Matrix<double, 6, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun6,
		const std::function<DScalar1<double, Eigen::Matrix<double, 8, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun8,
		const std::function<DScalar1<double, Eigen::Matrix<double, 12, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun12,
		const std::function<DScalar1<double, Eigen::Matrix<double, 18, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun18,
		const std::function<DScalar1<double, Eigen::Matrix<double, 24, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun24,
		const std::function<DScalar1<double, Eigen::Matrix<double, 30, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun30,
		const std::function<DScalar1<double, Eigen::Matrix<double, 60, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun60,
		const std::function<DScalar1<double, Eigen::Matrix<double, 81, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun81,
		const std::function<DScalar1<double, Eigen::Matrix<double, Eigen::Dynamic, 1, 0, SMALL_N, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &funN,
		const std::function<DScalar1<double, Eigen::Matrix<double, Eigen::Dynamic, 1, 0, 1000, 1>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &funBigN,
		const std::function<DScalar1<double, Eigen::VectorXd>				(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &funn
		);

	Eigen::MatrixXd hessian_from_energy(const int size, const int n_bases, const ElementAssemblyValues &vals, const Eigen::MatrixXd &displacement, const QuadratureVector &da,
		const std::function<DScalar2<double, Eigen::Matrix<double, 6, 1>, Eigen::Matrix<double, 6, 6>>		(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun6,
		const std::function<DScalar2<double, Eigen::Matrix<double, 8, 1>, Eigen::Matrix<double, 8, 8>>		(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun8,
		const std::function<DScalar2<double, Eigen::Matrix<double, 12, 1>, Eigen::Matrix<double, 12, 12>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun12,
		const std::function<DScalar2<double, Eigen::Matrix<double, 18, 1>, Eigen::Matrix<double, 18, 18>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun18,
		const std::function<DScalar2<double, Eigen::Matrix<double, 24, 1>, Eigen::Matrix<double, 24, 24>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun24,
		const std::function<DScalar2<double, Eigen::Matrix<double, 30, 1>, Eigen::Matrix<double, 30, 30>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun30,
		const std::function<DScalar2<double, Eigen::Matrix<double, 60, 1>, Eigen::Matrix<double, 60, 60>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun60,
		const std::function<DScalar2<double, Eigen::Matrix<double, 81, 1>, Eigen::Matrix<double, 81, 81>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &fun81,
		const std::function<DScalar2<double, Eigen::Matrix<double, Eigen::Dynamic, 1, 0, SMALL_N, 1>, Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, 0, SMALL_N, SMALL_N>>	(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &funN,
		const std::function<DScalar2<double, Eigen::VectorXd, Eigen::MatrixXd>								(const ElementAssemblyValues &, const Eigen::MatrixXd &, const QuadratureVector &)> &funn
		);

	double von_mises_stress_for_stress_tensor(const Eigen::MatrixXd &stress);
	void compute_diplacement_grad(const int size, const ElementBases &bs, const ElementAssemblyValues &vals, const Eigen::MatrixXd &local_pts, const int p, const Eigen::MatrixXd &displacement, Eigen::MatrixXd &displacement_grad);

	double convert_to_lambda(const bool is_volume, const double E, const double nu);
	double convert_to_mu(const double E, const double nu);

	class ElasticityTensor
	{
	public:
		void resize(const int size);

		double operator()(int i, int j) const;
		double &operator()(int i, int j);

		void set_from_entries(const std::vector<double> &entries);
		void set_from_lambda_mu(const double lambda, const double mu);
		void set_from_young_poisson(const double young, const double poisson);

		void set_orthotropic(
		double   Ex, double   Ey, double   Ez,
		double nuYX, double nuZX, double nuZY,
		double muYZ, double muZX, double muXY);
		void set_orthotropic(double Ex, double Ey, double nuYX, double muXY);

		template<int DIM>
		double compute_stress(const std::array<double, DIM> &strain, const int j) const;

	private:
		Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, 0, 6, 6> stifness_tensor_;
		int size_;
	};
}