// Copyright (c) 2020 Graphcore Ltd. All rights reserved.

#define BOOST_TEST_MODULE BlockSparseTest
#include <boost/test/unit_test.hpp>
#include <cstdlib>
#include <poplar/IPUModel.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/Util.hpp>
#include <poplin/codelets.hpp>
#include <popops/codelets.hpp>
#include <poputil/TileMapping.hpp>
#include <random>
#include <vector>

#include "popsparse/BSMatrix.hpp"
#include "popsparse/HyperGraphBlock.hpp"
#include "popsparse/experimental/BlockSparseMatMul.hpp"

using namespace poplar;
using namespace poplar::program;
using namespace popsparse::experimental;
using namespace poplibs_support;

// TODO: Move to some kind of utils
void getSparseMatrixBlocks(int rows, int cols, int rowsInBlock, int colsInBlock,
                           std::vector<unsigned char> &sparsity,
                           const boost::multi_array<float, 2> &denseMat,
                           boost::multi_array<float, 2> &blockData) {
  const int blockRows = rows / rowsInBlock;
  const int blockCols = cols / colsInBlock;

  int blockCount = 0;
  for (int br = 0; br < blockRows; br++) {
    for (int bc = 0; bc < blockCols; bc++) {
      if (sparsity[br * blockCols + bc] == 1) {
        blockCount++;
      }
    }
  }
  blockData.resize(boost::extents[blockCount][rowsInBlock * colsInBlock]);

  blockCount = 0;
  for (int br = 0; br < blockRows; br++) {
    for (int bc = 0; bc < blockCols; bc++) {
      if (sparsity[br * blockCols + bc] == 1) {
        int rowStart = br * rowsInBlock;
        int colStart = bc * colsInBlock;
        int rowEnd = rowStart + rowsInBlock;
        int colEnd = colStart + colsInBlock;
        int index = 0;
        for (int r = rowStart; r < rowEnd; r++) {
          for (int c = colStart; c < colEnd; c++) {
            blockData[blockCount][index++] = denseMat[r][c];
          }
        }
        blockCount++;
      }
    }
  }
}

boost::multi_array<float, 2> transpose(const boost::multi_array<float, 2> mat) {
  std::size_t rows = mat.shape()[0];
  std::size_t cols = mat.shape()[1];
  boost::multi_array<float, 2> transposedMat(boost::extents[cols][rows]);
  for (std::size_t r = 0; r < rows; ++r) {
    for (std::size_t c = 0; c < cols; ++c) {
      transposedMat[c][r] = mat[r][c];
    }
  }
  return transposedMat;
}

#define USE_RANDOM_VALUES 1

// TODO: Move to some kind of utils
void populateMatrixData1(boost::multi_array<float, 2> &host, int rows, int cols,
                         int rowsInBlock, int colsInBlock,
#if !USE_RANDOM_VALUES
                         int rowInBlockNz = 0, int colInBlockNz = 0,
                         int blockRowNz = -1, int blockColNz = -1
#else
                         int = 0, int = 0, int = 0, int = 0
#endif
) {

#if USE_RANDOM_VALUES
  std::mt19937 randomEngine;
  randomEngine.seed(102302);
#endif
  host.resize(boost::extents[rows][cols]);

  // We make only one element of a block as non-zero for simplicity
  // hostA is row-major
  for (int r = 0; r < rows; ++r) {
    for (int c = 0; c < cols; ++c) {
#if USE_RANDOM_VALUES
      host[r][c] =
          static_cast<float>(randomEngine()) / static_cast<float>(INT_MAX);
#else
      int rb = r % rowsInBlock;
      int cb = c % colsInBlock;
      int bc = c / colsInBlock;
      if (rb == rowInBlockNz && cb == colInBlockNz &&
          (blockRowNz < 0 || br == blockRowNz) &&
          (blockColNz < 0 || bc == blockColNz)) {
        host[r][c] = 1.0f;
      } else {
        host[r][c] = 0.0f;
      }
#endif
    }
  }
}

void populateMatrixDataWSparsity(boost::multi_array<float, 2> &host, int rows,
                                 int cols, int rowsInBlock, int colsInBlock,
                                 const unsigned char *sparsity) {

  std::mt19937 randomEngine;
  randomEngine.seed(102302);
  host.resize(boost::extents[rows][cols]);

  int blockCols = cols / colsInBlock;
  for (int r = 0; r < rows; ++r) {
    int br = r / rowsInBlock;
    for (int c = 0; c < cols; ++c) {
      int bc = c / colsInBlock;
      int idxSparsity = br * blockCols + bc;
      if (sparsity[idxSparsity]) {
        host[r][c] =
            static_cast<float>(randomEngine()) / static_cast<float>(INT_MAX);
      } else {
        host[r][c] = 0.0f;
      }
    }
  }
}

void populateMatricesData1(boost::multi_array<float, 2> &hostA,
                           boost::multi_array<float, 2> &hostB, int rowsA,
                           int colsA, int colsB, int rowsInBlockA,
                           int colsInBlockA, int colsInBlockB,
                           int rowInBlockNzA = 0, int colInBlockNzA = 0,
                           int blockRowNzA = -1, int blockColNzA = -1) {

  const int rowsB = colsA;
  const int rowsInBlockB = colsInBlockA;

  hostA.resize(boost::extents[rowsA][colsA]);
  hostB.resize(boost::extents[rowsB][colsB]);

  const int rowInBlockNzB = colInBlockNzA;
  const int colInBlockNzB = rowInBlockNzA;
  const int blockRowNzB = blockColNzA;
  const int blockColNzB = blockRowNzA;

  populateMatrixData1(hostA, rowsA, colsA, rowsInBlockA, colsInBlockA,
                      rowInBlockNzA, colInBlockNzA, blockRowNzA, blockColNzA);

  populateMatrixData1(hostB, rowsB, colsB, rowsInBlockB, colsInBlockB,
                      rowInBlockNzB, colInBlockNzB, blockRowNzB, blockColNzB);
}

void populateMatricesDataWSparsity(boost::multi_array<float, 2> &hostA,
                                   boost::multi_array<float, 2> &hostB,
                                   int rowsA, int colsA, int colsB,
                                   int rowsInBlockA, int colsInBlockA,
                                   int colsInBlockB,
                                   const unsigned char *sparsity) {

  const int rowsB = colsA;
  const int rowsInBlockB = colsInBlockA;

  hostA.resize(boost::extents[rowsA][colsA]);
  hostB.resize(boost::extents[rowsB][colsB]);

  populateMatrixData1(hostA, rowsA, colsA, rowsInBlockA, colsInBlockA);

  populateMatrixDataWSparsity(hostB, rowsB, colsB, rowsInBlockB, colsInBlockB,
                              sparsity);
}

void populateMatricesData1(const std::array<int, 3> &dim,
                           const std::array<int, 3> &block_size,
                           boost::multi_array<float, 2> &lhsMatrix,
                           boost::multi_array<float, 2> &rhsMatrix,
                           const std::array<int, 2> &nzBlockElem,
                           const std::array<int, 2> &nzBlock) {
  populateMatricesData1(lhsMatrix, rhsMatrix, dim[0], dim[1], dim[2],
                        block_size[0], block_size[1], block_size[3],
                        nzBlockElem[0], nzBlockElem[1], nzBlock[0], nzBlock[1]);
}

/*
Computes matrix multiplication.
All matrices are row-major
*/
std::vector<std::vector<float>>
matMul(const boost::multi_array<float, 2> &hostA,
       const boost::multi_array<float, 2> &hostB) {
  std::size_t rowsA = hostA.shape()[0];
  std::size_t colsA = hostA.shape()[1];
  std::size_t colsB = hostB.shape()[1];
  std::vector<std::vector<float>> hostC;
  for (std::size_t r = 0; r < rowsA; ++r) {
    hostC.push_back({});
    for (std::size_t c = 0; c < colsB; ++c) {
      float sum = 0.0f;
      for (std::size_t mulc = 0; mulc < colsA; ++mulc) {
        sum += hostA[r][mulc] * hostB[mulc][c];
      }
      hostC[r].push_back(sum);
    }
  }
  return hostC;
}

void checkDenseResult(
    const poplar::Type &dataType, int blockRowsC, int blockColsC,
    int rowsInBlockC, int colsInBlockC,
    const std::vector<std::vector<float>> &hostC,
    const std::vector<boost::multi_array<float, 2>> &blocksHostCParts) {
  const float epsilon = (dataType == FLOAT ? 0.001f : 0.1f);
  for (int br = 0, r = 0; br < blockRowsC; ++br) {
    for (int rb = 0; rb < rowsInBlockC; ++rb, ++r) {
      for (int bc = 0, c = 0; bc < blockColsC; ++bc) {
        int outBlockIdx = br * blockColsC + bc;
        for (int cb = 0; cb < colsInBlockC; ++cb, ++c) {
          float val = blocksHostCParts[outBlockIdx][rb][cb];
          float err = fabs(val - hostC[r][c]);
          if (err > epsilon)
            BOOST_TEST(err <= epsilon);
        }
      }
    }
  }
}

void checkDenseResult(const poplar::Type &dataType, int rowsC, int colsC,
                      const std::vector<std::vector<float>> &hostC,
                      const boost::multi_array<float, 2> &blocksHostC) {
  const float epsilon = (dataType == FLOAT ? 0.001f : 0.1f);
  for (int r = 0; r < rowsC; ++r) {
    for (int c = 0; c < colsC; ++c) {
      float val = blocksHostC[r][c];
      float err = fabs(val - hostC[r][c]);
      if (err > epsilon) {
        BOOST_TEST(err <= epsilon);
      }
    }
  }
}

void checkSparseResult(
    const poplar::Type &dataType, int blockRowsC, int blockColsC,
    int rowsInBlockC, int colsInBlockC,
    const std::vector<std::vector<float>> &hostC,
    const std::vector<boost::multi_array<float, 2>> &blocksHostCParts,
    const std::vector<unsigned char> &sparsityC) {
  const float epsilon = (dataType == FLOAT ? 0.001f : 0.05f);
  for (int br = 0, r = 0; br < blockRowsC; ++br) {
    for (int rb = 0; rb < rowsInBlockC; ++rb, ++r) {
      for (int bc = 0, c = 0; bc < blockColsC; ++bc) {
        int outBlockIdx = br * blockColsC + bc;
        for (int cb = 0; cb < colsInBlockC; ++cb, ++c) {
          float val = blocksHostCParts[outBlockIdx][rb][cb];
          float valTruth = hostC[r][c];
          if (sparsityC[br * blockColsC + bc] == 0) {
            valTruth = 0.0f;
          }
          float err = fabs(val - valTruth);
          if (err > epsilon)
            BOOST_TEST(err <= epsilon);
        }
      }
    }
  }
}

void checkSparseResult(const poplar::Type &dataType, int blockRowsC,
                       int blockColsC, int rowsInBlockC, int colsInBlockC,
                       const std::vector<std::vector<float>> &hostC,
                       const boost::multi_array<float, 2> &blocksHostC,
                       const std::vector<unsigned char> &sparsityC) {
  const float epsilon = (dataType == FLOAT ? 0.001f : 0.05f);
  int idxDense = 0;
  int idxSparse = 0;
  for (int br = 0; br < blockRowsC; ++br) {
    for (int bc = 0; bc < blockColsC; ++bc, ++idxDense) {
      if (sparsityC[idxDense]) {
        int idxInBlock = 0;
        for (int rb = 0; rb < rowsInBlockC; ++rb) {
          int r = br * rowsInBlockC + rb;
          for (int cb = 0; cb < colsInBlockC; ++cb, ++idxInBlock) {
            int c = bc * colsInBlockC + cb;
            float val = blocksHostC[idxSparse][idxInBlock];
            float valTruth = hostC[r][c];
            float err = fabs(val - valTruth);
            if (err > epsilon) {
              BOOST_TEST(err <= epsilon);
            }
          }
        }
        ++idxSparse;
      }
    }
  }
}

/*
Testing BlockSparseMatrix constructor and getBlockIdMatrix() method
*/
BOOST_AUTO_TEST_CASE(BlockSparseMatrix_test) {
  const int rows = 6;
  const int cols = 6;
  const int rows_block = 2;
  const int cols_block = 2;
  const int block_rows = rows / rows_block;
  const int block_cols = cols / cols_block;
  std::vector<unsigned char> sparsity(block_rows * block_cols, 0);
  sparsity[0 + 1] = 1;
  sparsity[3 + 0] = 1;
  sparsity[3 + 2] = 1;
  sparsity[6 + 1] = 1;

  BlockSparseMatrix bs1(rows, cols, rows_block, cols_block, false,
                        sparsity.data());

  std::vector<std::vector<int>> blockIdMatrix = bs1.getBlockIdMatrix();

  std::vector<std::vector<int>> blockIdMatrix_expected(3,
                                                       std::vector<int>(3, -1));
  blockIdMatrix_expected[0][1] = 0;
  blockIdMatrix_expected[1][0] = 1;
  blockIdMatrix_expected[1][2] = 2;
  blockIdMatrix_expected[2][1] = 3;

  BOOST_TEST(blockIdMatrix == blockIdMatrix_expected);
}

