// Copyright 2018, Intel Corporation

#include "tile/codegen/alias.h"

#include <algorithm>
#include <set>
#include <utility>

#include <boost/format.hpp>

#include "base/util/stream_container.h"
#include "base/util/throw.h"

namespace vertexai {
namespace tile {
namespace codegen {

using namespace stripe;  // NOLINT

namespace {

Affine UniqifyAffine(const Affine& orig, const std::string& prefix) {
  Affine ret;
  for (const auto& kvp : orig.getMap()) {
    if (kvp.first.empty()) {
      ret += kvp.second;
    } else {
      ret += Affine(prefix + kvp.first, kvp.second);
    }
  }
  return ret;
}

// Determines whether two accesses overlap (assuming base refinement
// and location checks have been performed).  This is used to
// determine the overlap of the internal ranges of two AliasInfos --
// i.e. whether they might overlap for any given instantiation of
// indices.  The implementation currently doesn't check all
// possibilities; it conservatively indicates an overlap in edge
// cases.
bool CheckRelativeOverlap(const AliasInfo& ai, const AliasInfo& bi) {
  IVLOG(4, boost::format("  CheckRelativeOverlap: a: '%1%', b: '%2%'") % ai % bi);
  if (ai.access.size() != bi.access.size()) {
    throw std::runtime_error{"Incompatible accesses"};
  }
  if ((ai.access.size() != ai.shape.dims.size()) || (bi.access.size() != bi.shape.dims.size())) {
    // N.B. This check might as well be an assert; there's no way this
    // should be inconsistent after the AliasInfo's been constructed.
    // But it's good to be careful.
    throw std::runtime_error{"Incompatible access shapes"};
  }
  bool ret = true;
  for (size_t i = 0; ret && i < ai.access.size(); ++i) {
    auto a_affine = ai.access[i];
    auto b_affine = bi.access[i];
    IVLOG(4, "    CheckRelativeOverlap[" << i << "]: Comparing " << a_affine << " with " << b_affine);
    auto& a_map = a_affine.mutateMap();
    auto& b_map = b_affine.mutateMap();
    for (auto a_it = a_map.begin(); a_it != a_map.end();) {
      auto b_it = b_map.find(a_it->first);
      if (b_it == b_map.end()) {
        ++a_it;
        continue;
      }
      auto min_factor = std::min(a_it->second, b_it->second);
      if (b_it->second == min_factor) {
        b_map.erase(b_it);
      } else {
        b_it->second -= min_factor;
      }
      if (a_it->second == min_factor) {
        a_it = a_map.erase(a_it);
      } else {
        a_it->second -= min_factor;
        ++a_it;
      }
    }
    IVLOG(4, "    CheckRelativeOverlap[" << i << "]: Simplified to " << a_affine << " and " << b_affine);
    if (!a_affine.isConstant() || !b_affine.isConstant()) {
      // TODO: Figure out how to compute the correct answer when the
      // resulting affines aren't constant.  In this situation, the
      // refinements are moving relative to each other across
      // different instantiations; we'd need to tell whether this ever
      // ends up with the refinements' colliding.
      return true;
    }
    IVLOG(4, "    CheckRelativeOverlap[" << i << "]: Simplified to " << a_affine << " and " << b_affine);
    std::int64_t a_first = a_affine.constant();
    std::int64_t a_limit = a_first + ai.shape.dims.at(i).size;
    std::int64_t b_first = b_affine.constant();
    std::int64_t b_limit = b_first + bi.shape.dims.at(i).size;
    ret &= b_first < a_limit;
    ret &= a_first < b_limit;
  }
  IVLOG(4, boost::format("  CheckRelativeOverlap: a: '%1%', b: '%2%' => %3%") % ai % bi % ret);
  return ret;
}

}  // namespace

std::ostream& operator<<(std::ostream& os, const AliasInfo& ai) {
  os << "(" << ai.base_name;
  os << ", " << ai.location;
  os << ", " << StreamContainer(ai.access);
  os << ", " << ai.shape;
  os << ")";
  return os;
}

bool CheckOverlap(const std::vector<Extent>& ae, const std::vector<Extent>& be) {
  IVLOG(4, boost::format("  CheckOverlap: a: '%1%', b: '%2%'") % StreamContainer(ae) % StreamContainer(be));
  if (ae.size() != be.size()) {
    throw std::runtime_error("Incompatible extents");
  }
  bool ret = true;
  for (size_t i = 0; i < ae.size(); i++) {
    ret &= be[i].min <= ae[i].max;
    ret &= ae[i].min <= be[i].max;
  }
  return ret;
}

AliasType AliasInfo::Compare(const AliasInfo& ai, const AliasInfo& bi) {
  IVLOG(3, "AliasInfo::Compare> a: " << ai.base_name << ", b: " << bi.base_name);
  IVLOG(4, "  a: " << ai);
  IVLOG(4, "  b: " << bi);
  if (ai.base_name != bi.base_name) {
    IVLOG(3, "  Different base tensors");
    return AliasType::None;
  }

  if (ai.location != bi.location) {
    IVLOG(3, boost::format("  Different banks, a: %1%, b: %2%") % ai.location % bi.location);
    return AliasType::None;
  }
  if (ai.access == bi.access) {
    IVLOG(3, boost::format("  Exact access, a: %1%, b: %2%") % StreamContainer(ai.access) % StreamContainer(bi.access));
    return AliasType::Exact;
  }
  if (!CheckRelativeOverlap(ai, bi)) {
    IVLOG(3, "  No overlap");
    return AliasType::None;
  }
  // TODO: We could compute the convex box enclosing each refinement and then check each
  // dimension to see if there is a splitting plane, and if so, safely declare alias None,
  // but it's unclear that that will happen enough for us to care, so just always return
  // Partial which is conservative.
  IVLOG(3, "  Partial");
  return AliasType::Partial;
}

bool AliasInfo::IsBanked() const { return !!base_ref->bank_dim; }

Affine AliasInfo::flat() const {
  Affine flat;
  for (size_t i = 0; i < shape.dims.size(); i++) {
    flat += access[i] * shape.dims[i].stride;
  }
  return flat;
}

AliasMap::AliasMap() : depth_(0), this_block_(nullptr), parent_block_(nullptr), parent_alias_map_(nullptr) {}

AliasMap::AliasMap(const AliasMap& outer, stripe::Block* block) : depth_(outer.depth_ + 1), parent_alias_map_(&outer) {
  idx_ranges_ = outer.idx_ranges_;
  this_block_ = block;
  parent_block_ = outer.this_block();
  // Make a prefix
  std::string prefix = str(boost::format("d%1%:") % depth_);
  // Get all the index data for new indexes
  for (const auto& idx : block->idxs) {
    if (idx.affine == stripe::Affine()) {
      // If it's a normal index, add it's information
      idx_ranges_[prefix + idx.name] = idx.range;
      idx_sources_[idx.name] = Affine(prefix + idx.name);
    } else {
      idx_sources_[idx.name] = idx.affine.sym_eval(outer.idx_sources_);
    }
  }
  // Make all inner alias data
  for (auto& ref : block->refs) {
    // Setup the place we are going to write to
    AliasInfo& info = info_[ref.into()];
    // Check if it's a refinement or a new buffer
    if (ref.dir != RefDir::None) {
      // Get the state from the outer context, fail if not found
      auto it = outer.info_.find(ref.from);
      if (it == outer.info_.end()) {
        throw_with_trace(std::runtime_error(
            str(boost::format("AliasMap::AliasMap: invalid ref.from during aliasing computation: '%1%' (ref: '%2%')") %
                ref.from % ref)));
      }
      // Copy data across
      info.base_block = it->second.base_block;
      info.base_ref = it->second.base_ref;
      info.base_name = it->second.base_name;
      info.access = it->second.access;
      info.location = AddDeviceUnits(it->second.location, ref.location);
    } else {
      // New alloc, initialize from scratch
      info.base_block = block;
      info.base_ref = &ref.mut();
      info.base_name = prefix + ref.into();
      info.access.resize(ref.access.size());
      info.location = ref.location;
    }
    if (info.access.size() != ref.access.size()) {
      throw_with_trace(std::runtime_error(
          str(boost::format("AliasMap::AliasMap: Mismatched access dimensions on refinement: %1% %2%") %
              info.base_name % ref.into())));
    }
    info.extents.resize(ref.access.size());
    for (size_t i = 0; i < ref.access.size(); i++) {
      info.access[i] += UniqifyAffine(ref.access[i], prefix);
      Extent& ext = info.extents[i];
      ext.min = 0;
      ext.max = 0;
      for (const auto& kvp : info.access[i].getMap()) {
        if (kvp.first == "") {
          ext.min += kvp.second;
          ext.max += kvp.second;
        } else {
          if (kvp.second > 0) {
            ext.max += (idx_ranges_[kvp.first] - 1) * kvp.second;
          } else {
            ext.min += (idx_ranges_[kvp.first] - 1) * kvp.second;
          }
        }
      }
      ext.max += ref.interior_shape.dims[i].size - 1;
    }
    IVLOG(5,
          boost::format("Extents for '%1%' in '%2%': %3%") % ref.into() % block->name % StreamContainer(info.extents));
    // Set shape
    info.shape = ref.interior_shape;
  }
}

std::unordered_map<std::string, size_t> AliasMap::RefUseCounts(const Block& block) const {
  // Compute statement use count of each buffer
  std::unordered_map<std::string, size_t> use_count;
  for (const auto& stmt : block.stmts) {
    std::set<std::string> buf_use;
    for (const auto& str : stmt->buffer_reads()) {
      buf_use.emplace(str);
    }
    for (const auto& str : stmt->buffer_writes()) {
      buf_use.emplace(str);
    }
    for (const auto& str : buf_use) {
      use_count[str]++;
    }
  }
  return use_count;
}

stripe::Affine AliasMap::translate(const stripe::Affine& in) const {
  IVLOG(3, "AliasMap::translate in = " << in);
  Affine cur = in;
  for (const auto& kvp : idx_ranges_) {
    if (kvp.second == 1) {
      cur.mutateMap().erase(kvp.first);
    }
  }
  IVLOG(3, "AliasMap::after removing 'irrelevant' indexes = " << cur);
  Affine out;
  for (const auto& kvp : idx_sources_) {
    IVLOG(4, "  " << kvp.first << " = " << kvp.second);
    int64_t mul = 0;
    for (const auto& kvp2 : kvp.second.getMap()) {
      if (mul == 0 && cur.getMap().count(kvp2.first)) {
        mul = cur.getMap().at(kvp2.first) / kvp2.second;
        break;
      }
    }
    out += Affine(kvp.first, mul);
    cur -= kvp.second * mul;
    IVLOG(4, "  mul = " << mul << ", cur = " << cur);
  }
  if (!cur.isConstant()) {
    throw std::runtime_error("AliasMap::translate, unable to translate " + in.toString());
  }
  out += cur.constant();
  IVLOG(3, "  out = " << out);
  return out;
}

void AliasMap::AddConstraintForIndex(stripe::Block* block,         //
                                     const AliasInfo& alias_info,  //
                                     size_t idx,                   //
                                     const std::string& idx_name,  //
                                     bool idx_passthru) const {
  int64_t top_index = alias_info.base_ref->interior_shape.dims[idx].size - 1;
  bool underflow = alias_info.extents[idx].min < 0;
  bool overflow = alias_info.extents[idx].max > top_index;
  if (underflow || overflow) {
    IVLOG(3, "AddConstraintForIndex: " << alias_info.base_name << ", " << idx_name);
    IVLOG(3, "extents = " << alias_info.extents[idx] << ", top_index = " << top_index);
    if (idx_name != "") {
      std::string global_idx_name = block->unique_idx_name(idx_name);
      block->idxs.emplace_back(Index{global_idx_name, 1, translate(alias_info.access[idx])});
      if (underflow) {
        block->constraints.push_back(idx_passthru ? Affine(global_idx_name)
                                                  : (Affine(global_idx_name) + Affine(idx_name)));
      }
      if (overflow) {
        block->constraints.push_back(idx_passthru ? (Affine(top_index) - Affine(global_idx_name))
                                                  : (Affine(top_index) - Affine(global_idx_name) - Affine(idx_name)));
      }
    } else {
      Affine exp = translate(alias_info.access[idx]);
      if (underflow) {
        block->constraints.push_back(exp);
      }
      if (overflow) {
        block->constraints.push_back(Affine(top_index) - exp);
      }
    }
  }
}

// Make a passthru index if it doesn't exist
std::string AliasMap::MakePassthruIdx(Block* block, const std::string& idx_name) {
  for (const auto idx : block->idxs) {
    if (idx.affine == Affine(idx_name)) {
      // The passthru index exists, return the existing index
      return idx.name;
    }
  }
  // doesn't exist
  std::string new_idx = block->unique_idx_name(idx_name);
  block->idxs.push_back({new_idx, 1, Affine(idx_name)});
  return new_idx;
}

// Add passthru index that are needed by exp from outer to this block
void AliasMap::AddPassthruIdx(Block* outer, Affine exp) {
  // Find a path from outer to inner
  std::vector<Block*> path;
  const AliasMap* am = this;
  while (am && am->this_block() != outer) {
    path.push_back(am->this_block());
    am = am->parent_alias_map();
  }
  if (am == nullptr) {
    throw std::runtime_error("Cannot find outer block.");
  }
  path.push_back(outer);
  // The idxs to be built in this level or inner
  std::vector<std::pair<std::string, size_t>> to_build;
  // The idxs to be built in the inner level or inner
  std::vector<std::pair<std::string, size_t>> to_build_next;
  // Resolve the initialze idxs
  for (const auto& it : exp.getMap()) {
    if (it.first != "") {
      size_t pos = it.first.find(':');
      std::string idx_name = it.first.substr(pos + 1);
      size_t idx_depth = std::stoi(it.first.substr(1, pos - 1));
      to_build.push_back(std::make_pair(idx_name, idx_depth));
    }
  }
  // Traverse from outer to inner
  for (int i = path.size() - 1; i >= 0; --i) {
    Block* this_block = path[i];
    size_t depth = depth_ - i;
    for (auto& it : to_build) {
      if (it.second == depth - 1) {
        std::string new_idx = MakePassthruIdx(this_block, it.first);
        to_build_next.push_back(std::make_pair(new_idx, depth));
      } else {
        to_build_next.push_back(std::make_pair(it.first, it.second));
      }
    }
    to_build = to_build_next;
    to_build_next.clear();
  }
}

void AliasMap::AddConstraintsForRef(const Refinement& ref) {
  const auto& old_ai = at(ref.into());
  // Add the passthru index that are used by the constraints
  for (size_t i = 0; i < ref.interior_shape.dims.size(); ++i) {
    AddPassthruIdx(parent_block_, old_ai.access[i]);
  }
  // Update alias map
  AliasMap parent_map(*(parent_alias_map_->parent_alias_map()), parent_block_);
  AliasMap this_map(parent_map, this_block_);
  const auto& new_ai = at(ref.into());
  for (size_t i = 0; i < ref.interior_shape.dims.size(); ++i) {
    this_map.AddConstraintForIndex(this_block_, new_ai, i, "", ref.interior_shape.dims[i].size <= 1);
  }
}

void AliasMap::AddConstraintsForBlock() {
  for (const auto& ref : this_block_->refs) {
    const auto& old_ai = at(ref.into());
    // Add the passthru index that are used by the constraints
    for (size_t i = 0; i < ref.interior_shape.dims.size(); ++i) {
      AddPassthruIdx(parent_block_, old_ai.access[i]);
    }
  }
  // Update alias map
  AliasMap parent_map(*(parent_alias_map_->parent_alias_map()), parent_block_);
  AliasMap this_map(parent_map, this_block_);
  for (const auto& ref : this_block_->refs) {
    const auto& new_ai = at(ref.into());
    for (size_t i = 0; i < ref.interior_shape.dims.size(); ++i) {
      this_map.AddConstraintForIndex(this_block_, new_ai, i, "", ref.interior_shape.dims[i].size <= 1);
    }
  }
}

}  // namespace codegen
}  // namespace tile
}  // namespace vertexai