namespace popsparse {
namespace experimental {

class HyperGraphBlockTest : public HyperGraphBlock {

public:
  HyperGraphBlockTest(BlockMatrix &A, BlockMatrix &B, poplar::Type inDataTypeIn,
                      poplar::Type outDataTypeIn,
                      poplar::Type partialDataTypeIn, int nTileIn,
                      int nTargetNodesVPerTileIn = TARGET_V_NODES_PER_TILE)
      : HyperGraphBlock(A, B, inDataTypeIn, outDataTypeIn, partialDataTypeIn,
                        nTileIn, nTargetNodesVPerTileIn) {}

  virtual ~HyperGraphBlockTest() = default;

protected:
  void partitionGraph() override {}

private:
  void logTileAssignment(const std::vector<int> &tileAssignment) override {}
};

} // namespace experimental
} // namespace popsparse

/*
Testing Hypergraph for MatMul nodes and edges - no reduction case

nodeA_   nodeB_   nodeC_
0 1    x 4 5    =  8  9
2 3      6 7      10 11

nodeV_
id  idxA_ idxB_
12  [0,1] [0,2]
13  [0,1] [1,3]
14  [2,3] [0,2]
15  [2,3] [1,3]

*/
BOOST_AUTO_TEST_CASE(HyperGraph_testMatMulNoReduction) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);

  const int rows = 2;
  const int cols = 2;
  const int rows_block = 1;
  const int cols_block = 1;
  const int block_rows = rows / rows_block;
  const int block_cols = cols / cols_block;

  std::vector<unsigned char> sparsityB(block_rows * block_cols, 1);

  BlockDenseMatrix A(rows, cols, rows_block, cols_block, false);

  BlockSparseMatrix B(rows, cols, rows_block, cols_block, false,
                      sparsityB.data());

  // No reduction - 1 tile
  HyperGraphBlockTest hg(A, B, FLOAT, FLOAT, FLOAT, 1, 1);

  hg.createGraphMatMul(graph, "C");

  auto &nodeA = hg.getNodeA();
  auto &nodeB = hg.getNodeB();
  auto &nodeC = hg.getNodeC();
  auto &nodeV = hg.getNodeV();

  auto &edgeA = hg.getEdgeA();
  auto &edgeB = hg.getEdgeB();
  auto &edgeC = hg.getEdgeC();

  BOOST_TEST(nodeA.size() == 4U);
  BOOST_TEST(nodeB.size() == 4U);
  BOOST_TEST(nodeC.size() == 4U);
  BOOST_TEST(nodeV.size() == 4U);

  BOOST_TEST(edgeA.size() == 4U);
  BOOST_TEST(edgeB.size() == 4U);
  BOOST_TEST(edgeC.size() == 4U);

  BOOST_TEST(nodeV[0].idxA.size() == 2U);
  BOOST_TEST(nodeV[0].idxB.size() == 2U);
  BOOST_TEST(nodeV[0].idxA == std::vector<unsigned>({0, 1}));
  BOOST_TEST(nodeV[0].idxB == std::vector<unsigned>({0, 2}));

  BOOST_TEST(nodeV[1].idxA.size() == 2U);
  BOOST_TEST(nodeV[1].idxB.size() == 2U);
  BOOST_TEST(nodeV[1].idxA == std::vector<unsigned>({0, 1}));
  BOOST_TEST(nodeV[1].idxB == std::vector<unsigned>({1, 3}));

  BOOST_TEST(nodeV[2].idxA.size() == 2U);
  BOOST_TEST(nodeV[2].idxB.size() == 2U);
  BOOST_TEST(nodeV[2].idxA == std::vector<unsigned>({2, 3}));
  BOOST_TEST(nodeV[2].idxB == std::vector<unsigned>({0, 2}));

  BOOST_TEST(nodeV[3].idxA.size() == 2U);
  BOOST_TEST(nodeV[3].idxB.size() == 2U);
  BOOST_TEST(nodeV[3].idxA == std::vector<unsigned>({2, 3}));
  BOOST_TEST(nodeV[3].idxB == std::vector<unsigned>({1, 3}));

  BOOST_TEST(edgeA[0].in.size() == 1U);
  BOOST_TEST(edgeA[0].in == std::vector<unsigned>({nodeA[0].id}));
  BOOST_TEST(edgeA[0].out.size() == 2U);
  BOOST_TEST(edgeA[0].out == std::vector<unsigned>({nodeV[0].id, nodeV[1].id}));

  BOOST_TEST(edgeA[1].in.size() == 1U);
  BOOST_TEST(edgeA[1].in == std::vector<unsigned>({nodeA[1].id}));
  BOOST_TEST(edgeA[1].out.size() == 2U);
  BOOST_TEST(edgeA[1].out == std::vector<unsigned>({nodeV[0].id, nodeV[1].id}));

  BOOST_TEST(edgeA[2].in.size() == 1U);
  BOOST_TEST(edgeA[2].in == std::vector<unsigned>({nodeA[2].id}));
  BOOST_TEST(edgeA[2].out.size() == 2U);
  BOOST_TEST(edgeA[2].out == std::vector<unsigned>({nodeV[2].id, nodeV[3].id}));

  BOOST_TEST(edgeA[3].in.size() == 1U);
  BOOST_TEST(edgeA[3].in == std::vector<unsigned>({nodeA[3].id}));
  BOOST_TEST(edgeA[3].out.size() == 2U);
  BOOST_TEST(edgeA[3].out == std::vector<unsigned>({nodeV[2].id, nodeV[3].id}));

  BOOST_TEST(edgeB[0].in.size() == 1U);
  BOOST_TEST(edgeB[0].in == std::vector<unsigned>({nodeB[0].id}));
  BOOST_TEST(edgeB[0].out.size() == 2U);
  BOOST_TEST(edgeB[0].out == std::vector<unsigned>({nodeV[0].id, nodeV[2].id}));

  BOOST_TEST(edgeB[1].in.size() == 1U);
  BOOST_TEST(edgeB[1].in == std::vector<unsigned>({nodeB[1].id}));
  BOOST_TEST(edgeB[1].out.size() == 2U);
  BOOST_TEST(edgeB[1].out == std::vector<unsigned>({nodeV[1].id, nodeV[3].id}));

  BOOST_TEST(edgeB[2].in.size() == 1U);
  BOOST_TEST(edgeB[2].in == std::vector<unsigned>({nodeB[2].id}));
  BOOST_TEST(edgeB[2].out.size() == 2U);
  BOOST_TEST(edgeB[2].out == std::vector<unsigned>({nodeV[0].id, nodeV[2].id}));

  BOOST_TEST(edgeB[3].in.size() == 1U);
  BOOST_TEST(edgeB[3].in == std::vector<unsigned>({nodeB[3].id}));
  BOOST_TEST(edgeB[3].out.size() == 2U);
  BOOST_TEST(edgeB[3].out == std::vector<unsigned>({nodeV[1].id, nodeV[3].id}));

  BOOST_TEST(edgeC[0].out.size() == 1U);
  BOOST_TEST(edgeC[0].out == std::vector<unsigned>({nodeC[0].id}));
  BOOST_TEST(edgeC[0].in.size() == 1U);
  BOOST_TEST(edgeC[0].in == std::vector<unsigned>({nodeV[0].id}));

  BOOST_TEST(edgeC[1].out.size() == 1U);
  BOOST_TEST(edgeC[1].out == std::vector<unsigned>({nodeC[1].id}));
  BOOST_TEST(edgeC[1].in.size() == 1U);
  BOOST_TEST(edgeC[1].in == std::vector<unsigned>({nodeV[1].id}));

  BOOST_TEST(edgeC[2].out.size() == 1U);
  BOOST_TEST(edgeC[2].out == std::vector<unsigned>({nodeC[2].id}));
  BOOST_TEST(edgeC[2].in.size() == 1U);
  BOOST_TEST(edgeC[2].in == std::vector<unsigned>({nodeV[2].id}));

  BOOST_TEST(edgeC[3].out.size() == 1U);
  BOOST_TEST(edgeC[3].out == std::vector<unsigned>({nodeC[3].id}));
  BOOST_TEST(edgeC[3].in.size() == 1U);
  BOOST_TEST(edgeC[3].in == std::vector<unsigned>({nodeV[3].id}));
}

/*
Testing Hypergraph for MatMul nodes and edges - reduction case

nodeA_   nodeB_   nodeC_
0 1    x 4 5    =  8  9
2 3      6 7      10 11

nodeV_
id  idxA_ idxB_
12  [0] [0]
13  [1] [2]
14  [0] [1]
15  [1] [3]
16  [2] [0]
17  [3] [2]
18  [2] [1]
19  [3] [3]

*/
BOOST_AUTO_TEST_CASE(HyperGraph_testMatMulWReduction) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);

  const int rows = 2;
  const int cols = 2;
  const int rows_block = 1;
  const int cols_block = 1;
  const int block_rows = rows / rows_block;
  const int block_cols = cols / cols_block;

  std::vector<unsigned char> sparsityB(block_rows * block_cols, 1);

  BlockDenseMatrix A(rows, cols, rows_block, cols_block, false);

  BlockSparseMatrix B(rows, cols, rows_block, cols_block, false,
                      sparsityB.data());

  // Reduction - all tiles
  HyperGraphBlockTest hg(A, B, FLOAT, FLOAT, FLOAT, target.getNumTiles());

  hg.createGraphMatMul(graph, "C");

  auto &nodeA = hg.getNodeA();
  auto &nodeB = hg.getNodeB();
  auto &nodeC = hg.getNodeC();
  auto &nodeV = hg.getNodeV();

  auto &edgeA = hg.getEdgeA();
  auto &edgeB = hg.getEdgeB();
  auto &edgeC = hg.getEdgeC();

  BOOST_TEST(nodeA.size() == 4U);
  BOOST_TEST(nodeB.size() == 4U);
  BOOST_TEST(nodeC.size() == 4U);
  BOOST_TEST(nodeV.size() == 8U);

  BOOST_TEST(edgeA.size() == 4U);
  BOOST_TEST(edgeB.size() == 4U);
  BOOST_TEST(edgeC.size() == 4U);

  BOOST_TEST(nodeV[0].idxA.size() == 1U);
  BOOST_TEST(nodeV[0].idxB.size() == 1U);
  BOOST_TEST(nodeV[0].idxA == std::vector<unsigned>({0}));
  BOOST_TEST(nodeV[0].idxB == std::vector<unsigned>({0}));

  BOOST_TEST(nodeV[1].idxA.size() == 1U);
  BOOST_TEST(nodeV[1].idxB.size() == 1U);
  BOOST_TEST(nodeV[1].idxA == std::vector<unsigned>({1}));
  BOOST_TEST(nodeV[1].idxB == std::vector<unsigned>({2}));

  BOOST_TEST(nodeV[2].idxA.size() == 1U);
  BOOST_TEST(nodeV[2].idxB.size() == 1U);
  BOOST_TEST(nodeV[2].idxA == std::vector<unsigned>({0}));
  BOOST_TEST(nodeV[2].idxB == std::vector<unsigned>({1}));

  BOOST_TEST(nodeV[3].idxA.size() == 1U);
  BOOST_TEST(nodeV[3].idxB.size() == 1U);
  BOOST_TEST(nodeV[3].idxA == std::vector<unsigned>({1}));
  BOOST_TEST(nodeV[3].idxB == std::vector<unsigned>({3}));

  BOOST_TEST(nodeV[4].idxA.size() == 1U);
  BOOST_TEST(nodeV[4].idxB.size() == 1U);
  BOOST_TEST(nodeV[4].idxA == std::vector<unsigned>({2}));
  BOOST_TEST(nodeV[4].idxB == std::vector<unsigned>({0}));

  BOOST_TEST(nodeV[5].idxA.size() == 1U);
  BOOST_TEST(nodeV[5].idxB.size() == 1U);
  BOOST_TEST(nodeV[5].idxA == std::vector<unsigned>({3}));
  BOOST_TEST(nodeV[5].idxB == std::vector<unsigned>({2}));

  BOOST_TEST(nodeV[6].idxA.size() == 1U);
  BOOST_TEST(nodeV[6].idxB.size() == 1U);
  BOOST_TEST(nodeV[6].idxA == std::vector<unsigned>({2}));
  BOOST_TEST(nodeV[6].idxB == std::vector<unsigned>({1}));

  BOOST_TEST(nodeV[7].idxA.size() == 1U);
  BOOST_TEST(nodeV[7].idxB.size() == 1U);
  BOOST_TEST(nodeV[7].idxA == std::vector<unsigned>({3}));
  BOOST_TEST(nodeV[7].idxB == std::vector<unsigned>({3}));

  BOOST_TEST(edgeA[0].in.size() == 1U);
  BOOST_TEST(edgeA[0].in == std::vector<unsigned>({nodeA[0].id}));
  BOOST_TEST(edgeA[0].out.size() == 2U);
  BOOST_TEST(edgeA[0].out == std::vector<unsigned>({nodeV[0].id, nodeV[2].id}));

  BOOST_TEST(edgeA[1].in.size() == 1U);
  BOOST_TEST(edgeA[1].in == std::vector<unsigned>({nodeA[1].id}));
  BOOST_TEST(edgeA[1].out.size() == 2U);
  BOOST_TEST(edgeA[1].out == std::vector<unsigned>({nodeV[1].id, nodeV[3].id}));

  BOOST_TEST(edgeA[2].in.size() == 1U);
  BOOST_TEST(edgeA[2].in == std::vector<unsigned>({nodeA[2].id}));
  BOOST_TEST(edgeA[2].out.size() == 2U);
  BOOST_TEST(edgeA[2].out == std::vector<unsigned>({nodeV[4].id, nodeV[6].id}));

  BOOST_TEST(edgeA[3].in.size() == 1U);
  BOOST_TEST(edgeA[3].in == std::vector<unsigned>({nodeA[3].id}));
  BOOST_TEST(edgeA[3].out.size() == 2U);
  BOOST_TEST(edgeA[3].out == std::vector<unsigned>({nodeV[5].id, nodeV[7].id}));

  BOOST_TEST(edgeB[0].in.size() == 1U);
  BOOST_TEST(edgeB[0].in == std::vector<unsigned>({nodeB[0].id}));
  BOOST_TEST(edgeB[0].out.size() == 2U);
  BOOST_TEST(edgeB[0].out == std::vector<unsigned>({nodeV[0].id, nodeV[4].id}));

  BOOST_TEST(edgeB[1].in.size() == 1U);
  BOOST_TEST(edgeB[1].in == std::vector<unsigned>({nodeB[1].id}));
  BOOST_TEST(edgeB[1].out.size() == 2U);
  BOOST_TEST(edgeB[1].out == std::vector<unsigned>({nodeV[2].id, nodeV[6].id}));

  BOOST_TEST(edgeB[2].in.size() == 1U);
  BOOST_TEST(edgeB[2].in == std::vector<unsigned>({nodeB[2].id}));
  BOOST_TEST(edgeB[2].out.size() == 2U);
  BOOST_TEST(edgeB[2].out == std::vector<unsigned>({nodeV[1].id, nodeV[5].id}));

  BOOST_TEST(edgeB[3].in.size() == 1U);
  BOOST_TEST(edgeB[3].in == std::vector<unsigned>({nodeB[3].id}));
  BOOST_TEST(edgeB[3].out.size() == 2U);
  BOOST_TEST(edgeB[3].out == std::vector<unsigned>({nodeV[3].id, nodeV[7].id}));

  BOOST_TEST(edgeC[0].out.size() == 1U);
  BOOST_TEST(edgeC[0].out == std::vector<unsigned>({nodeC[0].id}));
  BOOST_TEST(edgeC[0].in.size() == 2U);
  BOOST_TEST(edgeC[0].in == std::vector<unsigned>({nodeV[0].id, nodeV[1].id}));

  BOOST_TEST(edgeC[1].out.size() == 1U);
  BOOST_TEST(edgeC[1].out == std::vector<unsigned>({nodeC[1].id}));
  BOOST_TEST(edgeC[1].in.size() == 2U);
  BOOST_TEST(edgeC[1].in == std::vector<unsigned>({nodeV[2].id, nodeV[3].id}));

  BOOST_TEST(edgeC[2].out.size() == 1U);
  BOOST_TEST(edgeC[2].out == std::vector<unsigned>({nodeC[2].id}));
  BOOST_TEST(edgeC[2].in.size() == 2U);
  BOOST_TEST(edgeC[2].in == std::vector<unsigned>({nodeV[4].id, nodeV[5].id}));

  BOOST_TEST(edgeC[3].out.size() == 1U);
  BOOST_TEST(edgeC[3].out == std::vector<unsigned>({nodeC[3].id}));
  BOOST_TEST(edgeC[3].in.size() == 2U);
  BOOST_TEST(edgeC[3].in == std::vector<unsigned>({nodeV[6].id, nodeV[7].id}));
}