#include <engine/Vectormath.hpp>
#include <engine/Manifoldmath.hpp>
#include <utility/Constants.hpp>
#include <utility/Logging.hpp>
#include <utility/Exception.hpp>

#include <Eigen/Dense>

#include <GenEigsSolver.h>  // Also includes <MatOp/DenseGenMatProd.h>
#include <GenEigsRealShiftSolver.h>

#include <array>

namespace C = Utility::Constants;

#ifndef SPIRIT_USE_CUDA

namespace Engine
{
    namespace Manifoldmath
    {
        void project_parallel(vectorfield & vf1, const vectorfield & vf2)
        {
            vectorfield vf3 = vf1;
            project_orthogonal(vf3, vf2);
            // TODO: replace the loop with Vectormath Kernel
            #pragma omp parallel for
            for (unsigned int i = 0; i < vf1.size(); ++i)
                vf1[i] -= vf3[i];
        }

        void project_orthogonal(vectorfield & vf1, const vectorfield & vf2)
        {
            scalar x = Vectormath::dot(vf1, vf2);
            // TODO: replace the loop with Vectormath Kernel
            #pragma omp parallel for
            for (unsigned int i=0; i<vf1.size(); ++i)
                vf1[i] -= x*vf2[i];
        }

        void invert_parallel(vectorfield & vf1, const vectorfield & vf2)
        {
            scalar x = Vectormath::dot(vf1, vf2);
            // TODO: replace the loop with Vectormath Kernel
            #pragma omp parallel for
            for (unsigned int i=0; i<vf1.size(); ++i)
                vf1[i] -= 2*x*vf2[i];
        }

        void invert_orthogonal(vectorfield & vf1, const vectorfield & vf2)
        {
            vectorfield vf3 = vf1;
            project_orthogonal(vf3, vf2);
            // TODO: replace the loop with Vectormath Kernel
            #pragma omp parallel for
            for (unsigned int i = 0; i < vf1.size(); ++i)
                vf1[i] -= 2 * vf3[i];
        }

        void project_tangential(vectorfield & vf1, const vectorfield & vf2)
        {
            #pragma omp parallel for
            for (unsigned int i = 0; i < vf1.size(); ++i)
                vf1[i] -= vf1[i].dot(vf2[i]) * vf2[i];
        }

        scalar dist_geodesic(const vectorfield & v1, const vectorfield & v2)
        {
            scalar dist = 0;
            #pragma omp parallel for reduction(+:dist)
            for (unsigned int i = 0; i < v1.size(); ++i)
                dist += pow(Vectormath::angle(v1[i], v2[i]), 2);
            return sqrt(dist);
        }

        /*
        Calculates the 'tangent' vectors, i.e.in crudest approximation the difference between an image and the neighbouring
        */
        void Tangents(std::vector<std::shared_ptr<vectorfield>> configurations, const std::vector<scalar> & energies, std::vector<vectorfield> & tangents)
        {
            int noi = configurations.size();
            int nos = (*configurations[0]).size();

            for (int idx_img = 0; idx_img < noi; ++idx_img)
            {
                auto& image = *configurations[idx_img];

                // First Image
                if (idx_img == 0)
                {
                    auto& image_plus = *configurations[idx_img + 1];
                    Vectormath::set_c_a( 1, image_plus, tangents[idx_img]);
                    Vectormath::add_c_a(-1, image,      tangents[idx_img]);
                }
                // Last Image
                else if (idx_img == noi - 1)
                {
                    auto& image_minus = *configurations[idx_img - 1];
                    Vectormath::set_c_a( 1, image,       tangents[idx_img]);
                    Vectormath::add_c_a(-1, image_minus, tangents[idx_img]);
                }
                // Images Inbetween
                else
                {
                    auto& image_plus  = *configurations[idx_img + 1];
                    auto& image_minus = *configurations[idx_img - 1];

                    // Energies
                    scalar E_mid = 0, E_plus = 0, E_minus = 0;
                    E_mid   = energies[idx_img];
                    E_plus  = energies[idx_img + 1];
                    E_minus = energies[idx_img - 1];

                    // Vectors to neighbouring images
                    vectorfield t_plus(nos), t_minus(nos);

                    Vectormath::set_c_a( 1, image_plus, t_plus);
                    Vectormath::add_c_a(-1, image,      t_plus);

                    Vectormath::set_c_a( 1, image,       t_minus);
                    Vectormath::add_c_a(-1, image_minus, t_minus);

                    // Near maximum or minimum
                    if ((E_plus < E_mid && E_mid > E_minus) || (E_plus > E_mid && E_mid < E_minus))
                    {
                        // Get a smooth transition between forward and backward tangent
                        scalar E_max = std::max(std::abs(E_plus - E_mid), std::abs(E_minus - E_mid));
                        scalar E_min = std::min(std::abs(E_plus - E_mid), std::abs(E_minus - E_mid));

                        if (E_plus > E_minus)
                        {
                            Vectormath::set_c_a(E_max, t_plus,  tangents[idx_img]);
                            Vectormath::add_c_a(E_min, t_minus, tangents[idx_img]);
                        }
                        else
                        {
                            Vectormath::set_c_a(E_min, t_plus,  tangents[idx_img]);
                            Vectormath::add_c_a(E_max, t_minus, tangents[idx_img]);
                        }
                    }
                    // Rising slope
                    else if (E_plus > E_mid && E_mid > E_minus)
                    {
                        Vectormath::set_c_a(1, t_plus,  tangents[idx_img]);
                    }
                    // Falling slope
                    else if (E_plus < E_mid && E_mid < E_minus)
                    {
                        Vectormath::set_c_a(1, t_minus,  tangents[idx_img]);
                        //tangents = t_minus;
                        for (int i = 0; i < nos; ++i)
                        {
                            tangents[idx_img][i] = t_minus[i];
                        }
                    }
                    // No slope(constant energy)
                    else
                    {
                        Vectormath::set_c_a(1, t_plus,  tangents[idx_img]);
                        Vectormath::add_c_a(1, t_minus, tangents[idx_img]);
                    }
                }

                // Project tangents into tangent planes of spin vectors to make them actual tangents
                project_tangential(tangents[idx_img], image);

                // Normalise in 3N - dimensional space
                Manifoldmath::normalize(tangents[idx_img]);

            }// end for idx_img
        }// end Tangents
    }
}

#endif

namespace Engine
{
    namespace Manifoldmath
    {
        scalar norm(const vectorfield & vf)
        {
            scalar x = Vectormath::dot(vf, vf);
            return std::sqrt(x);
        }

        void normalize(vectorfield & vf)
        {
            scalar sc = 1.0/norm(vf);
            Vectormath::scale(vf, sc);
        }

        MatrixX tangential_projector(const vectorfield & image)
        {
            int nos = image.size();
            int size = 3*nos;

            // Get projection matrix M=1-S, blockwise S=x*x^T
            MatrixX proj = MatrixX::Identity(size, size);
            for (int i = 0; i < nos; ++i)
            {
                proj.block<3, 3>(3*i, 3*i) -= image[i] * image[i].transpose();
            }

            return proj;
        }