/*
Testing Hypergraph for MatMulOuter nodes and edges - no reduction case

nodeA_   nodeB_   nodeC_
0 1    x 4 5    =  8  9
2 3      6 7      x  10

nodeV_
id  idxA_ idxB_
12  [0,1] [0,2]
13  [0,1] [1,3]
14  [2,3] [1,3]

*/
BOOST_AUTO_TEST_CASE(HyperGraph_testMatMulOuterNoReduction) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);

  const int rows = 2;
  const int cols = 2;
  const int rows_block = 1;
  const int cols_block = 1;
  const int block_rows = rows / rows_block;
  const int block_cols = cols / cols_block;

  BlockDenseMatrix A(rows, cols, rows_block, cols_block, false);

  BlockDenseMatrix B(rows, cols, rows_block, cols_block, false);

  std::vector<unsigned char> sparsityC(block_rows * block_cols, 1);
  sparsityC[2 + 0] = 0;

  // No reduction - 1 tile
  HyperGraphBlockTest hg(A, B, FLOAT, FLOAT, FLOAT, 1, 1);

  hg.createGraphMatMulSparsifyResult(graph, sparsityC.data(), "C");

  auto &nodeA = hg.getNodeA();
  auto &nodeB = hg.getNodeB();
  auto &nodeC = hg.getNodeC();
  auto &nodeV = hg.getNodeV();

  auto &edgeA = hg.getEdgeA();
  auto &edgeB = hg.getEdgeB();
  auto &edgeC = hg.getEdgeC();

  BOOST_TEST(nodeA.size() == 4U);
  BOOST_TEST(nodeB.size() == 4U);
  BOOST_TEST(nodeC.size() == 3U);
  BOOST_TEST(nodeV.size() == 3U);

  BOOST_TEST(edgeA.size() == 4U);
  BOOST_TEST(edgeB.size() == 4U);
  BOOST_TEST(edgeC.size() == 3U);

  BOOST_TEST(nodeV[0].idxA.size() == 2U);
  BOOST_TEST(nodeV[0].idxB.size() == 2U);
  BOOST_TEST(nodeV[0].idxA == std::vector<unsigned>({0, 1}));
  BOOST_TEST(nodeV[0].idxB == std::vector<unsigned>({0, 2}));

  BOOST_TEST(nodeV[1].idxA.size() == 2U);
  BOOST_TEST(nodeV[1].idxB.size() == 2U);
  BOOST_TEST(nodeV[1].idxA == std::vector<unsigned>({0, 1}));
  BOOST_TEST(nodeV[1].idxB == std::vector<unsigned>({1, 3}));

  BOOST_TEST(nodeV[2].idxA.size() == 2U);
  BOOST_TEST(nodeV[2].idxB.size() == 2U);
  BOOST_TEST(nodeV[2].idxA == std::vector<unsigned>({2, 3}));
  BOOST_TEST(nodeV[2].idxB == std::vector<unsigned>({1, 3}));

  BOOST_TEST(edgeA[0].in.size() == 1U);
  BOOST_TEST(edgeA[0].in == std::vector<unsigned>({nodeA[0].id}));
  BOOST_TEST(edgeA[0].out.size() == 2U);
  BOOST_TEST(edgeA[0].out == std::vector<unsigned>({nodeV[0].id, nodeV[1].id}));

  BOOST_TEST(edgeA[1].in.size() == 1U);
  BOOST_TEST(edgeA[1].in == std::vector<unsigned>({nodeA[1].id}));
  BOOST_TEST(edgeA[1].out.size() == 2U);
  BOOST_TEST(edgeA[1].out == std::vector<unsigned>({nodeV[0].id, nodeV[1].id}));

  BOOST_TEST(edgeA[2].in.size() == 1U);
  BOOST_TEST(edgeA[2].in == std::vector<unsigned>({nodeA[2].id}));
  BOOST_TEST(edgeA[2].out.size() == 1U);
  BOOST_TEST(edgeA[2].out == std::vector<unsigned>({nodeV[2].id}));

  BOOST_TEST(edgeA[3].in.size() == 1U);
  BOOST_TEST(edgeA[3].in == std::vector<unsigned>({nodeA[3].id}));
  BOOST_TEST(edgeA[3].out.size() == 1U);
  BOOST_TEST(edgeA[3].out == std::vector<unsigned>({nodeV[2].id}));

  BOOST_TEST(edgeB[0].in.size() == 1U);
  BOOST_TEST(edgeB[0].in == std::vector<unsigned>({nodeB[0].id}));
  BOOST_TEST(edgeB[0].out.size() == 1U);
  BOOST_TEST(edgeB[0].out == std::vector<unsigned>({nodeV[0].id}));

  BOOST_TEST(edgeB[1].in.size() == 1U);
  BOOST_TEST(edgeB[1].in == std::vector<unsigned>({nodeB[1].id}));
  BOOST_TEST(edgeB[1].out.size() == 2U);
  BOOST_TEST(edgeB[1].out == std::vector<unsigned>({nodeV[1].id, nodeV[2].id}));

  BOOST_TEST(edgeB[2].in.size() == 1U);
  BOOST_TEST(edgeB[2].in == std::vector<unsigned>({nodeB[2].id}));
  BOOST_TEST(edgeB[2].out.size() == 1U);
  BOOST_TEST(edgeB[2].out == std::vector<unsigned>({nodeV[0].id}));

  BOOST_TEST(edgeB[3].in.size() == 1U);
  BOOST_TEST(edgeB[3].in == std::vector<unsigned>({nodeB[3].id}));
  BOOST_TEST(edgeB[3].out.size() == 2U);
  BOOST_TEST(edgeB[3].out == std::vector<unsigned>({nodeV[1].id, nodeV[2].id}));

  BOOST_TEST(edgeC[0].out.size() == 1U);
  BOOST_TEST(edgeC[0].out == std::vector<unsigned>({nodeC[0].id}));
  BOOST_TEST(edgeC[0].in.size() == 1U);
  BOOST_TEST(edgeC[0].in == std::vector<unsigned>({nodeV[0].id}));

  BOOST_TEST(edgeC[1].out.size() == 1U);
  BOOST_TEST(edgeC[1].out == std::vector<unsigned>({nodeC[1].id}));
  BOOST_TEST(edgeC[1].in.size() == 1U);
  BOOST_TEST(edgeC[1].in == std::vector<unsigned>({nodeV[1].id}));

  BOOST_TEST(edgeC[2].out.size() == 1U);
  BOOST_TEST(edgeC[2].out == std::vector<unsigned>({nodeC[2].id}));
  BOOST_TEST(edgeC[2].in.size() == 1U);
  BOOST_TEST(edgeC[2].in == std::vector<unsigned>({nodeV[2].id}));
}

void TestMatMul(const poplar::Type &dataType, int blockSize, int batchBlockSize,
                int blockRowsA, int blockColsA, int blockColsB,
                int rowInBlockNzA = 0, int colInBlockNzA = 0,
                int blockRowNzA = -1, int blockColNzA = -1);

/*
Testing asymmetric layout of blocks and asymmetric layout within a block
*/
BOOST_AUTO_TEST_CASE(MatMul_testF32_asymmetric) {
  TestMatMul(FLOAT, 8, 6, 2, 2, 2, 1, 2, 0, 1);
}

/*
Testing batch block of 8
*/
BOOST_AUTO_TEST_CASE(MatMul_testF32_batchBlock8) {
  TestMatMul(FLOAT, 8, 8, 2, 2, 2, 1, 2, 0, 1);
}

/*
Testing all combinations or number of block rows and columns from 1 to 2
*/
BOOST_AUTO_TEST_CASE(MatMul_testF32) {
  for (int blockRowsA = 1; blockRowsA <= 2; ++blockRowsA) {
    for (int blockColsA = 1; blockColsA <= 2; ++blockColsA) {
      for (int blockColsB = 1; blockColsB <= 2; ++blockColsB) {
        TestMatMul(FLOAT, 8, 6, blockRowsA, blockColsA, blockColsB, 1, 2);
      }
    }
  }
}

/*
Testing fp16
*/
BOOST_AUTO_TEST_CASE(MatMul_testF16) { TestMatMul(HALF, 16, 6, 1, 2, 1); }