        // This gives an orthogonal matrix of shape (3N, 2N), meaning M^T=M^-1 or M^T*M=1.
        // This assumes that the vectors of vf are normalized and that basis is 3N x 2N
        // It can be used to transform a vector into or back from the tangent space of a
        //      sphere w.r.t. euclidean 3N space.
        // It is generated by column-wise normalization of the Jacobi matrix for the
        //      transformation from (unit-)spherical coordinates to euclidean.
        // It therefore consists of the local basis vectors of the spherical coordinates
        //      of a unit sphere, represented in 3N, as the two columns of the matrix.
        void tangent_basis_spherical(const vectorfield & vf, MatrixX & basis)
        {
            Vector3 tmp, etheta, ephi;
            basis.setZero();
            for (unsigned int i=0; i < vf.size(); ++i)
            {
                if (vf[i][2] > 1-1e-8)
                {
                    tmp = Vector3{1, 0, 0};
                    basis.block<3,1>(3*i,2*i)   = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                    tmp = Vector3{0, 1, 0};
                    basis.block<3,1>(3*i,2*i+1) = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                }
                else if (vf[i][2] < -1+1e-8)
                {
                    tmp = Vector3{1, 0, 0};
                    basis.block<3,1>(3*i,2*i)   = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                    tmp = Vector3{0, -1, 0};
                    basis.block<3,1>(3*i,2*i+1) = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                }
                else
                {
                    scalar rxy = std::sqrt( 1 - vf[i][2]*vf[i][2] );
                    scalar z_rxy = vf[i][2] / rxy;

                    // Note: these are not unit vectors, but derivatives!
                    etheta = Vector3{  vf[i][0]*z_rxy, vf[i][1]*z_rxy, -rxy };
                    ephi   = Vector3{ -vf[i][1]/rxy,   vf[i][0]/rxy,    0   };

                    basis.block<3,1>(3*i,2*i)   = (etheta - etheta.dot(vf[i])*vf[i]).normalized();
                    basis.block<3,1>(3*i,2*i+1) = (ephi   - ephi.dot(vf[i])*vf[i]).normalized();
                }
            }
        }

        // This calculates the basis via calculation of cross products
        // This assumes that the vectors of vf are normalized and that basis is 3N x 2N
        void tangent_basis_cross(const vectorfield & vf, MatrixX & basis)
        {
            basis.setZero();
            for(int i=0; i < vf.size(); ++i)
            {
                if (std::abs(vf[i].z()) > 1-1e-8)
                {
                    basis.block<3,1>(3*i,2*i)   = Vector3{0,1,0}.cross(vf[i]).normalized();
                    basis.block<3,1>(3*i,2*i+1) = vf[i].cross(basis.block<3,1>(3*i,2*i));
                }
                else
                {
                    basis.block<3,1>(3*i,2*i)   = Vector3{0,0,1}.cross(vf[i]).normalized();
                    basis.block<3,1>(3*i,2*i+1) = vf[i].cross(basis.block<3,1>(3*i,2*i));
                }
            }
        }

        // This calculates the basis via orthonormalization to a random vector
        // This assumes that the vectors of vf are normalized and that basis is 3N x 2N
        void tangent_basis_righthanded(const vectorfield & vf, MatrixX & basis)
        {
            int size = vf.size();
            basis.setZero();

            // vf should be 3N
            // basis should be 3N x 2N

            // e1 and e2 will form a righthanded vectorset with the axis (though not orthonormal!)
            Vector3 e1, e2, v1;
            Vector3 ex{1,0,0}, ey{0,1,0}, ez{0,0,1};

            for(int i=0; i < size; ++i)
            {
                auto& axis = vf[i];

                // Choose orthogonalisation basis for Grahm-Schmidt
                //      We will need two vectors with which the axis always forms the
                //      same orientation (händigkeit des vektor-dreibeins)
                // If axis_z=0 its in the xy-plane
                //      the vectors should be: axis, ez, (axis x ez)
                if (axis[2] == 0)
                {
                    e1 = ez;
                    e2 = axis.cross(ez);
                }
                // Else its either above or below the xy-plane.
                //      if its above the xy-plane, it points in z-direction
                //      the vectors should be: axis, ex, -ey
                else if (axis[2] > 0)
                {
                    e1 = ex;
                    e2 = -ey;
                }
                //      if its below the xy-plane, it points in -z-direction
                //      the vectors should be: axis, ex, ey
                else if (axis[2] < 0)
                {
                    e1 = ex;
                    e2 = ey;
                }

                // First vector: orthogonalize e1 w.r.t. axis
                v1 = ( e1 - e1.dot(axis) * axis ).normalized();
                basis.block<3,1>(3*i,2*i)   = v1;

                // Second vector: orthogonalize e2 w.r.t. axis and v1
                basis.block<3,1>(3*i,2*i+1) = ( e2 - e2.dot(axis)*axis - e2.dot(v1)*v1 ).normalized();
            }
        }

        // This gives the Jacobian matrix for the transformation from (unit-)spherical
        // to euclidean coordinates. It consists of the derivative vectors d/d_theta
        // and d/d_phi as the two columns of the matrix.
        void spherical_to_cartesian_jacobian(const vectorfield & vf, MatrixX & jacobian)
        {
            Vector3 tmp, etheta, ephi;
            jacobian.setZero();
            for (unsigned int i=0; i < vf.size(); ++i)
            {
                if (vf[i][2] > 1-1e-8)
                {
                    tmp = Vector3{1, 0, 0};
                    jacobian.block<3,1>(3*i,2*i)   = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                    tmp = Vector3{0, 1, 0};
                    jacobian.block<3,1>(3*i,2*i+1) = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                }
                else if (vf[i][2] < -1+1e-8)
                {
                    tmp = Vector3{1, 0, 0};
                    jacobian.block<3,1>(3*i,2*i)   = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                    tmp = Vector3{0, -1, 0};
                    jacobian.block<3,1>(3*i,2*i+1) = (tmp - tmp.dot(vf[i])*vf[i]).normalized();
                }
                else
                {
                    scalar rxy = std::sqrt( 1 - vf[i][2]*vf[i][2] );
                    scalar z_rxy = vf[i][2] / rxy;

                    // Note: these are not unit vectors, but derivatives!
                    etheta = Vector3{  vf[i][0]*z_rxy, vf[i][1]*z_rxy, -rxy };
                    ephi   = Vector3{ -vf[i][1],       vf[i][0],        0   };

                    jacobian.block<3,1>(3*i,2*i)   = etheta - etheta.dot(vf[i])*vf[i];
                    jacobian.block<3,1>(3*i,2*i+1) = ephi   - ephi.dot(vf[i])*vf[i];
                }
            }
        }

        // The Hessian matrix of the transformation from spherical to euclidean coordinates
        void spherical_to_cartesian_hessian(const vectorfield & vf, MatrixX & gamma_x, MatrixX & gamma_y, MatrixX & gamma_z)
        {
            int nos = vf.size();
            gamma_x.setZero();
            gamma_y.setZero();
            gamma_z.setZero();

            for (unsigned int i=0; i < nos; ++i)
            {
                scalar z_rxy = vf[i][2] / std::sqrt(1+1e-6-vf[i][2]*vf[i][2]);

                gamma_x.block<2,2>(2*i,2*i) << -vf[i][0],       -vf[i][1]*z_rxy,
                                                -vf[i][1]*z_rxy, -vf[i][0];

                gamma_y.block<2,2>(2*i,2*i) << -vf[i][1],        vf[i][0]*z_rxy,
                                                vf[i][0]*z_rxy, -vf[i][1];

                gamma_z.block<2,2>(2*i,2*i) << -vf[i][2], 0,
                                                0,        0;
            }
        }

        // The (2Nx2N) Christoffel symbols of the transformation from (unit-)spherical coordinates to euclidean
        void spherical_to_cartesian_christoffel_symbols(const vectorfield & vf, MatrixX & gamma_theta, MatrixX & gamma_phi)
        {
            using std::cos;
            using std::sin;
            using std::acos;
            using std::tan;
            using std::atan2;

            int nos = vf.size();
            gamma_theta = MatrixX::Zero(2*nos, 2*nos);
            gamma_phi   = MatrixX::Zero(2*nos, 2*nos);

            for (unsigned int i=0; i < nos; ++i)
            {
                scalar theta = acos(vf[i][2]);
                scalar phi   = atan2(vf[i][1], vf[i][0]);
                scalar cot = 0;
                if (std::abs(theta) > 1e-4)
                    cot = -tan(C::Pi_2 + theta);

                gamma_theta(2*i+1,2*i+1) = -sin(theta)*cos(theta);

                gamma_phi(2*i+1,2*i)   = cot;
                gamma_phi(2*i,2*i+1)   = cot;
            }
        }

        void hessian_bordered(const vectorfield & image, const vectorfield & gradient, const MatrixX & hessian, MatrixX & tangent_basis, MatrixX & hessian_out)
        {
            // Calculates a 3Nx3N matrix in the bordered Hessian approach and transforms it into the tangent basis,
            // making the result a 2Nx2N matrix. The bordered Hessian's Lagrange multipliers assume a local extremum.

            int nos = image.size();
            MatrixX tmp_3N = hessian;

            VectorX lambda(nos);
            for (int i=0; i<nos; ++i)
                lambda[i] = image[i].dot(gradient[i]);

            for (int i=0; i<nos; ++i)
            {
                for (int j=0; j<3; ++j)
                {
                    tmp_3N(3*i+j,3*i+j) -= lambda(i);
                }
            }

            // Calculate the basis transformation matrix
            tangent_basis = MatrixX::Zero(3*nos, 2*nos);
            tangent_basis_spherical(image, tangent_basis);

            // Result is a 2Nx2N matrix
            hessian_out = tangent_basis.transpose() * tmp_3N * tangent_basis;
        }

        void hessian_projected(const vectorfield & image, const vectorfield & gradient, const MatrixX & hessian, MatrixX & tangent_basis, MatrixX & hessian_out)
        {
            // Calculates a 3Nx3N matrix in the projector approach and transforms it into the tangent basis,
            // making the result a 2Nx2N matrix

            int nos = image.size();
            hessian_out.setZero();

            // Calculate projector matrix
            auto P = tangential_projector(image);

            // Calculate tangential projection of Hessian
            hessian_out = P * hessian * P;

            // Calculate correction terms
            for (unsigned int i = 0; i < nos; ++i)
            {
                hessian_out.block<3, 3>(3*i, 3*i) -=    P.block<3, 3>(3*i, 3*i) * ( image[i].dot(gradient[i]) )
                                                    + ( P.block<3, 3>(3*i, 3*i) * gradient[i] ) * image[i].transpose();
            }

            // Calculate the basis transformation matrix
            tangent_basis = MatrixX::Zero(3*nos, 2*nos);
            tangent_basis_spherical(image, tangent_basis);

            // Result is a 2Nx2N matrix
            hessian_out = tangent_basis.transpose() * hessian_out * tangent_basis;
        }

        void hessian_weingarten(const vectorfield & image, const vectorfield & gradient, const MatrixX & hessian, MatrixX & tangent_basis, MatrixX & hessian_out)
        {
            // Calculates a 3Nx3N matrix in the Weingarten map approach and transforms it into the tangent basis,
            // making the result a 2Nx2N matrix

            int nos = image.size();
            hessian_out.setZero();

            // Calculate projector matrix
            auto P = tangential_projector(image);

            // Calculate tangential projection of Hessian
            hessian_out = P * hessian;

            // Add the Weingarten map
            for (unsigned int i = 0; i < nos; ++i)
            {
                MatrixX proj = MatrixX::Identity(3, 3);
                hessian_out.block<3, 3>(3*i, 3*i) -=  MatrixX::Identity(3, 3) * image[i].dot(gradient[i]);
            }

            // Calculate the basis transformation matrix
            tangent_basis = MatrixX::Zero(3*nos, 2*nos);
            tangent_basis_spherical(image, tangent_basis);

            // Result is a 2Nx2N matrix
            hessian_out = tangent_basis.transpose() * hessian_out * tangent_basis;
        }

        void hessian_spherical(const vectorfield & image, const vectorfield & gradient, const MatrixX & hessian, MatrixX & hessian_out)
        {
            // Calculates a 2Nx2N hessian matrix containing second order spherical derivatives

            int nos = image.size();

            MatrixX jacobian   = MatrixX::Zero(3*nos, 2*nos);
            MatrixX sph_hess_x = MatrixX::Zero(2*nos, 2*nos);
            MatrixX sph_hess_y = MatrixX::Zero(2*nos, 2*nos);
            MatrixX sph_hess_z = MatrixX::Zero(2*nos, 2*nos);

            // Calculate coordinate transformation jacobian
            Engine::Manifoldmath::spherical_to_cartesian_jacobian(image, jacobian);

            // Calculate coordinate transformation Hessian
            Engine::Manifoldmath::spherical_to_cartesian_hessian(image, sph_hess_x, sph_hess_y, sph_hess_z);

            // Calculate transformed Hessian
            hessian_out = jacobian.transpose() * hessian * jacobian;
            for (int i=0; i < nos; ++i)
            {
                hessian_out.block<2,2>(2*i, 2*i) += gradient[i][0] * sph_hess_x.block<2,2>(2*i, 2*i)
                                                    + gradient[i][1] * sph_hess_y.block<2,2>(2*i, 2*i)
                                                    + gradient[i][2] * sph_hess_z.block<2,2>(2*i, 2*i);
            }

        }

        void hessian_covariant(const vectorfield & image, const vectorfield & gradient, const MatrixX & hessian, MatrixX & hessian_out)
        {
            // Calculates a 2Nx2N covariant hessian matrix containing second order spherical derivatives
            // and correction terms (containing Christoffel symbols)

            int nos = image.size();

            // Calculate coordinate transformation jacobian
            MatrixX jacobian(3*nos, 2*nos);
            Engine::Manifoldmath::spherical_to_cartesian_jacobian(image, jacobian);

            // Calculate the gradient in spherical coordinates
            Eigen::Ref<const VectorX> grad = Eigen::Map<const VectorX>(gradient[0].data(), 3 * nos);
            VectorX gradient_spherical = jacobian.transpose() * grad;

            // Calculate the Hessian in spherical coordinates
            hessian_spherical(image, gradient, hessian, hessian_out);

            // Calculate the Christoffel symbols for spherical coordinates
            MatrixX christoffel_theta = MatrixX::Zero(2*nos, 2*nos);
            MatrixX christoffel_phi   = MatrixX::Zero(2*nos, 2*nos);
            Engine::Manifoldmath::spherical_to_cartesian_christoffel_symbols(image, christoffel_theta, christoffel_phi);

            // Calculate the covariant Hessian
            for (int i=0; i < nos; ++i)
            {
                hessian_out.block<2,2>(2*i, 2*i) -= gradient_spherical[2*i]   * christoffel_theta.block<2,2>(2*i, 2*i)
                                                    + gradient_spherical[2*i+1] * christoffel_phi.block<2,2>(2*i, 2*i);
            }
        }

    }
}


//          Copyright Gavin Band 2008 - 2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

#include <iostream>
#include <iomanip>
#include <boost/ptr_container/ptr_deque.hpp>
#include <boost/function.hpp>
#include <boost/timer/timer.hpp>
#include <boost/thread/thread.hpp>
#include <boost/bind.hpp>
#include <boost/format.hpp>
#include "unistd.h"
#include "config/config.hpp"
#if HAVE_CBLAS
	#include "cblas.h"
#endif

#include "Eigen/Core"
#include "Eigen/Eigenvalues"
#include "genfile/VariantIdentifyingData.hpp"
#include "genfile/VariantDataReader.hpp"
#include "genfile/SNPDataSourceProcessor.hpp"
#include "genfile/SingleSNPGenotypeProbabilities.hpp"
#include "genfile/vcf/get_set_eigen.hpp"
#include "statfile/BuiltInTypeStatSink.hpp"
#include "statfile/BuiltInTypeStatSource.hpp"
#include "worker/Task.hpp"
#include "worker/FunctionTask.hpp"
#include "appcontext/FileUtil.hpp"
#include "appcontext/get_current_time_as_string.hpp"
#include "components/RelatednessComponent/KinshipCoefficientComputer.hpp"
#include "components/RelatednessComponent/PCAComputer.hpp"
#include "components/RelatednessComponent/mean_centre_genotypes.hpp"

#define USING_BOOST_THREADPOOL 0
#if USING_BOOST_THREADPOOL
#include "boost/threadpool.hpp"
#else
#include "metro/concurrency/threadpool.hpp"
#endif

// #define DEBUG_KINSHIP_COEFFICIENT_COMPUTER 2

namespace impl {
	std::ostream& operator<<( std::ostream& out, SampleBounds const& b ) {
		out << "[" << b.begin_sample_i << "-" << b.end_sample_i << ", " << b.begin_sample_j << "-" << b.end_sample_j << "]" ;
		return out ;
	}

	// Return a tiling of a dxd matrix into (roughly) B blocks.
	// Blocks should be roughly evenly sized.
	// This implementation makes off-diagonal blocks half the size of diagonal blocks.
	std::vector< SampleBounds > get_matrix_lower_diagonal_tiling( std::size_t d, std::size_t B ) {
		// Because off-diagonal blocks involve about twice as much computation,
		// we make them half the size vertically.
		// If there are (N ( N+1)/2 ) square blocks this makes
		// (N (N+1) ) - N = N^2 blocks in total.
		// So we want to choose N so that N^2 is about twice the number of threads.
		// Def
		std::size_t const N = B ;
		std::vector< SampleBounds > result ;
		double K = double( d ) / N ;
//#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
		std::cerr << "get_matrix_lower_diagonal_tiling: " << d << "x" << d << " matrix, "
			<< N << " blocks across diagonal, "
			<< N*(N-1) << " blocks in lower triangle.\n" ;
//#endif
		// Set up some initial tasks.
		// Do j (the column) as the outer loop
		// so that tasks go in column-major direction.
		for( std::size_t j = 0; j < N; ++j ) {
			for( std::size_t i = 0; i < N; ++i ) {
				if( i == j ) {
					SampleBounds bounds ;
					bounds.begin_sample_i = std::ceil(i*K) ;
					bounds.end_sample_i = std::min( std::size_t( std::ceil( (i+1)*K ) ), d ) ;
					bounds.begin_sample_j = bounds.begin_sample_i ;
					bounds.end_sample_j = bounds.end_sample_i ;
					result.push_back( bounds ) ;
				}
				else {
					// Do two vertical stripes
					SampleBounds bounds ;
					bounds.begin_sample_i = std::ceil(i*K) ;
					bounds.end_sample_i = std::min( std::size_t( std::ceil( (i+1)*K ) ), d ) ;
					bounds.begin_sample_j = std::ceil(j*K) ;
					bounds.end_sample_j = std::ceil((j+0.5)*K) ;
					result.push_back( bounds ) ;

					bounds.begin_sample_j = bounds.end_sample_j ;
					bounds.end_sample_j = std::min( std::size_t( std::ceil( (j+1)*K ) ), d ) ;
					result.push_back( bounds ) ;
				}
			}
		}
//#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
		std::cerr << "get_matrix_lower_diagonal_tiling: tiling has " << result.size() << " blocks in total.\n" ;
//#endif
		return result ;
	}

	std::vector< SampleBounds > get_matrix_lower_diagonal_vertical_strip_tiling( std::size_t d, std::size_t const strip_width ) {
		std::vector< SampleBounds > result ;
//#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
		std::cerr << "get_matrix_lower_diagonal_vertical_strip_tiling: " << d << "x" << d << " matrix" ;
//#endif
		// Set up some initial tasks.
		// Do j (the column) as the outer loop
		// so that tasks go in column-major direction.
		for( std::size_t i = 0; i < d; i += strip_width ) {
			SampleBounds bounds ;
			bounds.begin_sample_i = i ;
			bounds.end_sample_i = d ;
			bounds.begin_sample_j = i ;
			bounds.end_sample_j = std::min( i+strip_width, d ) ;
			result.push_back( bounds ) ;
		}
//#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
		std::cerr << "get_matrix_lower_diagonal_vertical_strip_tiling: tiling has " << result.size() << " blocks in total.\n" ;
//#endif
		return result ;
	}

	std::vector< SampleBounds > get_matrix_lower_diagonal_horizontal_strip_tiling( std::size_t d, std::size_t const strip_width ) {
		std::vector< SampleBounds > result ;
//#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
		std::cerr << "get_matrix_lower_diagonal_horizontal_strip_tiling: " << d << "x" << d << " matrix" ;
//#endif
		// Set up some initial tasks.
		// Do j (the column) as the outer loop
		// so that tasks go in column-major direction.
		for( std::size_t i = 0; i < d; i += strip_width ) {
			SampleBounds bounds ;
			bounds.begin_sample_i = i ;
			bounds.end_sample_i = i + strip_width ;
			bounds.begin_sample_j = 0 ;
			bounds.end_sample_j = std::min( i+strip_width, d ) ;
			result.push_back( bounds ) ;
		}
//#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
		std::cerr << "get_matrix_lower_diagonal_horizontal_strip_tiling: tiling has " << result.size() << " blocks in total.\n" ;
//#endif
		return result ;
	}

	NTaskDispatcher::NTaskDispatcher(
		worker::Worker* worker
	):
		m_worker( worker )
	{
	}

	void NTaskDispatcher::set_number_of_tasks( std::size_t number_of_tasks ) {
		wait_until_complete() ;
		// TODO: do this better using the constructor
		// I don't have the docs here so don't know the right call.
		m_tasks.clear() ;
		for( std::size_t i = 0; i < number_of_tasks; ++i ) {
			m_tasks.push_back( 0 ) ;
		}
	}

	std::size_t NTaskDispatcher::number_of_tasks() const { return m_tasks.size() ; }

	void NTaskDispatcher::submit_task( std::size_t task_i, boost::function< void() > task ) {
		assert( task_i < m_tasks.size() ) ;
		if( !m_tasks.is_null( task_i ) ) {
			throw genfile::BadArgumentError(
				"impl::NTaskDispatcher::submit_task()",
				( boost::format( "task_i=%d" ) % task_i ).str(),
				"Task is still running.  Call wait_until_complete() first."
			) ;
		}
		worker::Task::UniquePtr new_task( new worker::FunctionTask( task ) ) ;
		m_tasks.replace( task_i, new_task ) ;
		m_worker->tell_to_perform_task( m_tasks[ task_i ] ) ;
	}

	void NTaskDispatcher::wait_until_complete() {
		for( std::size_t i = 0; i < m_tasks.size(); ++i ) {
			if( !m_tasks.is_null( i ) ) {
				m_tasks[ i ].wait_until_complete() ;
				m_tasks.replace( i, 0 ) ;
			}
		}
	}

	NormaliseGenotypesAndComputeXXtFast::UniquePtr NormaliseGenotypesAndComputeXXtFast::create(
		worker::Worker* worker,
		std::size_t const number_of_snps_per_computation
	) {
		return NormaliseGenotypesAndComputeXXtFast::UniquePtr( new NormaliseGenotypesAndComputeXXtFast( worker, number_of_snps_per_computation ) ) ;
	}

	NormaliseGenotypesAndComputeXXtFast::NormaliseGenotypesAndComputeXXtFast(
		worker::Worker* worker,
		std::size_t const number_of_snps_per_chunk
	):
		m_worker( worker ),
		m_dispatcher( new NTaskDispatcher( worker ) ),
		m_call_threshhold( 0.9 ),
		m_allele_frequency_threshhold( 0.001 ),
		m_number_of_snps_per_chunk( number_of_snps_per_chunk ),
		m_number_of_snps_per_computation( 4 * number_of_snps_per_chunk )
	{
		// We allocate lookup tables exactly once, here.
		// We follow, roughly, the scheme in plink 1.9 of having four lookup tables.
		// Here we encode pairs of genotypes in 4 bits (though we could do 3 as in plink).
		// Each computation then involves 16 variants, four from each lookup table.
		// This multi-table version reduces the number of times the result matrices
		// have to be traversed, and seems to make most difference in speed when
		// multithreading.
		m_lookup_tables.resize( 4, std::vector< double >( 1 << ( m_number_of_snps_per_chunk * 4 ), 0.0 ) ) ;
		m_nonmissingness_lookup_tables.resize( 4, std::vector< int >( 1 << ( m_number_of_snps_per_chunk * 4 ), 0.0 ) ) ;
	}