/*
Testing MatMul vertex - no reduction case
*/
void TestMatMul(const poplar::Type &dataType, int blockSize, int batchBlockSize,
                int blockRowsA, int blockColsA, int blockColsB,
                int rowInBlockNzA, int colInBlockNzA, int blockRowNzA,
                int blockColNzA) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);
  HyperGraphBlock::addCodelets(graph);

  const int rowsInBlockA = batchBlockSize;
  const int colsInBlockA = blockSize;
  const int rowsInBlockB = colsInBlockA;
  const int colsInBlockB = blockSize;
  const int rowsInBlockC = rowsInBlockA;
  const int colsInBlockC = colsInBlockB;

  const int blockRowsB = blockColsA;
  const int blockRowsC = blockRowsA;
  const int blockColsC = blockColsB;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int rowsB = rowsInBlockB * blockRowsB;
  const int colsB = colsInBlockB * blockColsB;

  unsigned outBlocks = blockRowsC * blockColsC;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;
  populateMatricesData1(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                        colsInBlockA, colsInBlockB, rowInBlockNzA,
                        colInBlockNzA, blockRowNzA, blockColNzA);

  std::vector<unsigned char> sparsityB(blockRowsB * blockColsB, 1);

  boost::multi_array<float, 2> hostBTransposed = transpose(hostB);

  // Blocks layout in blocksHostB are row-major
  // Individual blocks are column-major
  boost::multi_array<float, 2> blocksHostB;
  // Transpose B
  getSparseMatrixBlocks(colsB, rowsB, colsInBlockB, rowsInBlockB, sparsityB,
                        hostBTransposed, blocksHostB);

  BlockDenseMatrix A(rowsA, colsA, rowsInBlockA, colsInBlockA, false);

  BlockSparseMatrix B(colsB, rowsB, colsInBlockB, rowsInBlockB, true,
                      sparsityB.data());

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

  // LHS dense matrix
  poplar::Tensor tensorA = A.createTensor(graph, dataType, "A");
  A.setBlockTensor(tensorA);
  poputil::mapTensorLinearly(graph, tensorA);

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  // RHS sparse matrix
  poplar::Tensor tensorB = B.createTensor(graph, dataType, "B");
  B.setBlockTensor(tensorB);
  poputil::mapTensorLinearly(graph, tensorB);

  std::unique_ptr<char[]> blocksRawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, blocksHostB, dataType, blocksRawHostB.get());

  // No reduction
  HyperGraphBlockTest hg(A, B, dataType, dataType, dataType, 1, 1);
  hg.createGraphMatMul(graph, "C");

  // Put everything on tile 0
  std::vector<int> tileAssignment(hg.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts;
  poplar::program::Sequence matMulProg;
  hg.setTileAssignment(tileAssignment);
  hg.createComputeSetMatMul(graph, tensorCParts, matMulProg, "test");
  std::size_t tensorCPartsLen = tensorCParts.size();
  BOOST_TEST(tensorCPartsLen == outBlocks);

  std::vector<std::unique_ptr<char[]>> rawHostCParts;
  for (auto iter = tensorCParts.begin(); iter != tensorCParts.end(); ++iter) {
    rawHostCParts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        iter->second, std::string("C_") + std::to_string(iter->first), graph,
        uploadProg, downloadProg, streamMaps));
  }

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  // The result is row-major
  std::vector<boost::multi_array<float, 2>> blocksHostCParts;
  for (std::size_t i = 0; i < rawHostCParts.size(); ++i) {
    boost::multi_array<float, 2> blocksHostCPart(
        boost::extents[rowsInBlockC][colsInBlockC]);
    poplibs_test::util::copy(target, dataType, rawHostCParts[i].get(),
                             blocksHostCPart);
    blocksHostCParts.push_back(std::move(blocksHostCPart));
  }

  // hostC is row-major
  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  checkDenseResult(dataType, blockRowsC, blockColsC, rowsInBlockC, colsInBlockC,
                   hostC, blocksHostCParts);
}

void TestMatMulReduce(const poplar::Type &dataType, int blockSize,
                      int batchBlockSize);

BOOST_AUTO_TEST_CASE(MatMulReduce_testF32) { TestMatMulReduce(FLOAT, 8, 6); }

BOOST_AUTO_TEST_CASE(MatMulReduce_testF32batchBlock8) {
  TestMatMulReduce(FLOAT, 8, 8);
}

/*
Testing MatMul and Reduce vertices
*/
void TestMatMulReduce(const poplar::Type &dataType, int blockSize,
                      int batchSize) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);
  HyperGraphBlock::addCodelets(graph);

  const int blockRowsA = 1;
  const int blockColsA = 2;
  const int blockColsB = 1;

  const int rowsInBlockA = batchSize;
  const int colsInBlockA = blockSize;
  const int rowsInBlockB = colsInBlockA;
  const int colsInBlockB = blockSize;
  const int rowsInBlockC = rowsInBlockA;
  const int colsInBlockC = colsInBlockB;

  const int blockRowsB = blockColsA;
  const int blockRowsC = blockRowsA;
  const int blockColsC = blockColsB;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int rowsB = rowsInBlockB * blockRowsB;
  const int colsB = colsInBlockB * blockColsB;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;
  populateMatricesData1(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                        colsInBlockA, colsInBlockB, 1, 2);

  std::vector<unsigned char> sparsityB(blockRowsB * blockColsB, 1);

  boost::multi_array<float, 2> hostBTransposed = transpose(hostB);

  boost::multi_array<float, 2> blocksHostB;
  // Transpose B
  getSparseMatrixBlocks(colsB, rowsB, colsInBlockB, rowsInBlockB, sparsityB,
                        hostBTransposed, blocksHostB);

  BlockDenseMatrix A(rowsA, colsA, rowsInBlockA, colsInBlockA, false);

  BlockSparseMatrix B(colsB, rowsB, colsInBlockB, rowsInBlockB, true,
                      sparsityB.data());

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

  // LHS dense matrix
  poplar::Tensor tensorA = A.createTensor(graph, dataType, "A");
  A.setBlockTensor(tensorA);
  poputil::mapTensorLinearly(graph, tensorA);

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  // RHS sparse matrix
  poplar::Tensor tensorB = B.createTensor(graph, dataType, "B");
  B.setBlockTensor(tensorB);
  poputil::mapTensorLinearly(graph, tensorB);

  std::unique_ptr<char[]> blocksRawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, blocksHostB, dataType, blocksRawHostB.get());

  // Reduction
  HyperGraphBlockTest hg(A, B, dataType, dataType, dataType,
                         target.getTilesPerIPU());
  hg.createGraphMatMul(graph, "C");

  // Put everything on tile 0
  std::vector<int> tileAssignment(hg.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts;
  poplar::program::Sequence matMulProg;
  hg.setTileAssignment(tileAssignment);
  hg.createComputeSetMatMul(graph, tensorCParts, matMulProg, "test");
  std::size_t tensorCPartsLen = tensorCParts.size();
  BOOST_TEST(tensorCPartsLen == hg.getNodeV().size());

  poplar::program::Sequence reduceProg;
  hg.createComputeSetReduce(graph, tensorCParts, reduceProg, "test");

  std::vector<std::unique_ptr<char[]>> rawHostCParts;

  for (std::size_t i = 0; i < hg.matC->getBlockTensor().size(); ++i) {
    rawHostCParts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        hg.matC->getBlockTensor()[i], std::string("C_") + std::to_string(i),
        graph, uploadProg, downloadProg, streamMaps));
  }

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(reduceProg);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  // The result is row-major
  std::vector<boost::multi_array<float, 2>> blocksHostCParts;
  for (std::size_t i = 0; i < rawHostCParts.size(); ++i) {
    boost::multi_array<float, 2> blocksHostCPart(
        boost::extents[rowsInBlockC][colsInBlockC]);
    poplibs_test::util::copy(target, dataType, rawHostCParts[i].get(),
                             blocksHostCPart);
    blocksHostCParts.push_back(std::move(blocksHostCPart));
  }

  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  checkDenseResult(dataType, blockRowsC, blockColsC, rowsInBlockC, colsInBlockC,
                   hostC, blocksHostCParts);
}

void TestMatMulOuter(const poplar::Type &dataType, bool needTranspose);

BOOST_AUTO_TEST_CASE(MatMulOuter_testF32) { TestMatMulOuter(FLOAT, false); }

BOOST_AUTO_TEST_CASE(MatMulOuterTranspose_testF32) {
  TestMatMulOuter(FLOAT, true);
}

/*
Testing MatMulOuter vertex - no reduction case
*/
void TestMatMulOuter(const poplar::Type &dataType, bool needTranspose) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);
  HyperGraphBlock::addCodelets(graph);

  const int blockSize = 8;
  const int batchSize = 6;

  const int blockRowsA = 2;
  const int blockColsA = 2;
  const int blockColsB = 2;

  const int rowsInBlockA = batchSize;
  const int colsInBlockA = blockSize;
  const int rowsInBlockB = colsInBlockA;
  const int colsInBlockB = blockSize;
  const int rowsInBlockC = rowsInBlockA;
  const int colsInBlockC = colsInBlockB;

  const int blockRowsB = blockColsA;
  const int blockRowsC = blockRowsA;
  const int blockColsC = blockColsB;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int rowsB = rowsInBlockB * blockRowsB;
  const int colsB = colsInBlockB * blockColsB;

  unsigned outBlocks = blockRowsC * blockColsC;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;
  populateMatricesData1(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                        colsInBlockA, colsInBlockB, 1, 2, 0, 1);

  std::vector<unsigned char> sparsityC(blockRowsC * blockColsC, 1);

  // Mask [1][0] element out
  const int blockRowZero = 1;
  const int blockColZero = 0;
  sparsityC[blockRowZero * blockColsB + blockColZero] = 0;
  --outBlocks;

  BlockDenseMatrix A(rowsA, colsA, rowsInBlockA, colsInBlockA, false);

  BlockDenseMatrix B(colsB, rowsB, colsInBlockB, rowsInBlockB, needTranspose);

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

  // LHS dense matrix
  poplar::Tensor tensorA = A.createTensor(graph, dataType, "A");
  A.setBlockTensor(tensorA);
  poputil::mapTensorLinearly(graph, tensorA);

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  // RHS dense matrix
  poplar::Tensor tensorB = B.createTensor(graph, dataType, "B");
  B.setBlockTensor(tensorB);
  poputil::mapTensorLinearly(graph, tensorB);

  std::unique_ptr<char[]> rawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  if (needTranspose) {
    boost::multi_array<float, 2> hostBTransposed = transpose(hostB);
    poplibs_test::util::copy(target, hostBTransposed, dataType, rawHostB.get());
  } else {
    poplibs_test::util::copy(target, hostB, dataType, rawHostB.get());
  }

  // No reduction
  HyperGraphBlockTest hg(A, B, dataType, dataType, dataType, 1, 1);
  hg.createGraphMatMulSparsifyResult(graph, sparsityC.data(), "C");

  // Put everything on tile 0
  std::vector<int> tileAssignment(hg.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts;
  poplar::program::Sequence matMulProg;
  hg.setTileAssignment(tileAssignment);
  hg.createComputeSetMatMul(graph, tensorCParts, matMulProg, "test");
  std::size_t tensorCPartsLen = tensorCParts.size();
  BOOST_TEST(tensorCPartsLen == outBlocks);

  std::vector<std::unique_ptr<char[]>> rawHostCParts;
  for (auto iter = tensorCParts.begin(); iter != tensorCParts.end(); ++iter) {
    rawHostCParts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        iter->second, std::string("C_") + std::to_string(iter->first), graph,
        uploadProg, downloadProg, streamMaps));
  }

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  // The result is row-major
  std::vector<boost::multi_array<float, 2>> blocksHostCParts;
  for (int br = 0, outBlockIdx = 0; br < blockRowsC; ++br) {
    for (int bc = 0; bc < blockColsC; ++bc) {
      boost::multi_array<float, 2> blocksHostCPart(
          boost::extents[rowsInBlockC][colsInBlockC]);
      std::fill_n(blocksHostCPart.data(), blocksHostCPart.num_elements(), 0.0f);
      if (sparsityC[br * blockColsC + bc] > 0) {
        poplibs_test::util::copy(target, dataType,
                                 rawHostCParts[outBlockIdx].get(),
                                 blocksHostCPart);
        ++outBlockIdx;
      }
      blocksHostCParts.push_back(std::move(blocksHostCPart));
    }
  }

  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  checkSparseResult(dataType, blockRowsC, blockColsC, rowsInBlockC,
                    colsInBlockC, hostC, blocksHostCParts, sparsityC);
}

void TestSparseTensorReuse4Transpose(const poplar::Type &dataType,
                                     int blockSize, int blockRowsA,
                                     int blockColsA, int blockColsB,
                                     int rowInBlockNzA, int colInBlockNzA,
                                     int blockRowNzA, int blockColNzA);

BOOST_AUTO_TEST_CASE(SparseTensorReuse4Transpose_testF32) {
  TestSparseTensorReuse4Transpose(FLOAT, 8, 2, 3, 2, 1, 2, 0, 1);
}