	NormaliseGenotypesAndComputeXXtFast::Matrix const& NormaliseGenotypesAndComputeXXtFast::result() const {
		return m_result ;
	}

	NormaliseGenotypesAndComputeXXtFast::IntegerMatrix const& NormaliseGenotypesAndComputeXXtFast::nonmissingness() const {
		return m_nonmissingness ;
	}

	std::size_t NormaliseGenotypesAndComputeXXtFast::number_of_snps_included() const {
		return m_snp_count ;
	}

	std::string NormaliseGenotypesAndComputeXXtFast::get_summary() const {
		return (
			boost::format(
				"NormaliseGenotypesAndComputeXXtFast():\n"
				" - %d SNPs per chunk\n"
				" - %d chunks per array acces\n"
				" - four lookup tables of size: %d\n"
				" - minimum allele frequency: %.3f"
			)
				% m_number_of_snps_per_chunk
				% (m_number_of_snps_per_computation / m_number_of_snps_per_chunk)
				% m_lookup_tables[0].size()
				% m_allele_frequency_threshhold
		).str()
		;
	}
	void NormaliseGenotypesAndComputeXXtFast::begin_processing_snps(
		std::size_t number_of_samples,
		genfile::SNPDataSource::Metadata const&
	) {
		m_result.resize( number_of_samples, number_of_samples ) ;
		m_result.setZero() ;
		m_nonmissingness.resize( number_of_samples, number_of_samples ) ;
		m_nonmissingness.setZero() ;
		std::size_t numberOfTasks = m_worker->get_number_of_worker_threads() ;
#if 0
		m_matrix_tiling = get_matrix_lower_diagonal_vertical_strip_tiling( number_of_samples, 1 ) ;
#else
		m_matrix_tiling = get_matrix_lower_diagonal_tiling(
			number_of_samples,
			numberOfTasks
		) ;
#endif
		m_dispatcher->set_number_of_tasks( m_matrix_tiling.size() ) ;
		m_combined_genotypes.resize( number_of_samples, 0 ) ;
		m_per_snp_genotypes.resize( number_of_samples, 0 ) ;
		m_snp_count = 0 ;
	}

	namespace impl {
		template< typename T >
		double add( T a, T b ) {
			return a + b ;
		}
	}

#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
	namespace {
		template< typename T >
		void print_lookup_table( std::ostream& out, std::vector< T > const& lookup_table, std::size_t const max = 64 ) {
			out << "print_lookup_table(): lookup table of size " << lookup_table.size() << ":" ;
			for( std::size_t i = 0; i < std::min( max, lookup_table.size() ); ++i ) {
				if( i % 8 == 0 ) {
					out << "\n" ;
				} else {
					out << " " ;
				}
				out << ( boost::format( " - 0x%02x: %.2f" ) % i % lookup_table[i] ) ;
			}
			out << "\n" ;
		}
	}
#endif

	void NormaliseGenotypesAndComputeXXtFast::processed_snp(
		genfile::VariantIdentifyingData const& id_data,
		genfile::VariantDataReader::SharedPtr data_reader
	) {
		// std::vector< std::size_t > m_genotypes
		// std::vector< int >
		m_per_snp_genotypes.assign( m_result.rows(), 0 ) ;

		// I find it simplest here to encode genotypes as
		// 0 (missing), 1 (AA homozygote), 2 (heterozygote), 3 (BB homozygote).
		data_reader->get(
			":genotypes:",
			genfile::vcf::get_threshholded_calls( m_per_snp_genotypes, m_call_threshhold, 0, 1, 2, 3 )
		) ;

		// compute allele frequency
		double allele2_count = 0.0 ;
		double nonmissing_count = 0.0 ;
		for( std::size_t i = 0; i < m_combined_genotypes.size(); ++i ) {
			bool const missing = ( m_per_snp_genotypes[i] == 0 ) ;
			if( !missing ) {
				allele2_count += m_per_snp_genotypes[i] - 1.0 ;
				nonmissing_count += 1 ;
			}
		}
		// Data estimate:
		// double const allele_frequency = allele2_count / ( nonmissing_count * 2.0 ) ;
		// Regularised estimate based on adding one of each allele,
		// this matches Price et al Nature Genetics 2006
		double const posterior_allele_frequency = (1 + allele2_count) / ( 2 + ( 2 * nonmissing_count)) ;
		double const allele_frequency = allele2_count / ( 2 * nonmissing_count ) ;
		if( nonmissing_count > 0 && std::min( allele_frequency, 1.0 - allele_frequency ) > m_allele_frequency_threshhold ) {
			// Ok we will process this SNP.
			// We batch SNPs together to form a total of m_number_of_snps_per_computation SNPs.
			// This means we visit the result matrix (which for large samples won't fit in cache)
			// as few times as possible.
			// Here we are not quite as efficient as plink: I use four bits instead of three to
			// encode the two genotypes.
			std::size_t const lookup_snp_index = m_snp_count % m_number_of_snps_per_computation ;
			if( lookup_snp_index == 0 ) {
				m_dispatcher->wait_until_complete() ;
				std::copy( m_per_snp_genotypes.begin(), m_per_snp_genotypes.end(), m_combined_genotypes.begin() ) ;
				// Clear the lookup tables
				for( std::size_t i = 0; i < m_lookup_tables.size(); ++i ) {
					std::fill( m_lookup_tables[i].begin(), m_lookup_tables[i].end(), 0.0 ) ;
					std::fill( m_nonmissingness_lookup_tables[i].begin(), m_nonmissingness_lookup_tables[i].end(), 0.0 ) ;
				}
			} else {
				// We allow 4 bits per genotype, although each genotype takes up only the lower two bits.
				// This allows us to encode pairs of genotypes in consecutive pairs of bits
				// for simpler lookup of the lookup table.
				for( std::size_t i = 0; i < m_combined_genotypes.size(); ++i ) {
					m_combined_genotypes[i] |= m_per_snp_genotypes[i] << (4 * lookup_snp_index) ;
				}
			}

			std::size_t const lookup_table_index = ( m_snp_count % m_number_of_snps_per_computation ) / m_number_of_snps_per_chunk ;
			double const mean = 2.0 * allele_frequency ;
			double const sd = std::sqrt( ( 2.0 * posterior_allele_frequency * ( 1.0 - posterior_allele_frequency )) ) ;
#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
			std::cerr << "NormaliseGenotypesAndComputeXXtFast::processed_snp(): filling lookup table " << lookup_table_index << ".\n" ;
#endif

			add_snp_to_lookup_table(
				lookup_snp_index % m_number_of_snps_per_chunk,
				mean, sd,
				&m_lookup_tables[ lookup_table_index ],
				&m_nonmissingness_lookup_tables[ lookup_table_index ]
			) ;

			++m_snp_count ;
#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
			boost::format fmt( "%.2f" ) ;
			std::cerr << "NormaliseGenotypesAndComputeXXtFast::processed_snp(): m_snp_count: " << m_snp_count << ".\n" ;
			std::cerr << "NormaliseGenotypesAndComputeXXtFast::processed_snp(): allele_frequency: " << fmt % allele_frequency << ".\n" ;
			std::cerr << "NormaliseGenotypesAndComputeXXtFast::processed_snp(): mean: " << fmt % mean << ", sd = " << fmt % sd << ".\n" ;
			print_lookup_table( std::cerr, m_lookup_tables[0], 32 ) ;
			print_lookup_table( std::cerr, m_lookup_tables[1], 32 ) ;
			print_lookup_table( std::cerr, m_lookup_tables[2], 32 ) ;
			print_lookup_table( std::cerr, m_lookup_tables[3], 32 ) ;
			//print_lookup_table( std::cerr, m_nonmissingness_lookup_table, 256 ) ;
#endif
			if( m_snp_count % m_number_of_snps_per_computation == 0 ) {
#if DEBUG_KINSHIP_COEFFICIENT_COMPUTER
			std::cerr << "NormaliseGenotypesAndComputeXXtFast::processed_snp(): submitting tasks...\n" ;
#endif
				submit_tasks( m_combined_genotypes ) ;
			}
		}
	}

	void NormaliseGenotypesAndComputeXXtFast::add_snp_to_lookup_table(
		std::size_t const snp_lookup_index,
		double const mean,
		double const sd,
		std::vector< double >* lookup_table,
		std::vector< int >* nonmissingness_lookup_table
	) {
		// We lay out the lookup table as follows.
		// The two genotypes at the first SNP occupy the first 4 bits,
		// the two genotypes at the second SNP occupy the next 4 bits,
		// etc.
		// Thus after k SNPs the size of the useful part of the table is 16^k,
		// which is computed as (1 << 4k).
		// Note that because genotype value 0 corresponds to missing data,
		// The block as computed after SNP k-1 is equal to the block
		// for missing values for SNP k.
		if( snp_lookup_index == 0 ) {
			// bottom 2 bits correspond to top genotype = 0.
			for( std::size_t g = 0; g < 4; ++g ) {
				(*nonmissingness_lookup_table)[g] = 0 ;
				(*lookup_table)[ g ] = 0 ;
			}
			for( std::size_t g = 1; g < 4; ++g ) {
				(*nonmissingness_lookup_table)[ (g << 2) + 0 ] = 0 ;
				(*nonmissingness_lookup_table)[ (g << 2) + 1 ] = 1 ;
				(*nonmissingness_lookup_table)[ (g << 2) + 2 ] = 1 ;
				(*nonmissingness_lookup_table)[ (g << 2) + 3 ] = 1 ;
				(*lookup_table)[ (g << 2) + 0 ] = 0.0 ;
				(*lookup_table)[ (g << 2) + 1 ] = ( ( 0.0 - mean ) / sd ) * ( ( (g-1) - mean ) / sd ) ;
				(*lookup_table)[ (g << 2) + 2 ] = ( ( 1.0 - mean ) / sd ) * ( ( (g-1) - mean ) / sd ) ;
				(*lookup_table)[ (g << 2) + 3 ] = ( ( 2.0 - mean ) / sd ) * ( ( (g-1) - mean ) / sd ) ;
			}
		} else {
			// We have already computed the initial block of size 16^k.
			// Now we compute 15 additional blocks for 16 in total.
			// In the computation we break these up as 4 blocks of size
			// This means this SNP has no effect and we need only copy.
			std::size_t const inner_block_size = 1 << ( 4 * snp_lookup_index ) ;
			std::size_t const outer_block_size = 4 * inner_block_size ;
			for( std::size_t outer_g = 0; outer_g < 4; ++outer_g ) {
				for( std::size_t inner_g = (outer_g == 0) ? 1 : 0; inner_g < 4; ++inner_g ) {
					double const value = ( inner_g == 0 || outer_g == 0 )
						? 0.0
						: (((inner_g-1.0) - mean)/sd) * (((outer_g-1.0) - mean)/sd) ;
					int const nonmissingness = ( inner_g == 0 || outer_g == 0 )
						? 0
						: 1 ;
					std::transform(
						lookup_table->begin(),
						lookup_table->begin() + inner_block_size,
						lookup_table->begin() + outer_block_size * outer_g + inner_block_size * inner_g,
						boost::bind( &impl::add< double >, _1, value )
					) ;
					std::transform(
						nonmissingness_lookup_table->begin(),
						nonmissingness_lookup_table->begin() + inner_block_size,
						nonmissingness_lookup_table->begin() + outer_block_size * outer_g + inner_block_size * inner_g,
						boost::bind( &impl::add< int >, _1, nonmissingness )
					) ;
				}
			}
		}
	}