/*
Testing that created by API sparse tensor can be used for both transposed and
non-transposed scenarios

  1) X * W = Y

  2) dY * Wt = dX

  Using notations:

  X == A, W == B, Y == C
  dX == A1, Wt == Bt, dY == C1

*/
void TestSparseTensorReuse4Transpose(const poplar::Type &dataType,
                                     int blockSize, int blockRowsA,
                                     int blockColsA, int blockColsB,
                                     int rowInBlockNzA, int colInBlockNzA,
                                     int blockRowNzA, int blockColNzA) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);
  HyperGraphBlock::addCodelets(graph);

  const int batchSize = 6;

  const int rowsInBlockA = batchSize;
  const int colsInBlockA = blockSize;
  const int rowsInBlockB = colsInBlockA;
  const int colsInBlockB = blockSize;
  const int rowsInBlockC = rowsInBlockA;
  const int colsInBlockC = colsInBlockB;

  ///////////////////
  // A * B = C
  const int blockRowsB = blockColsA;
  const int blockRowsC = blockRowsA;
  const int blockColsC = blockColsB;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int rowsB = rowsInBlockB * blockRowsB;
  const int colsB = colsInBlockB * blockColsB;

  unsigned outBlocks = blockRowsC * blockColsC;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;
  populateMatricesData1(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                        colsInBlockA, colsInBlockB, rowInBlockNzA,
                        colInBlockNzA, blockRowNzA, blockColNzA);

  std::vector<unsigned char> sparsityB(blockRowsB * blockColsB, 1);

  boost::multi_array<float, 2> blocksHostB;
  getSparseMatrixBlocks(rowsB, colsB, rowsInBlockB, colsInBlockB, sparsityB,
                        hostB, blocksHostB);

  BlockDenseMatrix A(rowsA, colsA, rowsInBlockA, colsInBlockA, false);

  BlockSparseMatrix B(rowsB, colsB, colsInBlockB, rowsInBlockB, false,
                      sparsityB.data());

  ///////////////////
  // A1 * Bt = C1
  const int blockRowsA1 = blockRowsA;
  const int blockColsA1 = blockRowsC;

  const int blockRowsC1 = blockRowsA1;
  const int blockColsC1 = blockRowsB;

  const int rowsInBlockA1 = rowsInBlockA;
  const int colsInBlockA1 = colsInBlockC;

  const int rowsA1 = rowsInBlockA1 * blockRowsA1;
  const int colsA1 = colsInBlockA1 * blockColsA1;

  unsigned outBlocks1 = blockRowsC1 * blockColsC1;

  boost::multi_array<float, 2> hostA1;
  populateMatrixData1(hostA1, rowsA1, colsA1, rowsInBlockA1, colsInBlockA1,
                      // Swap rows and columns to match transposed B
                      colInBlockNzA, rowInBlockNzA, blockColNzA, blockRowNzA);

  boost::multi_array<float, 2> hostBt = transpose(hostB);

  std::vector<unsigned char> sparsityBt(blockColsB * blockRowsB, 1);

  BlockDenseMatrix A1(rowsA1, colsA1, rowsInBlockA1, colsInBlockA1, false);

  // Bt[colsB x rowsB] - here rows and cols are swapped again, because of "true"
  // flag
  BlockSparseMatrix Bt(rowsB, colsB, rowsInBlockB, colsInBlockB, true,
                       sparsityBt.data());

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

  /////////////////// A * B = C
  poplar::Tensor tensorA = A.createTensor(graph, dataType, "A");
  A.setBlockTensor(tensorA);
  poputil::mapTensorLinearly(graph, tensorA);

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  poplar::Tensor tensorB = B.createTensor(graph, dataType, "B");
  B.setBlockTensor(tensorB);
  poputil::mapTensorLinearly(graph, tensorB);

  std::unique_ptr<char[]> blocksRawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, blocksHostB, dataType, blocksRawHostB.get());

  // No reduction
  HyperGraphBlockTest hg(A, B, dataType, dataType, dataType, 1, 1);
  hg.createGraphMatMul(graph, "C");

  std::vector<int> tileAssignment(hg.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts;
  poplar::program::Sequence matMulProg;
  hg.setTileAssignment(tileAssignment);
  hg.createComputeSetMatMul(graph, tensorCParts, matMulProg, "test");
  std::size_t tensorCPartsLen = tensorCParts.size();
  BOOST_TEST(tensorCPartsLen == outBlocks);

  std::vector<std::unique_ptr<char[]>> rawHostCParts;
  for (auto iter = tensorCParts.begin(); iter != tensorCParts.end(); ++iter) {
    rawHostCParts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        iter->second, std::string("C_") + std::to_string(iter->first), graph,
        uploadProg, downloadProg, streamMaps));
  }

  /////////////////// A1 * Bt = C1
  poplar::Tensor tensorA1 = A1.createTensor(graph, dataType, "A1");
  A1.setBlockTensor(tensorA1);
  poputil::mapTensorLinearly(graph, tensorA1);

  std::unique_ptr<char[]> rawHostA1 =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA1, "A1", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA1, dataType, rawHostA1.get());

#if 0 // Reference path for tensor Bt

  boost::multi_array<float, 2> blocksHostBt;
  // Here rows and cols are swapped again, because hostB is not tranposed
  getSparseMatrixBlocks(rowsB, colsB, rowsInBlockB, colsInBlockB, sparsityBt,
                        hostB, blocksHostBt);

  poplar::Tensor tensorBt = Bt.createTensor(graph, dataType, "Bt");
  Bt.setBlockTensor(tensorBt);
  poputil::mapTensorLinearly(graph, tensorBt);

  std::unique_ptr<char[]> blocksRawHostBt = poplibs_test::util::allocateHostMemoryForTensor(
        tensorBt, "Bt", graph, uploadProg,
        downloadProg, streamMaps);
  poplibs_test::util::copy(target, blocksHostBt, dataType,
                             blocksRawHostBt.get());

#else

  // Reusing existing tensor in transposed matrix
  Bt.setBlockTensor(tensorB);

#endif

  // No reduction
  HyperGraphBlockTest hg1(A1, Bt, dataType, dataType, dataType, 1, 1);
  hg1.createGraphMatMul(graph, "C1");

  std::vector<int> tileAssignment1(hg1.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts1;
  poplar::program::Sequence matMulProg1;
  hg1.setTileAssignment(tileAssignment1);
  hg1.createComputeSetMatMul(graph, tensorCParts1, matMulProg1, "matMul1");
  std::size_t tensorCPartsLen1 = tensorCParts1.size();
  BOOST_TEST(tensorCPartsLen1 == outBlocks1);

  std::vector<std::unique_ptr<char[]>> rawHostC1Parts;
  for (auto iter = tensorCParts1.begin(); iter != tensorCParts1.end(); ++iter) {
    rawHostC1Parts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        iter->second, std::string("C1_") + std::to_string(iter->first), graph,
        uploadProg, downloadProg, streamMaps));
  }
  ///////////////////

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(matMulProg1);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  /////////////////// C
  std::vector<boost::multi_array<float, 2>> blocksHostCParts;
  for (std::size_t i = 0; i < rawHostCParts.size(); ++i) {
    boost::multi_array<float, 2> blocksHostCPart(
        boost::extents[rowsInBlockC][colsInBlockC]);
    poplibs_test::util::copy(target, dataType, rawHostCParts[i].get(),
                             blocksHostCPart);
    blocksHostCParts.push_back(std::move(blocksHostCPart));
  }

  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  checkDenseResult(dataType, blockRowsC, blockColsC, rowsInBlockC, colsInBlockC,
                   hostC, blocksHostCParts);

  /////////////////// C1
  std::vector<boost::multi_array<float, 2>> blocksHostC1Parts;
  for (std::size_t i = 0; i < rawHostC1Parts.size(); ++i) {
    boost::multi_array<float, 2> blocksHostC1Part(
        boost::extents[rowsInBlockC][colsInBlockC]);
    poplibs_test::util::copy(target, dataType, rawHostC1Parts[i].get(),
                             blocksHostC1Part);
    blocksHostC1Parts.push_back(std::move(blocksHostC1Part));
  }

  std::vector<std::vector<float>> hostC1 = matMul(hostA1, hostBt);

  checkDenseResult(dataType, blockRowsC1, blockColsC1, rowsInBlockC,
                   colsInBlockC, hostC1, blocksHostC1Parts);
}

void TestDenseTensorReuse4Transpose(const poplar::Type &dataType, int blockSize,
                                    int blockRowsA, int blockColsA,
                                    int blockColsB, int rowInBlockNzA,
                                    int colInBlockNzA, int blockRowNzA,
                                    int blockColNzA);

BOOST_AUTO_TEST_CASE(DenseTensorReuse4Transpose_testF32) {
  TestDenseTensorReuse4Transpose(FLOAT, 8, 2, 3, 2, 1, 2, 0, 1);
}

/*
Testing that created by API dense tensor can be used for both transposed and
non-transposed scenarios

  1) X * W = Y

  2) Xt * dY = dW

  Using notations:

  X == A, W == B, Y == C
  Xt == At, dY == B1, Y == C1

*/
void TestDenseTensorReuse4Transpose(const poplar::Type &dataType, int blockSize,
                                    int blockRowsA, int blockColsA,
                                    int blockColsB, int rowInBlockNzA,
                                    int colInBlockNzA, int blockRowNzA,
                                    int blockColNzA) {
  auto device = createTestDeviceFullSize(TEST_TARGET);
  const auto &target = device.getTarget();
  Graph graph(target);
  HyperGraphBlock::addCodelets(graph);

  const int batchSize = 8;

  const int rowsInBlockA = batchSize;
  const int colsInBlockA = blockSize;
  const int rowsInBlockB = colsInBlockA;
  const int colsInBlockB = blockSize;
  const int rowsInBlockC = rowsInBlockA;
  const int colsInBlockC = colsInBlockB;

  /////////////////// A * B = C
  //
  const int blockRowsB = blockColsA;
  const int blockRowsC = blockRowsA;
  const int blockColsC = blockColsB;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int rowsB = rowsInBlockB * blockRowsB;
  const int colsB = colsInBlockB * blockColsB;

  unsigned outBlocks = blockRowsC * blockColsC;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;
  populateMatricesData1(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                        colsInBlockA, colsInBlockB, rowInBlockNzA,
                        colInBlockNzA, blockRowNzA, blockColNzA);

  std::vector<unsigned char> sparsityB(blockRowsB * blockColsB, 1);

  boost::multi_array<float, 2> blocksHostB;
  getSparseMatrixBlocks(rowsB, colsB, rowsInBlockB, colsInBlockB, sparsityB,
                        hostB, blocksHostB);

  BlockDenseMatrix A(rowsA, colsA, rowsInBlockA, colsInBlockA, false);

  BlockSparseMatrix B(rowsB, colsB, colsInBlockB, rowsInBlockB, false,
                      sparsityB.data());

  /////////////////// At * B1 = C1
  const int blockRowsB1 = blockRowsA;
  const int blockColsB1 = blockColsC;

  const int blockRowsC1 = blockColsA;
  const int blockColsC1 = blockColsB;

  const int rowsInBlockB1 = rowsInBlockA;
  const int colsInBlockB1 = colsInBlockC;

  const int rowsB1 = rowsInBlockB1 * blockRowsB1;
  const int colsB1 = colsInBlockB1 * blockColsB1;

  unsigned outBlocks1 = blockRowsC1 * blockColsC1;

  boost::multi_array<float, 2> hostB1;
  populateMatrixData1(hostB1, rowsB1, colsB1, rowsInBlockB1, colsInBlockB1,
                      // Swap rows and columns twice here
                      rowInBlockNzA, colInBlockNzA, blockRowNzA, blockColNzA);

  boost::multi_array<float, 2> hostAt = transpose(hostA);

  std::vector<unsigned char> sparsityC1(blockColsC1 * blockRowsC1, 1);

  // Mask [1][0] element out
  const int blockRowZero = 1;
  const int blockColZero = 0;
  sparsityC1[blockRowZero * blockColsB + blockColZero] = 0;
  --outBlocks1;

  // At[colsA x rowsA]
  BlockDenseMatrix At(colsA, rowsA, colsInBlockA, rowsInBlockA, false);

  BlockDenseMatrix B1(rowsB1, colsB1, rowsInBlockB1, colsInBlockB1, false);

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

  /////////////////// A * B = C
  poplar::Tensor tensorA = A.createTensor(graph, dataType, "A");
  A.setBlockTensor(tensorA);
  poputil::mapTensorLinearly(graph, tensorA);

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  poplar::Tensor tensorB = B.createTensor(graph, dataType, "B");
  B.setBlockTensor(tensorB);
  poputil::mapTensorLinearly(graph, tensorB);

  std::unique_ptr<char[]> blocksRawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, blocksHostB, dataType, blocksRawHostB.get());

  // No reduction
  HyperGraphBlockTest hg(A, B, dataType, dataType, dataType, 1, 1);
  hg.createGraphMatMul(graph, "C");

  std::vector<int> tileAssignment(hg.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts;
  poplar::program::Sequence matMulProg;
  hg.setTileAssignment(tileAssignment);
  hg.createComputeSetMatMul(graph, tensorCParts, matMulProg, "test");
  std::size_t tensorCPartsLen = tensorCParts.size();
  BOOST_TEST(tensorCPartsLen == outBlocks);

  std::vector<std::unique_ptr<char[]>> rawHostCParts;
  for (auto iter = tensorCParts.begin(); iter != tensorCParts.end(); ++iter) {
    rawHostCParts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        iter->second, std::string("C_") + std::to_string(iter->first), graph,
        uploadProg, downloadProg, streamMaps));
  }

  /////////////////// At * B1 = C1
  poplar::Tensor tensorB1 = B1.createTensor(graph, dataType, "B1");
  B1.setBlockTensor(tensorB1);
  poputil::mapTensorLinearly(graph, tensorB1);

  std::unique_ptr<char[]> rawHostB1 =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB1, "B1", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostB1, dataType, rawHostB1.get());

#if 0 // Reference path for tensor At

  poplar::Tensor tensorAt = At.createTensor(graph, dataType, "At");
  At.setBlockTensor(tensorAt);
  poputil::mapTensorLinearly(graph, tensorAt);

  std::unique_ptr<char[]> rawHostAt =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorAt, "At", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostAt, dataType, rawHostAt.get());

#else

  // Reusing existing tensor in transposed matrix
  poplar::Tensor tensorAt = tensorA.transpose();
  At.setBlockTensor(tensorAt);

#endif

  // No reduction
  HyperGraphBlockTest hg1(At, B1, dataType, dataType, dataType, 1, 1);
  hg1.createGraphMatMulSparsifyResult(graph, sparsityC1.data(), "C1");

  std::vector<int> tileAssignment1(hg1.getTotalNodes(), 0);

  std::map<unsigned int, poplar::Tensor> tensorCParts1;
  poplar::program::Sequence matMulProg1;
  hg1.setTileAssignment(tileAssignment1);
  hg1.createComputeSetMatMul(graph, tensorCParts1, matMulProg1, "matMul1");
  std::size_t tensorCPartsLen1 = tensorCParts1.size();
  BOOST_TEST(tensorCPartsLen1 == outBlocks1);

  std::vector<std::unique_ptr<char[]>> rawHostC1Parts;
  for (auto iter = tensorCParts1.begin(); iter != tensorCParts1.end(); ++iter) {
    rawHostC1Parts.push_back(poplibs_test::util::allocateHostMemoryForTensor(
        iter->second, std::string("C1_") + std::to_string(iter->first), graph,
        uploadProg, downloadProg, streamMaps));
  }
  ///////////////////

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(matMulProg1);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  /////////////////// C
  std::vector<boost::multi_array<float, 2>> blocksHostCParts;
  for (std::size_t i = 0; i < rawHostCParts.size(); ++i) {
    boost::multi_array<float, 2> blocksHostCPart(
        boost::extents[rowsInBlockC][colsInBlockC]);
    poplibs_test::util::copy(target, dataType, rawHostCParts[i].get(),
                             blocksHostCPart);
    blocksHostCParts.push_back(std::move(blocksHostCPart));
  }

  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  checkDenseResult(dataType, blockRowsC, blockColsC, rowsInBlockC, colsInBlockC,
                   hostC, blocksHostCParts);

  /////////////////// C1
  std::vector<boost::multi_array<float, 2>> blocksHostC1Parts;
  for (int br = 0, outBlockIdx = 0; br < blockRowsC1; ++br) {
    for (int bc = 0; bc < blockColsC1; ++bc) {
      boost::multi_array<float, 2> blocksHostC1Part(
          boost::extents[rowsInBlockC][colsInBlockC]);
      std::fill_n(blocksHostC1Part.data(), blocksHostC1Part.num_elements(),
                  0.0f);
      if (sparsityC1[br * blockColsC1 + bc] > 0) {
        poplibs_test::util::copy(target, dataType,
                                 rawHostC1Parts[outBlockIdx].get(),
                                 blocksHostC1Part);
        ++outBlockIdx;
      }
      blocksHostC1Parts.push_back(std::move(blocksHostC1Part));
    }
  }

  std::vector<std::vector<float>> hostC1 = matMul(hostAt, hostB1);

  checkSparseResult(dataType, blockRowsC1, blockColsC1, rowsInBlockC,
                    colsInBlockC, hostC1, blocksHostC1Parts, sparsityC1);
}

/*
Testing block-sparse API
dense x sparse = dense case
*/
void TestDSDAPI(const poplar::Type &dataType, int blockSize, int batchSize,
                const std::string &partitionMethod = "block-naive") {
  auto device = createTestDevice(TEST_TARGET, 1, 16);
  const auto &target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);
  poplin::addCodelets(graph);

  const int blockRowsA = 2;
  const int blockColsA = 3;
  const int blockColsB = 2;

  const int rowsInBlockA = batchSize;
  const int colsInBlockA = blockSize;
  const int rowsInBlockB = colsInBlockA;
  const int colsInBlockB = blockSize;

  const int blockRowsB = blockColsA;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int rowsB = rowsInBlockB * blockRowsB;
  const int colsB = colsInBlockB * blockColsB;
  const int rowsC = rowsA;
  const int colsC = colsB;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;

  std::vector<unsigned char> sparsityB(blockRowsB * blockColsB, 1);
  // Put couple of zero blocks
  sparsityB[2] = 0;
  sparsityB[4] = 0;

  populateMatricesDataWSparsity(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                                colsInBlockA, colsInBlockB, sparsityB.data());

  boost::multi_array<float, 2> blocksHostB;
  getSparseMatrixBlocks(rowsB, colsB, rowsInBlockB, colsInBlockB, sparsityB,
                        hostB, blocksHostB);

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

  BSMatMulParams bsParams({rowsA, colsA, colsB},
                          {rowsInBlockA, colsInBlockA, colsInBlockB}, sparsityB,
                          false, dataType, dataType, dataType);

  // LHS dense matrix
  poplar::Tensor tensorA = createBSMatMulInputLHS(graph, bsParams, "A");
  // RHS sparse matrix
  poplar::Tensor tensorB = createBSMatMulInputRHS(graph, bsParams, "B");

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  std::unique_ptr<char[]> blocksRawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, blocksHostB, dataType, blocksRawHostB.get());

  poplar::program::Sequence matMulProg;

  poplar::OptionFlags options = {
      {"partitionMethod", partitionMethod},
  };
  poplar::Tensor tensorC =
      bsMatMul(graph, bsParams, matMulProg, tensorA, tensorB, options);

  std::unique_ptr<char[]> rawHostC =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorC, "C", graph, uploadProg, downloadProg, streamMaps);

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  boost::multi_array<float, 2> blocksHostC(boost::extents[rowsC][colsC]);
  poplibs_test::util::copy(target, dataType, rawHostC.get(), blocksHostC);

  checkDenseResult(dataType, rowsC, colsC, hostC, blocksHostC);
}

/*
Testing block-sparse API
dense x dense = sparse case
*/
void TestDDSAPI(const poplar::Type &dataType, int blockSize, int batchSize,
                const std::string &partitionMethod = "block-naive") {
  auto device = createTestDevice(TEST_TARGET, 1, 16);
  const auto &target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);
  poplin::addCodelets(graph);

  const int blockRowsA = 2;
  const int blockColsA = 3;
  const int blockColsB = 4;

  const int rowsInBlockA = batchSize;
  const int colsInBlockA = blockSize;
  const int colsInBlockB = blockSize;

  const int rowsA = rowsInBlockA * blockRowsA;
  const int colsA = colsInBlockA * blockColsA;
  const int colsB = colsInBlockB * blockColsB;

  const int blockRowsC = blockRowsA;
  const int blockColsC = blockColsB;
  const int rowsInBlockC = rowsInBlockA;
  const int colsInBlockC = colsInBlockB;

  boost::multi_array<float, 2> hostA;
  boost::multi_array<float, 2> hostB;

  std::vector<unsigned char> sparsityC(blockRowsC * blockColsC, 1);
  // Put couple of zero blocks
  sparsityC[0] = 0;
  sparsityC[3] = 0;

  populateMatricesData1(hostA, hostB, rowsA, colsA, colsB, rowsInBlockA,
                        colsInBlockA, colsInBlockB);

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

  BSMatMulParams bsParams({rowsA, colsA, colsB},
                          {rowsInBlockA, colsInBlockA, colsInBlockB}, sparsityC,
                          dataType, dataType, dataType);

  // LHS dense matrix
  poplar::Tensor tensorA = createBSMatMulInputLHS(graph, bsParams, "A");
  // RHS dense matrix
  poplar::Tensor tensorB = createBSMatMulInputRHS(graph, bsParams, "B");

  std::unique_ptr<char[]> rawHostA =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorA, "A", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostA, dataType, rawHostA.get());

  std::unique_ptr<char[]> rawHostB =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorB, "B", graph, uploadProg, downloadProg, streamMaps);
  poplibs_test::util::copy(target, hostB, dataType, rawHostB.get());

  poplar::program::Sequence matMulProg;

  poplar::OptionFlags options = {
      {"partitionMethod", partitionMethod},
  };

  poplar::Tensor tensorC =
      bsMatMul(graph, bsParams, matMulProg, tensorA, tensorB, options);

  std::unique_ptr<char[]> rawHostC =
      poplibs_test::util::allocateHostMemoryForTensor(
          tensorC, "C", graph, uploadProg, downloadProg, streamMaps);

  Sequence allSequence;
  allSequence.add(uploadProg);
  allSequence.add(matMulProg);
  allSequence.add(downloadProg);

  const OptionFlags engineOptions{{"debug.allowOutOfMemory", "true"}};

  Engine engine(graph, allSequence, engineOptions);
  poplibs_test::util::attachStreams(engine, streamMaps);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run(0);
  });

  std::vector<std::vector<float>> hostC = matMul(hostA, hostB);

  boost::multi_array<float, 2> blocksHostC(
      boost::extents[tensorC.dim(0)][tensorC.dim(1)]);
  poplibs_test::util::copy(target, dataType, rawHostC.get(), blocksHostC);

  checkSparseResult(dataType, blockRowsC, blockColsC, rowsInBlockC,
                    colsInBlockC, hostC, blocksHostC, sparsityC);
}

BOOST_AUTO_TEST_CASE(DenseSparseDenseAPI_testF32_block) {
  TestDSDAPI(FLOAT, 8, 8, "block");
}
BOOST_AUTO_TEST_CASE(DenseSparseDenseAPI_testF32_block_naive) {
  TestDSDAPI(FLOAT, 8, 8, "block-naive");
}

BOOST_AUTO_TEST_CASE(DenseDenseSparseAPI_testF32) { TestDDSAPI(FLOAT, 8, 8); }


/*
 * python_mod.cpp
 *
 *  Created on: Jul 17, 2018
 *      Author: sujiwo
 */

#include <boost/python.hpp>
#include <iostream>
#include <Eigen/Eigen>
#include "VMap.h"
#include "MapPoint.h"
#include "KeyFrame.h"
#include "Frame.h"
#include "Localizer.h"
#include "MapBuilder2.h"
#include "pymat.h"

using namespace boost::python;
using namespace std;
using namespace Eigen;

static NDArrayConverter matcvt;

template<typename _Tp, int _rows>
void copyEigenVector2cv (const Eigen::Matrix<_Tp,_rows,1> &v, cv::Mat &m)
{
	assert(m.type() == cv::DataType<_Tp>::type);
	cv::Mat __(v.rows(), 1, cv::DataType<_Tp>::type, (void*)v.data(), v.stride()*sizeof(_Tp));
	__.copyTo(m);
}

Vector3d toVector3(const boost::python::list &p)
{
	return Vector3d (extract<double>(p[0]), extract<double>(p[1]), extract<double>(p[2]));
}

//Vector3d toVector3(const PyArrayObject *v)
//{
//	Vector3d vr;
//	vr.x() = extract<double>(v[0]);
//}

Quaterniond toQuaternion(const boost::python::list &qs)
{
	Quaterniond q;
	q.x() = extract<double>(qs[0]);
	q.y() = extract<double>(qs[1]);
	q.z() = extract<double>(qs[2]);
	q.w() = extract<double>(qs[3]);
	return q;
}

class VMapPy : public VMap
{
public:

	kfid find (PyObject *inp)
	{
		cv::Mat minp = matcvt.toMat(inp);
		kfid k = localizer->detect(minp);
		return k;
	}

	bool setCameraParams (const string &filename)
	{
		try {
			CameraPinholeParams par
				= CameraPinholeParams::loadCameraParamsFromFile(filename);
			if (localizer)
				localizer->setCameraParameter(par);
			return true;
		} catch (exception &e) {
			return false;
		}
	}

	dict info () const
	{
		dict mapinfo;
		mapinfo["numKeyframes"] = this->numOfKeyFrames();
		mapinfo["numMappoints"] = this->numOfMapPoints();
		return mapinfo;
	}

	dict camera() const
	{
		const CameraPinholeParams &cm = localizer->getCamera();
		dict camera;
		camera["fx"] = cm.fx;
		camera["fy"] = cm.fy;
		camera["cx"] = cm.cx;
		camera["cy"] = cm.cy;
		return camera;
	}

	PyObject*
	allMapPoints ()
	{
		uint N = this->numOfMapPoints();
		cv::Mat allPts(N, 3, CV_64F);

		for (mpid i: this->getMapPointList()) {
			Vector3d p = this->mappoint(i)->getPosition();
//			cv::Mat r = allPts.row(i);
//			copyEigenVector2cv(p, r);
			allPts.row(i).at<double>(0) = p.x();
			allPts.row(i).at<double>(1) = p.y();
			allPts.row(i).at<double>(2) = p.z();
		}

		return matcvt.toNDArray(allPts);
	}

	PyObject*
	allKeyFrames ()
	{
		uint N = this->numOfKeyFrames();
		cv::Mat allKfs(N, 7, CV_64F);

		for (kfid k: this->getKeyFrameList()) {
			Vector3d p = this->keyframe(k)->position();
			Quaterniond q = this->keyframe(k)->orientation();
			allKfs.row(k).col(0) = p.x();
			allKfs.row(k).col(1) = p.y();
			allKfs.row(k).col(2) = p.z();
			allKfs.row(k).col(3) = q.x();
			allKfs.row(k).col(4) = q.y();
			allKfs.row(k).col(5) = q.z();
			allKfs.row(k).col(6) = q.w();
		}

		return matcvt.toNDArray(allKfs);
	}

	bool load (const std::string &p)
	{
		try {
			VMap::load (p);
			localizer = new Localizer(this);
			return true;
		} catch(exception &e) {
			cerr << e.what() << endl;
			return false;
		}
	}

protected:
	Localizer *localizer;
};

class InputFramePy : public InputFrame
{
public:
	InputFramePy() {}

	InputFramePy(PyObject *_img, boost::python::list &posl, boost::python::list &ql)
	{
		position = toVector3(posl);
		orientation = toQuaternion(ql);
		image = matcvt.toMat(_img);
	}

	boost::python::list
	getpos()
	{
		boost::python::list l;
		l.append(position[0]);
		l.append(position[1]);
		l.append(position[2]);
		return l;
	}
};

class MapBuilderPy : public MapBuilder2
{
public:
	void initialize (const InputFramePy &f1, const InputFramePy &f2)
	{
		MapBuilder2::initialize(f1, f2);
	}

//	void initialize (PyObject *f1, PyObject *f2)
//	{
//
//	}

	void track (const InputFrame &f)
	{

	}

	VMap* getMap()
	{ return cMap; }

	void addCameraParam (const CameraPinholeParams &c)
	{ cMap->addCameraParameter(c); }

};