	void NormaliseGenotypesAndComputeXXtFast::submit_tasks(
		std::vector< uint64_t > const& genotypes
	) {
		for( std::size_t tile_i = 0; tile_i < m_matrix_tiling.size(); ++tile_i ) {
	#if 1
			m_dispatcher->submit_task(
				tile_i,
				boost::bind(
					&NormaliseGenotypesAndComputeXXtFast::compute_block,
					this,
					boost::cref( m_matrix_tiling[ tile_i ] ),
					boost::cref( genotypes )
				)
			) ;
	#else
			compute_block( m_matrix_tiling[tile_i], genotypes, lookup_table, nonmissingness_lookup_table ) ;
	#endif
		}
	}

	void NormaliseGenotypesAndComputeXXtFast::end_processing_snps() {
		if( m_snp_count % m_number_of_snps_per_computation > 0 ) {
			submit_tasks( m_combined_genotypes ) ;
		}
		m_dispatcher->wait_until_complete() ;
		std::cerr << "snp count is: " << m_snp_count << ".\n" ;
		std::cerr << "non-missingness array is:\n" << m_nonmissingness.block(0,0,10,10) << ".\n" ;
		m_result.array() /= m_nonmissingness.array().cast< double >() ;
	}

	void NormaliseGenotypesAndComputeXXtFast::compute_block(
		SampleBounds const& sample_bounds,
		std::vector< uint64_t > const& genotypes
	) {
		// only compute the lower diagonal.
		// We also want to make the function go down the rows of the result
		// in the inner loop, since it's stored column-major, so put row index in the inner loop.
		for( std::size_t j = sample_bounds.begin_sample_j; j < sample_bounds.end_sample_j; ++j ) {
			for( std::size_t i = std::max( j, std::size_t( sample_bounds.begin_sample_i ) ); i < sample_bounds.end_sample_i; ++i ) {
				uint64_t const combined_genotype = ( genotypes[i] << 2 ) + genotypes[j] ;
				std::size_t const lookup_index0 = combined_genotype & 0xFFFF ;
				std::size_t const lookup_index1 = ( combined_genotype >> 16 ) & 0xFFFF ;
				std::size_t const lookup_index2 = ( combined_genotype >> 32 ) & 0xFFFF ;
				std::size_t const lookup_index3 = ( combined_genotype >> 48 ) & 0xFFFF ;
				m_result(i,j) += m_lookup_tables[0][ lookup_index0 ] ;
				m_result(i,j) += m_lookup_tables[1][ lookup_index1 ] ;
				m_result(i,j) += m_lookup_tables[2][ lookup_index2 ] ;
				m_result(i,j) += m_lookup_tables[3][ lookup_index3 ] ;
				m_nonmissingness(i,j) += m_nonmissingness_lookup_tables[0][ lookup_index0 ] ;
				m_nonmissingness(i,j) += m_nonmissingness_lookup_tables[1][ lookup_index1 ] ;
				m_nonmissingness(i,j) += m_nonmissingness_lookup_tables[2][ lookup_index2 ] ;
				m_nonmissingness(i,j) += m_nonmissingness_lookup_tables[3][ lookup_index3 ] ;
			}
		}
	}
}

KinshipCoefficientComputer::KinshipCoefficientComputer(
	appcontext::OptionProcessor const& options,
	genfile::CohortIndividualSource const& samples,
	appcontext::UIContext& ui_context,
	Computation::UniquePtr computation
):
	m_options( options ),
	m_ui_context( ui_context ),
	m_samples( samples ),
	m_computation( computation )
{
	assert( m_computation.get() ) ;
}

void KinshipCoefficientComputer::begin_processing_snps( std::size_t number_of_samples, genfile::SNPDataSource::Metadata const& metadata ) {
	m_computation->begin_processing_snps( number_of_samples, metadata ) ;
}

void KinshipCoefficientComputer::processed_snp( genfile::VariantIdentifyingData const& id_data, genfile::VariantDataReader::SharedPtr data_reader ) {
	m_computation->processed_snp( id_data, data_reader ) ;
}

void KinshipCoefficientComputer::end_processing_snps() {
	m_computation->end_processing_snps() ;

	std::string description = "Number of SNPs: "
		+ genfile::string_utils::to_string( m_computation->number_of_snps_included() )
		+ "\nNumber of samples: "
		+ genfile::string_utils::to_string( m_samples.get_number_of_individuals() )
	;
#if 1
	send_results(
		m_computation->number_of_snps_included(),
		m_computation->nonmissingness().cast< double >(),
		m_computation->result().cast< double >(),
		"KinshipCoefficientComputer",
		description
	) ;
	#endif
}


#include <algorithm>
#include <cstdlib>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include <boost/dynamic_bitset.hpp>
#include <routingkit/geo_position_to_node.h>

#include "DataStructures/Geometry/Area.h"
#include "DataStructures/Geometry/LatLng.h"
#include "DataStructures/Geometry/Point.h"
#include "DataStructures/Graph/Attributes/LatLngAttribute.h"
#include "DataStructures/Graph/Graph.h"
#include "Tools/CommandLine/CommandLineParser.h"
#include "Tools/Constants.h"

inline void printUsage() {
  std::cout <<
      "Usage: GenerateMatchingODPairs -g1 <file> -g2 <file> -o1 <file> -o2 <file>\n"
      "This program takes two graphs as input, matches the vertices in both graphs with\n"
      "each other and outputs for each graph an OD-file representing the same journeys.\n"
      "  -n <num>          the number of OD-pairs to be generated\n"
      "  -d <meters>       only match vertices within a maximum distance of <meters>\n"
      "  -s <seed>         the seed for the random number generator\n"
      "  -g1 <file>        the first input graph in binary format\n"
      "  -g2 <file>        the second input graph in binary format\n"
      "  -p <file>         restrict the OD-pairs to an area from an OSM POLY file\n"
      "  -o1 <file>        place the output for the first graph in <file>\n"
      "  -o2 <file>        place the output for the second graph in <file>\n"
      "  -help             display this help and exit\n";
}

using VPTree = RoutingKit::GeoPositionToNode; // An implementation of a vantage-point tree.

// Returns the point in the specified vp-tree that is closest to the query point p.
inline int findNearestNeighbor(const VPTree& tree, const LatLng& p, const int maxDist) {
  auto neighbors = tree.find_all_nodes_within_radius(p.latInDeg(), p.lngInDeg(), maxDist);
  std::sort(neighbors.begin(), neighbors.end(), [](const auto& u, const auto& v) {
    return u.distance < v.distance;
  });
  const int numNeighbors = neighbors.size();
  if (numNeighbors == 1 || (numNeighbors > 1 && neighbors[0].distance < neighbors[1].distance))
    return neighbors[0].id;
  else
    return INVALID_INDEX;
}

int main(int argc, char* argv[]) {
  try {
    CommandLineParser clp(argc, argv);
    if (clp.isSet("help")) {
      printUsage();
      return EXIT_SUCCESS;
    }

    const auto infile1 = clp.getValue<std::string>("g1");
    const auto infile2 = clp.getValue<std::string>("g2");
    const auto polyfile = clp.getValue<std::string>("p");
    const auto outfile1 = clp.getValue<std::string>("o1");
    const auto outfile2 = clp.getValue<std::string>("o2");
    const int numPairs = clp.getValue<int>("n");
    const int maxDist = clp.getValue<int>("d", 10);

    // Read both input graphs from disk.
    using Graph = StaticGraph<VertexAttrs<LatLngAttribute>>;
    std::ifstream in1(infile1, std::ios::binary);
    std::ifstream in2(infile2, std::ios::binary);
    if (!in1.good())
      throw std::invalid_argument("file not found -- '" + infile1 + "'");
    if (!in2.good())
      throw std::invalid_argument("file not found -- '" + infile2 + "'");
    Graph graph1(in1);
    Graph graph2(in2);
    in1.close();
    in2.close();

    // Build vp-trees for both graphs.
    std::vector<float> lat1(graph1.numVertices());
    std::vector<float> lat2(graph2.numVertices());
    std::vector<float> lng1(graph1.numVertices());
    std::vector<float> lng2(graph2.numVertices());
    FORALL_VERTICES(graph1, v) {
      lat1[v] = graph1.latLng(v).latInDeg();
      lng1[v] = graph1.latLng(v).lngInDeg();
    }
    FORALL_VERTICES(graph2, v) {
      lat2[v] = graph2.latLng(v).latInDeg();
      lng2[v] = graph2.latLng(v).lngInDeg();
    }
    VPTree tree1(lat1, lng1);
    VPTree tree2(lat2, lng2);

    // Check which vertices are inside the area of interest.
    boost::dynamic_bitset<> isVertexInsideArea1(graph1.numVertices());
    boost::dynamic_bitset<> isVertexInsideArea2(graph2.numVertices());
    isVertexInsideArea1.set();
    isVertexInsideArea2.set();
    if (!polyfile.empty()) {
      Area area;
      area.importFromOsmPolyFile(polyfile);
      const auto box = area.boundingBox();
      FORALL_VERTICES(graph1, v) {
        const Point p = {graph1.latLng(v).longitude(), graph1.latLng(v).latitude()};
        isVertexInsideArea1[v] = box.contains(p) && area.contains(p);
      }
      FORALL_VERTICES(graph2, v) {
        const Point p = {graph2.latLng(v).longitude(), graph2.latLng(v).latitude()};
        isVertexInsideArea2[v] = box.contains(p) && area.contains(p);
      }
    }

    // Match the vertices in both graphs with each other.
    std::vector<std::pair<int, int>> matchedVertices;
    FORALL_VERTICES(graph1, u) {
      if (isVertexInsideArea1[u]) {
        const int v = findNearestNeighbor(tree2, graph1.latLng(u), maxDist);
        if (v != INVALID_INDEX && isVertexInsideArea2[v]) {
          const int w = findNearestNeighbor(tree1, graph2.latLng(v), maxDist);
          if (w == u)
            matchedVertices.emplace_back(u, v);
        }
      }
    }

    std::cout << "Number of vertices in graph 1: " << graph1.numVertices() << std::endl;
    std::cout << "Number of vertices in graph 2: " << graph2.numVertices() << std::endl;
    std::cout << "Number of edges in graph 1: " << graph1.numEdges() << std::endl;
    std::cout << "Number of edges in graph 2: " << graph2.numEdges() << std::endl;
    const int numVerticesInsideArea1 = isVertexInsideArea1.count();
    const int numVerticesInsideArea2 = isVertexInsideArea2.count();
    const double matched1 = 100.0 * matchedVertices.size() / numVerticesInsideArea1;
    const double matched2 = 100.0 * matchedVertices.size() / numVerticesInsideArea2;
    if (!polyfile.empty()) {
      std::cout << "Vertices of interest in graph 1: " << numVerticesInsideArea1 << std::endl;
      std::cout << "Vertices of interest in graph 2: " << numVerticesInsideArea2 << std::endl;
    }
    std::cout << std::fixed << std::setprecision(2);
    std::cout << "Matched vertices in graph 1: " << std::setw(6) << matched1 << "%" << std::endl;
    std::cout << "Matched vertices in graph 2: " << std::setw(6) << matched2 << "%" << std::endl;

    // Open both output OD-files.
    std::ofstream out1(outfile1 + ".csv");
    std::ofstream out2(outfile2 + ".csv");
    if (!out1.good())
      throw std::invalid_argument("file cannot be opened -- '" + outfile1 + ".csv'");
    if (!out2.good())
      throw std::invalid_argument("file cannot be opened -- '" + outfile2 + ".csv'");
    out1 << "# Input graph: " << infile1 << std::endl;
    out2 << "# Input graph: " << infile2 << std::endl;
    out1 << "# Methodology: matched" << std::endl;
    out2 << "# Methodology: matched" << std::endl;
    out1 << "origin,destination" << std::endl;
    out2 << "origin,destination" << std::endl;

    // Pick the OD-pairs from the matched vertices uniformly at random.
    std::default_random_engine rand(clp.getValue<int>("s", 19900325));
    std::uniform_int_distribution<> dist(0, matchedVertices.size() - 1);
    for (int i = 0; i < numPairs; ++i) {
      const int o = dist(rand);
      const int d = dist(rand);
      out1 << matchedVertices[o].first << ',' << matchedVertices[d].first << std::endl;
      out2 << matchedVertices[o].second << ',' << matchedVertices[d].second << std::endl;
    }

    out1.close();
    out2.close();
  } catch (std::exception& e) {
    std::cerr << argv[0] << ": " << e.what() << std::endl;
    std::cerr << "Try '" << argv[0] <<" -help' for more information." << std::endl;
    return EXIT_FAILURE;
  }
  return EXIT_SUCCESS;
}