void
test_read (PyObject *arr)
{
	cv::Mat x = matcvt.toMat(arr);
	cerr << x << endl;
}

PyObject*
test_eye()
{
	cv::Mat M = cv::Mat::eye(3,3,CV_64F);
	cerr << M << endl;
	return matcvt.toNDArray(M);
}

BOOST_PYTHON_MODULE(photogrampy)
{
	Py_Initialize();
	import_array();
//	import_ufunc();

	class_ <VMapPy> ("VMap")
		.def("load", &VMapPy::load)
		.add_property("info", &VMapPy::info)
		.add_property("camera", &VMapPy::camera)
		.def("allMapPoints", &VMapPy::allMapPoints)
		.def("allKeyFrames", &VMapPy::allKeyFrames)
		.def("setCameraParams", &VMapPy::setCameraParams)
		.def("find", &VMapPy::find);

	class_ <InputFramePy> ("InputFrame",
		init<PyObject*,boost::python::list&,boost::python::list&>())
		.def("getpos", &InputFramePy::getpos);

	class_ <MapBuilderPy> ("MapBuilder")
		.def("initialize", &MapBuilderPy::initialize);

	def("test_read", &test_read);

	def("test_eye", &test_eye);
}


// Copyright (c) 2016 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "amount.h"
#include "test/test_bitcoin.h"

#include <boost/test/unit_test.hpp>

#include <array>