#ifndef INTERNAL_EXTRACTOR_EDGE_HPP
#define INTERNAL_EXTRACTOR_EDGE_HPP

#include "extractor/guidance/road_classification.hpp"
#include "extractor/guidance/turn_lane_types.hpp"
#include "extractor/node_based_edge.hpp"
#include "extractor/travel_mode.hpp"
#include "osrm/coordinate.hpp"
#include "util/strong_typedef.hpp"
#include "util/typedefs.hpp"

#include <boost/assert.hpp>
#include <mapbox/variant.hpp>
#include <utility>

namespace osrm
{
namespace extractor
{

namespace detail
{
// Use a single float to pack two positive values:
//  * by_edge: + sign, value >= 0, value used as the edge weight
//  * by_meter: - sign, value > 0, value used as a denominator in distance / value
struct ByEdgeOrByMeterValue
{
    struct ValueByEdge
    {
    } static const by_edge;
    struct ValueByMeter
    {
    } static const by_meter;

    ByEdgeOrByMeterValue() : value(0.f) {}

    ByEdgeOrByMeterValue(ValueByEdge, double input)
    {
        BOOST_ASSERT(input >= 0.f);
        value = static_cast<value_type>(input);
    }

    ByEdgeOrByMeterValue(ValueByMeter, double input)
    {
        BOOST_ASSERT(input > 0.f);
        value = -static_cast<value_type>(input);
    }

    double operator()(double distance) { return value >= 0 ? value : -distance / value; }

  private:
    using value_type = float;
    value_type value;
};
}

struct InternalExtractorEdge
{
    using WeightData = detail::ByEdgeOrByMeterValue;
    using DurationData = detail::ByEdgeOrByMeterValue;

    explicit InternalExtractorEdge()
        : result(MIN_OSM_NODEID,
                 MIN_OSM_NODEID,
                 SPECIAL_NODEID,
                 0,
                 0,
                 false, // forward
                 false, // backward
                 false, // roundabout
                 false, // circular
                 true,  // can be startpoint
                 false, // local access only
                 false, // split edge
                 TRAVEL_MODE_INACCESSIBLE,
                 guidance::TurnLaneType::empty,
                 guidance::RoadClassification()),
          weight_data(), duration_data()
    {
    }

    explicit InternalExtractorEdge(OSMNodeID source,
                                   OSMNodeID target,
                                   NodeID name_id,
                                   WeightData weight_data,
                                   DurationData duration_data,
                                   bool forward,
                                   bool backward,
                                   bool roundabout,
                                   bool circular,
                                   bool startpoint,
                                   bool restricted,
                                   bool is_split,
                                   TravelMode travel_mode,
                                   LaneDescriptionID lane_description,
                                   guidance::RoadClassification road_classification,
                                   util::Coordinate source_coordinate)
        : result(source,
                 target,
                 name_id,
                 0,
                 0,
                 forward,
                 backward,
                 roundabout,
                 circular,
                 startpoint,
                 restricted,
                 is_split,
                 travel_mode,
                 lane_description,
                 std::move(road_classification)),
          weight_data(std::move(weight_data)), duration_data(std::move(duration_data)),
          source_coordinate(std::move(source_coordinate))
    {
    }

    // data that will be written to disk
    NodeBasedEdgeWithOSM result;
    // intermediate edge weight
    WeightData weight_data;
    // intermediate edge duration
    DurationData duration_data;
    // coordinate of the source node
    util::Coordinate source_coordinate;

    // necessary static util functions for stxxl's sorting
    static InternalExtractorEdge min_osm_value()
    {
        return InternalExtractorEdge(MIN_OSM_NODEID,
                                     MIN_OSM_NODEID,
                                     SPECIAL_NODEID,
                                     WeightData(),
                                     DurationData(),
                                     false, // forward
                                     false, // backward
                                     false, // roundabout
                                     false, // circular
                                     true,  // can be startpoint
                                     false, // local access only
                                     false, // split edge
                                     TRAVEL_MODE_INACCESSIBLE,
                                     INVALID_LANE_DESCRIPTIONID,
                                     guidance::RoadClassification(),
                                     util::Coordinate());
    }
    static InternalExtractorEdge max_osm_value()
    {
        return InternalExtractorEdge(MAX_OSM_NODEID,
                                     MAX_OSM_NODEID,
                                     SPECIAL_NODEID,
                                     WeightData(),
                                     DurationData(),
                                     false, // forward
                                     false, // backward
                                     false, // roundabout
                                     false, // circular
                                     true,  // can be startpoint
                                     false, // local access only
                                     false, // split edge
                                     TRAVEL_MODE_INACCESSIBLE,
                                     INVALID_LANE_DESCRIPTIONID,
                                     guidance::RoadClassification(),
                                     util::Coordinate());
    }

    static InternalExtractorEdge min_internal_value()
    {
        auto v = min_osm_value();
        v.result.source = 0;
        v.result.target = 0;
        return v;
    }
    static InternalExtractorEdge max_internal_value()
    {
        auto v = max_osm_value();
        v.result.source = std::numeric_limits<NodeID>::max();
        v.result.target = std::numeric_limits<NodeID>::max();
        return v;
    }
};
}
}

#endif // INTERNAL_EXTRACTOR_EDGE_HPP


#ifndef _LDAPLUSPLUS_LDA_HPP_
#define _LDAPLUSPLUS_LDA_HPP_

#include <condition_variable>
#include <list>
#include <memory>
#include <mutex>
#include <vector>
#include <thread>
#include <tuple>

#include <Eigen/Core>

#include "ldaplusplus/events/Events.hpp"
#include "ldaplusplus/em/EStepInterface.hpp"
#include "ldaplusplus/em/MStepInterface.hpp"
#include "ldaplusplus/Parameters.hpp"

namespace ldaplusplus {

/**
 * LDA contains the logic of using an expectation step, a maximization step and
 * some model parameters to train and make use of an LDA model.
 *
 * 1. It is agnostic of the underlying implementations it uses and thus allows
 *    for experimentation through a common facade.
 * 2. It uses multiple threads to compute the time consuming expectation step.
 * 3. It aggregates all the events and redispatches them on the same thread
 *    through a single event dispatcher.
 * 4. It provides a very simple interface (borrowed from scikit-learn)
 */
template <typename Scalar = double>
class LDA
{
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> MatrixX;
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1> VectorX;

    public:
        /**
         * Create an LDA with the given model parameters, expectation and
         * maximization steps default iterations and worker threads.
         *
         * @param model_parameters A pointer to a struct containing the model
         *                         parameters (for instance ModelParameters and
         *                         SupervisedModelParameters)
         * @param e_step           A pointer to an expectation step
         *                         implementation
         * @param m_step           A pointer to a maximization step
         *                         implementation
         * @param iterations       The number of epochs to run when using
         *                         LDA::fit
         * @param workers          The number of worker threads to create for
         *                         computing the expectation step
         */
        LDA(
            std::shared_ptr<parameters::Parameters> model_parameters,
            std::shared_ptr<em::EStepInterface<Scalar> > e_step,
            std::shared_ptr<em::MStepInterface<Scalar> > m_step,
            size_t iterations = 20,
            size_t workers = 1
        );

        /**
         * Create a move constructor that doesn't try to copy or move mutexes.
         */
        LDA(LDA &&lda);

        /**
         * Compute a supervised topic model for word counts X and classes y.
         *
         * Perform as many EM iterations as configured and stop when reaching
         * max_iter_ or any other stopping criterion.
         *
         * An EigenClassificationCorpus will be created from the passed
         * parameters.
         *
         * @param X The word counts in column-major order
         * @param y The classes as integers
         */
        void fit(const Eigen::MatrixXi &X, const Eigen::VectorXi &y);

        /**
         * Compute an unsupervised topic model for word counts X.
         *
         * Perform as many EM iterations as configured and stop when reaching
         * max_iter_ or any other stopping criterion.
         *
         * An EigenCorpus will be created from the passed
         * parameters.
         *
         * @param X The word counts in column-major order
         */
        void fit(const Eigen::MatrixXi &X);

        /**
         * Perform a single EM iteration.
         *
         * An EigenClassificationCorpus will be created from the passed
         * parameters.
         *
         * @param X The word counts in column-major order
         * @param y The classes as integers
         */
        void partial_fit(const Eigen::MatrixXi &X, const Eigen::VectorXi &y);

        /**
         * Perform a single EM iteration.
         *
         * @param corpus The implementation of Corpus that contains the
         *               observed variables.
         */
        void partial_fit(std::shared_ptr<corpus::Corpus> corpus);

        /**
         * Run the expectation step and return the topic mixtures for the
         * documents defined by the word counts X.
         *
         * @param  X The word counts in column-major order
         * @return The variational parameter \f$\gamma\f$ for every document
         *         that approximates the count of words generated by each topic
         */
        MatrixX transform(const Eigen::MatrixXi &X);

        /**
         * Treat the SupervisedModelParameters::eta as a linear model and
         * compute the distances from the planes of the documents in the topic
         * space.
         *
         * Use LDA::transform to obtain the \f$\gamma\f$ for every document and
         * then assume that the \f$\eta\f$ parameters of the
         * SupervisedModelParameters are a linear model. Compute the dot
         * product between the normal vectors and the normalized topic mixtures
         * for each document. The more positive the value for a given class the
         * more confident is the model that a document belongs in this class.
         *
         * @param  X The word counts in column-major order
         * @return A matrix of class scores (positive => confident) for each
         *         document
         */
        MatrixX decision_function(const Eigen::MatrixXi &X);

        /**
         * Use the model to predict the class indexes for the word counts X.
         *
         * Use LDA::decision_function to get class scores and then compute the
         * *argmax* for every document.
         *
         * @param  X The word counts in column-major order
         * @return A matrix of class indexes (the predicted class for each
         *         document)
         */
        Eigen::VectorXi predict(const Eigen::MatrixXi &X);

        /**
         * Return both the class predictions and the transformed data using a
         * single LDA expectation step.
         *
         * @param  X The word counts in column-major order
         * @return A tuple containing a matrix with the \f$\gamma\f$
         *         variational parameters for every document and a vector
         *         containing the class index predictions for every document.
         */
        std::tuple<MatrixX, Eigen::VectorXi> transform_predict(const Eigen::MatrixXi &X);

        /**
         * Get the event dispatcher for this LDA instance.
         */
        std::shared_ptr<events::EventDispatcherInterface> get_event_dispatcher() {
            return event_dispatcher_;
        }

        /**
         * Get a constant reference to the model's parameters.
         */
        const std::shared_ptr<parameters::Parameters> model_parameters() {
            return model_parameters_;
        }

        template <typename P>
        const std::shared_ptr<P> model_parameters() {
            return std::static_pointer_cast<P>(model_parameters_);
        }

    protected:
        /**
         * Generate a Corpus from a pair of X, y matrices
         */
        std::shared_ptr<corpus::Corpus> get_corpus(
            const Eigen::MatrixXi &X,
            const Eigen::VectorXi &y
        );

        /**
         * Generate a Corpus from just the word count matrix.
         */
        std::shared_ptr<corpus::Corpus> get_corpus(const Eigen::MatrixXi &X);

        /**
         * Create a worker thread pool.
         */
        void create_worker_pool();

        /**
         * Destroy the worker thread pool
         */
        void destroy_worker_pool();

        /**
         * Forward the events generated in the worker threads to this event
         * dispatcher in this thread.
         */
        void process_worker_events() {
            std::static_pointer_cast<events::ThreadSafeEventDispatcher>(
                event_dispatcher_
            )->process_events();
        }

        /**
         * Extract the variational parameters and the document index from the
         * worker queue.
         */
        std::tuple<std::shared_ptr<parameters::Parameters>, size_t> extract_vp_from_queue();

        /**
         * A doc_e_step worker thread.
         */
        void doc_e_step_worker();

        /**
         * Implement the decision function using already transformed data.
         * Topic representations instead of BOW.
         *
         * @param X The \f$\gamma\f$ variational parameter for each document in
         *          a column-major ordered matrix.
         * @return A matrix of class scores (positive => confident) for each
         *         document
         */
        MatrixX decision_function(const MatrixX &X);

        /**
         * Transform the decision function to class predictions.
         */
        Eigen::VectorXi predict(const MatrixX &scores);

    private:
        /**
         * Pass the event dispatcher down to the implementations so that they
         * can communicate with the outside world.
         */
        void set_up_event_dispatcher();

        // The model parameters
        std::shared_ptr<parameters::Parameters> model_parameters_;

        // The LDA implementation
        std::shared_ptr<em::EStepInterface<Scalar> > e_step_;
        std::shared_ptr<em::MStepInterface<Scalar> > m_step_;

        // Member variables that affect the behaviour of fit
        size_t iterations_;

        // The thread related member variables
        std::vector<std::thread> workers_;
        std::mutex queue_in_mutex_;
        std::list<std::tuple<std::shared_ptr<corpus::Corpus>, size_t> > queue_in_;
        std::mutex queue_out_mutex_;
        std::condition_variable queue_out_cv_;
        std::list<std::tuple<std::shared_ptr<parameters::Parameters>, size_t> > queue_out_;

        // An event dispatcher that we will use to communicate with the
        // external components
        std::shared_ptr<events::EventDispatcherInterface> event_dispatcher_;
};

} // namespace ldaplusplus
#endif  // _LDAPLUSPLUS_LDA_HPP_


// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#include "CTCInferenceCodeletTestConnection.hpp"

#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/IPUModel.hpp>
#include <poplar/Type.hpp>

#include <poplibs_support/LogArithmetic.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/CTCInference.hpp>
#include <poplibs_test/Util.hpp>
#include <poputil/VertexTemplates.hpp>

#include <boost/multi_array.hpp>

#include <limits>

using namespace poplar;
using namespace poplar::program;
using namespace poplibs_test::ctc;
using namespace poplibs_test;
using namespace poplibs_test::util;
using namespace poplibs_support;
using namespace poputil;

namespace poplibs_test {
namespace ctc {

template <typename PartialsType>
std::tuple<BeamHistory, std::vector<BeamProbability<PartialsType>>,
           std::vector<unsigned>>
runUpdateCodelet(Graph &graph, TestDevice &device, DeviceType deviceType,
                 Type inType, Type partialsType,
                 const std::vector<Candidate<PartialsType>> &candidates,
                 unsigned timestep, const BeamHistory &beamHistory,
                 const std::vector<unsigned> &beamLengthIn,
                 const std::vector<BeamProbability<PartialsType>> &beamProbs,
                 unsigned blankClass, bool profile) {
  const auto target = graph.getTarget();

  const auto totalCandidates = candidates.size();
  const auto beamwidth = beamHistory.symbols.size();
  const auto maxT = beamHistory.symbols[0].size();

  auto beamAddend =
      graph.addVariable(UNSIGNED_INT, {maxT, beamwidth}, "beamAddend");
  auto beamParent =
      graph.addVariable(UNSIGNED_INT, {maxT, beamwidth}, "beamParent");
  auto beamLength =
      graph.addVariable(UNSIGNED_INT, {2 * beamwidth}, "beamLength");

  auto currentTimestep =
      graph.addConstant(UNSIGNED_INT, {}, timestep, "currentTimestep");
  auto dataLength = graph.addConstant(UNSIGNED_INT, {}, timestep);
  auto complete = graph.addVariable(UNSIGNED_INT, {}, "completeFlag");

  graph.setTileMapping(beamAddend, 0);
  graph.setTileMapping(beamParent, 0);
  graph.setTileMapping(beamLength, 0);

  graph.setTileMapping(currentTimestep, 0);
  graph.setTileMapping(dataLength, 0);
  graph.setTileMapping(complete, 0);

  auto cs = graph.addComputeSet("cs");
  auto vertex = graph.addVertex(
      cs, templateVertex("popnn::CTCUpdate", partialsType, UNSIGNED_INT));
  graph.setTileMapping(vertex, 0);

  graph.connect(vertex["beamAddend"], beamAddend.flatten());
  graph.connect(vertex["beamParent"], beamParent.flatten());
  graph.connect(vertex["beamLength"], beamLength);

  graph.connect(vertex["currentTimestep"], currentTimestep);
  graph.connect(vertex["dataLength"], dataLength);
  graph.connect(vertex["complete"], complete);

  graph.setInitialValue(vertex["beamwidth"], beamwidth);

  Sequence uploadProg, downloadProg;
  std::vector<std::pair<std::string, char *>> tmap;

  // Inputs
  auto rawCandidates = createAndConnectCandidates(
      graph, vertex, "candidate", partialsType, {totalCandidates}, uploadProg,
      downloadProg, tmap, true);

  std::vector<unsigned> candidateParentIn{};
  std::vector<unsigned> candidateAddendIn{};
  std::vector<float> candidateBeamProbNonBlankIn{};
  std::vector<float> candidateBeamProbBlankIn{};

  for (unsigned c = 0; c < totalCandidates; c++) {
    candidateParentIn.push_back(candidates[c].beam);
    candidateAddendIn.push_back(candidates[c].addend);
    candidateBeamProbNonBlankIn.push_back(candidates[c].pnb);
    candidateBeamProbBlankIn.push_back(candidates[c].pb);
  }

  copy(target, candidateParentIn, UNSIGNED_INT, rawCandidates.parent.get());
  copy(target, candidateAddendIn, UNSIGNED_INT, rawCandidates.addend.get());
  copy(target, candidateBeamProbNonBlankIn, partialsType,
       rawCandidates.probNonBlank.get());
  copy(target, candidateBeamProbBlankIn, partialsType,
       rawCandidates.probBlank.get());

  // InOut
  std::unique_ptr<char[]> rawBeamAddend, rawBeamParent, rawBeamLength,
      rawComplete;

  rawBeamAddend = allocateHostMemoryForTensor(beamAddend, "beamAddend", graph,
                                              uploadProg, downloadProg, tmap);
  rawBeamParent = allocateHostMemoryForTensor(beamParent, "beamParent", graph,
                                              uploadProg, downloadProg, tmap);
  rawBeamLength = allocateHostMemoryForTensor(beamLength, "beamLength", graph,
                                              uploadProg, downloadProg, tmap);
  rawComplete = allocateHostMemoryForTensor(complete, "complete", graph,
                                            uploadProg, downloadProg, tmap);

  // TODO last beam output isn't verified by this test, but will probably be
  // optimised out in future
  auto rawBeamProbs = createAndConnectBeamProbs(
      graph, vertex, partialsType, {beamwidth},
      BeamScalars::BLANK_AND_NON_BLANK, uploadProg, downloadProg, tmap);

  std::vector<unsigned> beamAddendIn{};
  std::vector<unsigned> beamParentIn{};
  std::vector<double> beamProbNonBlankIn{};
  std::vector<double> beamProbBlankIn{};
  std::vector<double> beamProbTotalIn{};
  std::vector<unsigned> completeIn = {0};

  for (unsigned t = 0; t < maxT; t++) {
    for (unsigned b = 0; b < beamwidth; b++) {
      beamAddendIn.push_back(beamHistory.symbols[b][t]);
      if (beamHistory.parents[b][t]) {
        beamParentIn.push_back(*beamHistory.parents[b][t]);
      } else {
        // TODO need this column in real impl?
        beamParentIn.push_back(std::numeric_limits<unsigned>::max());
      }
    }
  }

  for (unsigned b = 0; b < beamwidth; b++) {
    beamProbNonBlankIn.push_back(beamProbs[b].pnb);
    beamProbBlankIn.push_back(beamProbs[b].pb);
    beamProbTotalIn.push_back(log::add(beamProbs[b].pb, beamProbs[b].pnb));
  }

  copy(target, beamAddendIn, UNSIGNED_INT, rawBeamAddend.get());
  copy(target, beamParentIn, UNSIGNED_INT, rawBeamParent.get());
  copy(target, beamProbNonBlankIn, partialsType, rawBeamProbs.pnb.get());
  copy(target, beamProbBlankIn, partialsType, rawBeamProbs.pb.get());
  copy(target, beamProbTotalIn, partialsType, rawBeamProbs.pTotal.get());
  copy(target, beamLengthIn, UNSIGNED_INT, rawBeamLength.get());
  copy(target, completeIn, UNSIGNED_INT, rawComplete.get());

  OptionFlags engineOptions;
  if (profile) {
    engineOptions.set("debug.instrumentCompute", "true");
  }
  Sequence prog;
  prog.add(Execute(cs));
  Engine engine(graph, Sequence{uploadProg, prog, downloadProg}, engineOptions);
  attachStreams(engine, tmap);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run();
  });

  std::vector<unsigned> beamAddendOut(maxT * beamwidth);
  std::vector<unsigned> beamParentOut(maxT * beamwidth);
  // TODO partialsType == float
  std::vector<float> beamProbNonBlankOut(beamwidth);
  std::vector<float> beamProbBlankOut(beamwidth);
  std::vector<unsigned> beamLengthOut(2 * beamwidth);

  copy(target, UNSIGNED_INT, rawBeamAddend.get(), beamAddendOut);
  copy(target, UNSIGNED_INT, rawBeamParent.get(), beamParentOut);
  copy(target, partialsType, rawBeamProbs.pnb.get(), beamProbNonBlankOut);
  copy(target, partialsType, rawBeamProbs.pb.get(), beamProbBlankOut);
  copy(target, UNSIGNED_INT, rawBeamLength.get(), beamLengthOut);

  if (profile && deviceType != DeviceType::Cpu) {
    engine.printProfileSummary(std::cout,
                               OptionFlags{{"showExecutionSteps", "true"}});
  }

  std::vector<BeamProbability<float>> beamProbOut;
  BeamHistory beamHistoryOut(beamwidth, maxT);
  beamHistoryOut.nextIndexToAssign = timestep + 1;

  for (unsigned b = 0; b < beamwidth; b++) {
    beamProbOut.push_back({beamProbNonBlankOut[b], beamProbBlankOut[b]});
    for (unsigned t = 0; t < maxT; t++) {
      beamHistoryOut.symbols[b][t] = beamAddendOut[b + beamwidth * t];
      beamHistoryOut.parents[b][t] = beamParentOut[b + beamwidth * t];
    }
  }
  return {beamHistoryOut, beamProbOut, beamLengthOut};
}

template std::tuple<BeamHistory, std::vector<BeamProbability<float>>,
                    std::vector<unsigned>>
runUpdateCodelet(Graph &graph, TestDevice &device, DeviceType deviceType,
                 Type inType, Type partialsType,
                 const std::vector<Candidate<float>> &candidates,
                 unsigned timestep, const BeamHistory &beamHistory,
                 const std::vector<unsigned> &beamLengthIn,
                 const std::vector<BeamProbability<float>> &beamProbs,
                 unsigned blankClass, bool profile);

} // namespace ctc
} // namespace poplibs_test