BOOST_FIXTURE_TEST_SUITE(feerate_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(GetFeeTest) {
    CFeeRate feeRate;

    feeRate = CFeeRate(Amount(0));
    // Must always return 0
    BOOST_CHECK_EQUAL(feeRate.GetFee(0), Amount(0));
    BOOST_CHECK_EQUAL(feeRate.GetFee(1e5), Amount(0));

    feeRate = CFeeRate(Amount(1000));
    // Must always just return the arg
    BOOST_CHECK_EQUAL(feeRate.GetFee(0), Amount(0));
    BOOST_CHECK_EQUAL(feeRate.GetFee(1), Amount(1));
    BOOST_CHECK_EQUAL(feeRate.GetFee(121), Amount(121));
    BOOST_CHECK_EQUAL(feeRate.GetFee(999), Amount(999));
    BOOST_CHECK_EQUAL(feeRate.GetFee(1000), Amount(1000));
    BOOST_CHECK_EQUAL(feeRate.GetFee(9000), Amount(9000));

    feeRate = CFeeRate(Amount(-1000));
    // Must always just return -1 * arg
    BOOST_CHECK_EQUAL(feeRate.GetFee(0), Amount(0));
    BOOST_CHECK_EQUAL(feeRate.GetFee(1), Amount(-1));
    BOOST_CHECK_EQUAL(feeRate.GetFee(121), Amount(-121));
    BOOST_CHECK_EQUAL(feeRate.GetFee(999), Amount(-999));
    BOOST_CHECK_EQUAL(feeRate.GetFee(1000), Amount(-1000));
    BOOST_CHECK_EQUAL(feeRate.GetFee(9000), Amount(-9000));

    feeRate = CFeeRate(Amount(123));
    // Truncates the result, if not integer
    BOOST_CHECK_EQUAL(feeRate.GetFee(0), Amount(0));
    // Special case: returns 1 instead of 0
    BOOST_CHECK_EQUAL(feeRate.GetFee(8), Amount(1));
    BOOST_CHECK_EQUAL(feeRate.GetFee(9), Amount(1));
    BOOST_CHECK_EQUAL(feeRate.GetFee(121), Amount(14));
    BOOST_CHECK_EQUAL(feeRate.GetFee(122), Amount(15));
    BOOST_CHECK_EQUAL(feeRate.GetFee(999), Amount(122));
    BOOST_CHECK_EQUAL(feeRate.GetFee(1000), Amount(123));
    BOOST_CHECK_EQUAL(feeRate.GetFee(9000), Amount(1107));

    feeRate = CFeeRate(Amount(-123));
    // Truncates the result, if not integer
    BOOST_CHECK_EQUAL(feeRate.GetFee(0), Amount(0));
    // Special case: returns -1 instead of 0
    BOOST_CHECK_EQUAL(feeRate.GetFee(8), Amount(-1));
    BOOST_CHECK_EQUAL(feeRate.GetFee(9), Amount(-1));

    // Check ceiling results
    feeRate = CFeeRate(Amount(18));
    // Truncates the result, if not integer
    BOOST_CHECK_EQUAL(feeRate.GetFeeCeiling(0), Amount(0));
    BOOST_CHECK_EQUAL(feeRate.GetFeeCeiling(100), Amount(2));
    BOOST_CHECK_EQUAL(feeRate.GetFeeCeiling(200), Amount(4));
    BOOST_CHECK_EQUAL(feeRate.GetFeeCeiling(1000), Amount(18));

    // Check full constructor
    // default value
    BOOST_CHECK(CFeeRate(Amount(-1), 1000) == CFeeRate(Amount(-1)));
    BOOST_CHECK(CFeeRate(Amount(0), 1000) == CFeeRate(Amount(0)));
    BOOST_CHECK(CFeeRate(Amount(1), 1000) == CFeeRate(Amount(1)));
    // lost precision (can only resolve satoshis per kB)
    BOOST_CHECK(CFeeRate(Amount(1), 1001) == CFeeRate(Amount(0)));
    BOOST_CHECK(CFeeRate(Amount(2), 1001) == CFeeRate(Amount(1)));
    // some more integer checks
    BOOST_CHECK(CFeeRate(Amount(26), 789) == CFeeRate(Amount(32)));
    BOOST_CHECK(CFeeRate(Amount(27), 789) == CFeeRate(Amount(34)));
    // Maximum size in bytes, should not crash
    CFeeRate(MAX_MONEY, std::numeric_limits<size_t>::max() >> 1).GetFeePerK();
}

BOOST_AUTO_TEST_SUITE_END()


// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#include "TransformationEstimation.h"

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/LU>
#include <Eigen/SVD>
#include <Eigen/Dense>
#include <Open3D/Geometry/PointCloud.h>
#include <Open3D/Utility/Eigen.h>
#include <iostream>

namespace Eigen {

#ifndef EIGEN_PARSED_BY_DOXYGEN

    // These helpers are required since it allows to use mixed types as parameters
    // for the Umeyama. The problem with mixed parameters is that the return type
    // cannot trivially be deduced when float and double types are mixed.
    namespace internal {

        // Compile time return type deduction for different MatrixBase types.
        // Different means here different alignment and parameters but the same underlying
        // real scalar type.
        template<typename MatrixType, typename OtherMatrixType>
            struct umeyama_transform_matrix_type2
            {
                enum {
                    MinRowsAtCompileTime = EIGEN_SIZE_MIN_PREFER_DYNAMIC(MatrixType::RowsAtCompileTime, OtherMatrixType::RowsAtCompileTime),

                    // When possible we want to choose some small fixed size value since the result
                    // is likely to fit on the stack. So here, EIGEN_SIZE_MIN_PREFER_DYNAMIC is not what we want.
                    HomogeneousDimension = int(MinRowsAtCompileTime) == Dynamic ? Dynamic : int(MinRowsAtCompileTime)+1
                };

                typedef Matrix<typename traits<MatrixType>::Scalar,
                        HomogeneousDimension,
                        HomogeneousDimension,
                        AutoAlign | (traits<MatrixType>::Flags & RowMajorBit ? RowMajor : ColMajor),
                        HomogeneousDimension,
                        HomogeneousDimension
                            > type;
            };

    }

#endif

    // Similar to the original Eigen::Umeyama  https://eigen.tuxfamily.org/dox/Umeyama_8h_source.html
    template <typename Derived, typename OtherDerived>
    Eigen::Matrix4d umeyama_constrained(const Eigen::MatrixXd& _src, const Eigen::MatrixXd& _dst, bool with_scaling = true) {
        typedef typename Eigen::internal::umeyama_transform_matrix_type2<Derived, OtherDerived>::type TransformationMatrixType;
        typedef typename Eigen::internal::traits<TransformationMatrixType>::Scalar Scalar;
        typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;

        enum { Dimension = EIGEN_SIZE_MIN_PREFER_DYNAMIC(Derived::RowsAtCompileTime, OtherDerived::RowsAtCompileTime) };

        typedef Eigen::Matrix<Scalar, Dimension, 1> VectorType;
        typedef Eigen::Matrix<Scalar, Dimension, Dimension> MatrixType;
        typedef typename Eigen::internal::plain_matrix_type_row_major<Derived>::type RowMajorMatrixType;

        const Index m = _src.rows(); // dimension
        const Index n = _src.cols(); // number of measurements

        Eigen::MatrixXd src(m, n);
        src << _src.topRows(2), Eigen::MatrixXd::Zero(1, n);
        Eigen::MatrixXd dst(m, n);
        dst << _dst.topRows(2), Eigen::MatrixXd::Zero(1, n);

        // required for demeaning ...
        const RealScalar one_over_n = RealScalar(1) / static_cast<RealScalar>(n);

        // computation of mean
        const VectorType src_mean = src.rowwise().sum() * one_over_n;
        const VectorType dst_mean = dst.rowwise().sum() * one_over_n;

        // demeaning of src and dst points
        const RowMajorMatrixType src_demean = src.colwise() - src_mean;
        const RowMajorMatrixType dst_demean = dst.colwise() - dst_mean;

        // Eq. (36)-(37)
        const Scalar src_var = src_demean.rowwise().squaredNorm().sum() * one_over_n;

        // Eq. (38)
        const MatrixType sigma = one_over_n * dst_demean * src_demean.transpose();

        Eigen::JacobiSVD<MatrixType> svd(sigma, Eigen::ComputeFullU | Eigen::ComputeFullV);

        // Initialize the resulting transformation with an identity matrix...
        TransformationMatrixType Rt = TransformationMatrixType::Identity(m+1,m+1);

        // Eq. (39)
        VectorType S = VectorType::Ones(m);

        if  ( svd.matrixU().determinant() * svd.matrixV().determinant() < 0 ) {
            // S(m-1) = -1;
            Rt.col(m).head(m) = dst_mean-src_mean;
            return Rt;
        }

        // Eq. (40) and (43)
        Rt.block(0,0,m,m).noalias() = svd.matrixU() * S.asDiagonal() * svd.matrixV().transpose();

        if (with_scaling)
        {
            // Eq. (42)
            const Scalar c = Scalar(1)/src_var * svd.singularValues().dot(S);

            // Eq. (41)
            Rt.col(m).head(m) = dst_mean;
            Rt.col(m).head(m).noalias() -= c*Rt.topLeftCorner(m,m)*src_mean;
            Rt.block(0,0,m,m) *= c;
        }
        else
        {
            Rt.col(m).head(m) = dst_mean;
            Rt.col(m).head(m).noalias() -= Rt.topLeftCorner(m,m)*src_mean;
        }
        return Rt;
    }
}

namespace open3d {
namespace registration {

double TransformationEstimationPointToPoint::ComputeRMSE(
        const geometry::PointCloud &source,
        const geometry::PointCloud &target,
        const CorrespondenceSet &corres) const {
    if (corres.empty()) return 0.0;
    double err = 0.0;
    for (const auto &c : corres) {
        err += (source.points_[c[0]] - target.points_[c[1]]).squaredNorm();
    }
    return std::sqrt(err / (double)corres.size());
}

Eigen::Matrix4d TransformationEstimationPointToPoint::ComputeTransformation(
        const geometry::PointCloud &source,
        const geometry::PointCloud &target,
        const CorrespondenceSet &corres) const {
    if (corres.empty()) return Eigen::Matrix4d::Identity();
    Eigen::MatrixXd source_mat(3, corres.size());
    Eigen::MatrixXd target_mat(3, corres.size());
    for (size_t i = 0; i < corres.size(); i++) {
        source_mat.block<3, 1>(0, i) = source.points_[corres[i][0]];
        target_mat.block<3, 1>(0, i) = target.points_[corres[i][1]];
    }
    if (with_constraint_) {
        return Eigen::umeyama_constrained<Eigen::MatrixXd, Eigen::MatrixXd>(source_mat, target_mat, with_scaling_);
    }
    else {
        return Eigen::umeyama(source_mat, target_mat, with_scaling_);
    }
}

double TransformationEstimationPointToPlane::ComputeRMSE(
        const geometry::PointCloud &source,
        const geometry::PointCloud &target,
        const CorrespondenceSet &corres) const {
    if (corres.empty() || target.HasNormals() == false) return 0.0;
    double err = 0.0, r;
    for (const auto &c : corres) {
        r = (source.points_[c[0]] - target.points_[c[1]])
                    .dot(target.normals_[c[1]]);
        err += r * r;
    }
    return std::sqrt(err / (double)corres.size());
}

Eigen::Matrix4d TransformationEstimationPointToPlane::ComputeTransformation(
        const geometry::PointCloud &source,
        const geometry::PointCloud &target,
        const CorrespondenceSet &corres) const {
    if (corres.empty() || target.HasNormals() == false)
        return Eigen::Matrix4d::Identity();

    auto compute_jacobian_and_residual = [&](int i, Eigen::Vector6d &J_r,
                                             double &r) {
        const Eigen::Vector3d &vs = source.points_[corres[i][0]];
        const Eigen::Vector3d &vt = target.points_[corres[i][1]];
        const Eigen::Vector3d &nt = target.normals_[corres[i][1]];
        r = (vs - vt).dot(nt);
        J_r.block<3, 1>(0, 0) = vs.cross(nt);
        J_r.block<3, 1>(3, 0) = nt;
    };

    Eigen::Matrix6d JTJ;
    Eigen::Vector6d JTr;
    double r2;
    std::tie(JTJ, JTr, r2) =
            utility::ComputeJTJandJTr<Eigen::Matrix6d, Eigen::Vector6d>(
                    compute_jacobian_and_residual, (int)corres.size());

    bool is_success;
    Eigen::Matrix4d extrinsic;
    std::tie(is_success, extrinsic) =
            utility::SolveJacobianSystemAndObtainExtrinsicMatrix(JTJ, JTr);

    return is_success ? extrinsic : Eigen::Matrix4d::Identity();
}

}  // namespace registration
}  // namespace open3d


/*!@file
 * @copyright This code is licensed under the 3-clause BSD license.
 *   Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.
 *   See LICENSE.txt for details.
 */

#include <boost/test/unit_test.hpp>

#include "Molassembler/Temple/BoundedNodeTrie.h"
#include "Molassembler/Temple/OrderedPair.h"
#include "Molassembler/Temple/Poset.h"
#include "Molassembler/Temple/Random.h"
#include "Molassembler/Temple/Stringify.h"
#include "Molassembler/Temple/constexpr/Jsf.h"

#include <algorithm>
#include <random>

using namespace Scine::Molassembler;

extern Temple::Generator<> generator;

BOOST_AUTO_TEST_CASE(OrderedPairTests, *boost::unit_test::label("Temple")) {
  Temple::OrderedPair<unsigned> a {14, 3};
  BOOST_CHECK(a.front() < a.back());
  BOOST_CHECK(a.first < a.second);
  BOOST_CHECK(
    std::is_sorted(std::begin(a), std::end(a))
  );
}

BOOST_AUTO_TEST_CASE(PosetTests, *boost::unit_test::label("Temple")) {
  constexpr unsigned nTests = 10;
  constexpr unsigned nValues = 20;
  constexpr unsigned maxValue = 40;

  auto compareFirstDigit = [](const unsigned a, const unsigned b) -> bool {
    return (a / 10) < (b / 10);
  };

  std::vector<unsigned> unorderedValues;
  for(unsigned testNum = 0; testNum < nTests; ++testNum) {
    unorderedValues = Temple::Random::getN<unsigned>(0, maxValue, nValues, generator.engine);
    Temple::Poset<unsigned> f {unorderedValues};

    f.orderUnordered(compareFirstDigit);
    f.orderUnordered(std::less<>());
    f.finalize();

    BOOST_CHECK(
      std::is_sorted(
        std::begin(f),
        std::end(f),
        [](const auto& a, const auto& b) -> bool {
          return std::less<>()(
            a.values.front(),
            b.values.front()
          );
        }
      )
    );
  }
}

template<
  typename ChoiceIndex,
  typename Generator
>
struct ChooseFunctor {
  Generator& prngRef;

  ChooseFunctor(Generator& prng) : prngRef(prng) {}

  int euclideanModulo(const int a, const int base) {
    return ((a % base) + base) % base;
  }

  int distance(const int i, const int j, const int U) {
    assert(i >= 0 && j >= 0 && U > 1);

    return std::min(
      euclideanModulo(i - j, U),
      euclideanModulo(j - i, U)
    );
  }

  double merit(const ChoiceIndex choice, const boost::dynamic_bitset<>& childExists) {
    double value = 0.0;

    const ChoiceIndex U = childExists.size();

    for(ChoiceIndex other = 0; other < U; ++other) {
      if(other == choice) {
        continue;
      }

      if(childExists.test(other)) {
        int choiceDistance = distance(choice, other, U);
        assert(choiceDistance >= 0);
        assert(choiceDistance < std::numeric_limits<ChoiceIndex>::max());
        value += choiceDistance;
      }
    }

    assert(value >= 0);

    return value;
  }

  ChoiceIndex operator () (const std::vector<ChoiceIndex>& viableChildren, const boost::dynamic_bitset<>& existingChildren) {
    assert(!viableChildren.empty());

    /* In case not all children already exist, it makes sense to calculate
     * merits for all choices at this level.
     */
    if(!existingChildren.all()) {
      std::vector<ChoiceIndex> bestChoices;
      double bestMerit = 0;
      for(ChoiceIndex viableChild : viableChildren) {
        if(!existingChildren.test(viableChild)) {
          double choiceMerit = merit(viableChild, existingChildren);
          if(choiceMerit > bestMerit) {
            bestChoices = {viableChild};
            bestMerit = choiceMerit;
          } else if(choiceMerit == bestMerit) {
            bestChoices.push_back(viableChild);
          }
        }
      }

      assert(!bestChoices.empty());

      return bestChoices.at(
        std::uniform_int_distribution<ChoiceIndex>(
          0,
          bestChoices.size() - 1
        )(prngRef.engine)
      );
    }

    /* All children exist, so no point in calculating merits: we choose from
     * the viable ones
     */
    return viableChildren.at(
      std::uniform_int_distribution<ChoiceIndex>(0, viableChildren.size() - 1)(prngRef.engine)
    );
  }
};

template<typename ChoiceIndex, typename Generator>
auto make_ChooseFunctor(
  const Temple::BoundedNodeTrie<ChoiceIndex>& /* trie */,
  Generator& prng
) {
  return ChooseFunctor<ChoiceIndex, Generator>(prng);
}

BOOST_AUTO_TEST_CASE(TrieTests, *boost::unit_test::label("Temple")) {
  {
    std::vector<std::uint8_t> boundaries {4, 2, 3};

    using TrieType = Temple::BoundedNodeTrie<std::uint8_t>;

    TrieType trie {boundaries};

    std::vector<TrieType::ChoiceList> sampleValueLists {
      {0, 0, 0},
      {0, 1, 2},
      {3, 0, 1}
    };

    BOOST_CHECK(trie.capacity() == 24);

    for(auto& sampleValueList : sampleValueLists) {
      BOOST_CHECK(!trie.contains(sampleValueList));
    }

    for(unsigned i = 0; i < sampleValueLists.size(); ++i) {
      trie.insert(sampleValueLists.at(i));

      for(unsigned j = 0; j <= i; ++j) {
        BOOST_CHECK(trie.contains(sampleValueLists.at(j)));
      }
      for(unsigned j = i + 1; j < sampleValueLists.size(); ++j) {
        BOOST_CHECK(!trie.contains(sampleValueLists.at(j)));
      }

      BOOST_CHECK(trie.size() == i + 1);
    }
  }
  {
    using TrieType = Temple::BoundedNodeTrie<std::uint8_t>;
    TrieType::ChoiceList boundaries;
    boundaries.push_back(4);

    TrieType trie {boundaries};

    std::vector<TrieType::ChoiceList> sampleValueLists {
      {0},
      {3},
      {2}
    };

    BOOST_CHECK(trie.capacity() == 4);

    for(auto& sampleValueList : sampleValueLists) {
      BOOST_CHECK(!trie.contains(sampleValueList));
    }

    for(unsigned i = 0; i < sampleValueLists.size(); ++i) {
      trie.insert(sampleValueLists.at(i));

      for(unsigned j = 0; j <= i; ++j) {
        BOOST_CHECK(trie.contains(sampleValueLists.at(j)));
      }
      for(unsigned j = i + 1; j < sampleValueLists.size(); ++j) {
        BOOST_CHECK(!trie.contains(sampleValueLists.at(j)));
      }

      BOOST_CHECK(trie.size() == i + 1);
    }
  }
  {
    // Generate new list every time test
    using TrieType = Temple::BoundedNodeTrie<std::uint8_t>;
    TrieType::ChoiceList boundaries {6, 4, 2};
    TrieType trie {boundaries};
    auto chooseFunctor = make_ChooseFunctor(trie, generator);

    std::set<TrieType::ChoiceList> listsSet;

    TrieType::ChoiceList list;
    while(trie.size() != trie.capacity()) {
      list = trie.generateNewEntry(chooseFunctor);

      BOOST_REQUIRE(listsSet.count(list) == 0);
      listsSet.insert(list);
    }
  }
}


/**
 ** Copyright (c) 2014 Illumina, Inc.
 **
 ** This file is part of Illumina's Enhanced Artificial Genome Engine (EAGLE),
 ** covered by the "BSD 2-Clause License" (see accompanying LICENSE file)
 **
 ** \author Lilian Janin
 **/

#ifndef EAGLE_GENOME_ERROR_MODEL_PLUGIN_HH
#define EAGLE_GENOME_ERROR_MODEL_PLUGIN_HH

#include <fstream>
#include <boost/assign.hpp>
#include <boost/format.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/discrete_distribution.hpp>

#include "model/Nucleotides.hh"
#include "common/Logger.hh"

namespace eagle
{
namespace genome
{

class ClusterErrorModelContext;

class ErrorModelPlugin
{
public:
    ErrorModelPlugin( const std::vector< std::string >& errorModelOptions, const std::string& errorModelId )
        : parsedUserOptions_( parseUserOptions( errorModelOptions, errorModelId ) )
    {
    }
    virtual ~ErrorModelPlugin()
    {
    }
    virtual void apply( boost::mt19937& randomGen, const double errorRate, unsigned int& randomErrorType, char& bclBase, ClusterErrorModelContext& clusterErrorModelContext ) = 0;

protected:
    template <class T>
    T getParsedValue( const std::string& key, const T& defaultValue )
    {
        std::map< std::string, std::pair< std::string, bool > >::iterator valueItr = parsedUserOptions_.find( key );
        if (valueItr != parsedUserOptions_.end())
        {
            T val;
            try
            {
                val = boost::lexical_cast<T>( valueItr->second.first );
            }
            catch (boost::bad_lexical_cast &)
            {
                EAGLE_ERROR( (boost::format("Invalid type for error model command line option \"%s\"") % key).str() );
            }
            valueItr->second.second = true; // Remember that this key+value has been used
            return val;
        }
        return defaultValue;
    }

    void reportUnusedCommandLineOptionsForThisPlugin();

private:
    std::map< std::string, std::pair< std::string, bool > > parseUserOptions( const std::vector< std::string >& errorModelOptions, const std::string& errorModelId );
    std::map< std::string, std::pair< std::string, bool > > parsedUserOptions_;
};

class LongreadBaseDuplicationModel: public ErrorModelPlugin
{
public:
    LongreadBaseDuplicationModel( const std::vector< std::string >& errorModelOptions );
    virtual void apply( boost::mt19937& randomGen, const double errorRate, unsigned int& randomErrorType, char& bclBase, ClusterErrorModelContext& clusterErrorModelContext );

private:
    const double prob_;
};

class LongreadDeletionModel: public ErrorModelPlugin
{
public:
    LongreadDeletionModel( const std::vector< std::string >& errorModelOptions );
    void apply( boost::mt19937& randomGen, const double errorRate, unsigned int& randomErrorType, char& bclBase, ClusterErrorModelContext& clusterErrorModelContext );

private:
    const double prob_;
    boost::random::discrete_distribution<> deletionLengthDist_;
    unsigned int basesLeftToDelete_;
};

} // namespace genome
} // namespace eagle

#endif // EAGLE_GENOME_ERROR_MODEL_PLUGIN_HH


//==============================================================================
//         Copyright 2003 - 2011   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_TRIGONOMETRIC_CONSTANTS_PIO_4LO_HPP_INCLUDED
#define NT2_TRIGONOMETRIC_CONSTANTS_PIO_4LO_HPP_INCLUDED

#include <nt2/include/functor.hpp>
#include <boost/simd/constant/hierarchy.hpp>
#include <boost/simd/constant/register.hpp>

namespace nt2
{
  namespace tag
  {
   /*!
     @brief Pio_4lo generic tag

     Represents the Pio_4lo constant in generic contexts.

     @par Models:
        Hierarchy
   **/
    BOOST_SIMD_CONSTANT_REGISTER( Pio_4lo, double
                                , 0, 0xb2bbbd2eUL
                                , 0x3c81a62633145c07ULL
                                )
  }
  namespace ext
  {
   template<class Site, class... Ts>
   BOOST_FORCEINLINE generic_dispatcher<tag::Pio_4lo, Site> dispatching_Pio_4lo(adl_helper, boost::dispatch::meta::unknown_<Site>, boost::dispatch::meta::unknown_<Ts>...)
   {
     return generic_dispatcher<tag::Pio_4lo, Site>();
   }
   template<class... Args>
   struct impl_Pio_4lo;
  }
  /*!
    This constant is such that, for pairs of types (Tup, T)
    (namely (float,  double) and (double, long double)) the sum:

    abs(Tup(Pio_4lo<T>())+Tup(Pio_4<T>())-Pio_4<Tup>())  is  lesser than
    a few Eps<Tup>().

    This is used to improve accurracy when computing sums of the kind
    Pio_4 + x with x small,  by replacing them by Pio_4 + (Pio_4lo+x)

    @par Semantic:

    For type T0:

    @code
    T r = Pio_4lo<T>();
    @endcode

    is similar to:

    @code
    if T is double
      r = 3.061616997868383e-17
    else if T is float
      r = -2.1855694e-08
    @endcode

    @return a value of type T0
  **/
  BOOST_SIMD_CONSTANT_IMPLEMENTATION(tag::Pio_4lo, Pio_4lo);
}

#endif