#include <Windows.h>
#include <experimental/filesystem>

#include <assert.h>
#include <math.h>
#include <time.h>

#include <algorithm>
#include <fstream>
#include <functional>
#include <iostream>
#include <numeric>
#include <sstream>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

#include <boost/proto/proto.hpp>
#include <boost/typeof/std/ostream.hpp>

template <int I>
struct placeholder
{ };

boost::proto::terminal<placeholder<0>>::type const _1 = { { } };
boost::proto::terminal<placeholder<1>>::type const _2 = { { } };

boost::proto::terminal<std::ostream&>::type cout_ = { std::cout };

template <typename Expr>
void evaluate(const Expr& expr) {
	boost::proto::default_context context;
	boost::proto::eval(expr, context);
}

struct calculator_context
	: boost::proto::callable_context<calculator_context const> {
	// Values to replace the placeholders
	std::vector<double> args;

	// Define the result type of the calculator.
	// (This makes the calculator_context "callable")
	using result_type = double;

	// Handle the placeholders
	template <int I>
	double operator()(boost::proto::tag::terminal, placeholder<I>) const {
		return this->args[I];
	}
};

int main(int argc, char *argv[])
{
	calculator_context ctx;
	ctx.args.emplace_back(45); // The value of _1 is 45
	ctx.args.emplace_back(50); // The value of _2 is 50

	// Create an arithmetic expression and immediately evaluate it
	const auto d = boost::proto::eval((_2 - _1) / _2 * 100, ctx);

	// This prints "10"
	std::cout << "d: " << d << "\n";

	return EXIT_SUCCESS;
}


// Copyright Louis Dionne 2013-2016
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)

#include <boost/hana/assert.hpp>
#include <boost/hana/equal.hpp>
#include <boost/hana/insert.hpp>
#include <boost/hana/integral_constant.hpp>
#include <boost/hana/set.hpp>
#include <boost/hana/string.hpp>
#include <boost/hana/type.hpp>
namespace hana = boost::hana;

int main() {
    constexpr auto xs = hana::make_set(hana::int_c<0>, hana::type_c<int>);
    BOOST_HANA_CONSTANT_CHECK(
        hana::insert(xs, BOOST_HANA_STRING("abc")) ==
        hana::make_set(hana::int_c<0>, hana::type_c<int>, BOOST_HANA_STRING("abc"))
    );

    BOOST_HANA_CONSTANT_CHECK(
        hana::insert(xs, hana::int_c<0>) == hana::make_set(hana::int_c<0>, hana::type_c<int>)
    );
}


#include <boost/ratio/mpl/rational_constant.hpp>


/**********************************************************************
*  Copyright (c) 2008-2014, Alliance for Sustainable Energy.
*  All rights reserved.
*
*  This library is free software; you can redistribute it and/or
*  modify it under the terms of the GNU Lesser General Public
*  License as published by the Free Software Foundation; either
*  version 2.1 of the License, or (at your option) any later version.
*
*  This library is distributed in the hope that it will be useful,
*  but WITHOUT ANY WARRANTY; without even the implied warranty of
*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
*  Lesser General Public License for more details.
*
*  You should have received a copy of the GNU Lesser General Public
*  License along with this library; if not, write to the Free Software
*  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
**********************************************************************/

#include "Matrix.hpp"

#include "../core/Optional.hpp"
#include "../math/FloatCompare.hpp"

#include <random>

// this should all be moved to a utilities/core/Random.h
#include <boost/random/uniform_real.hpp>
#include <boost/random/variate_generator.hpp>

namespace openstudio{

  bool operator==(const Matrix& lhs, const Matrix& rhs)
  {
    bool result = false;
    if (lhs.size1() == rhs.size1()){
      if (lhs.size2() == rhs.size2()){
        result = true;
        for (unsigned i = 0; i < lhs.size1(); ++i){
          for (unsigned j = 0; j < lhs.size2(); ++j){
            if (lhs(i,j) != rhs(i,j)){
              return false;
            }
          }
        }
      }
    }
    return result;
  }

  bool operator!=(const Matrix& lhs, const Matrix& rhs)
  {
    return !(lhs==rhs);
  }

  /// linear interpolation of the function v = f(x, y) at point xi, yi
  /// assumes that x and y are strictly increasing
  double interp(const Vector& x, const Vector& y, const Matrix& v, double xi, double yi, InterpMethod interpMethod, ExtrapMethod extrapMethod)
  {
    double result = 0.0;

    unsigned M = x.size();
    unsigned N = y.size();

    if ((M != v.size1()) || (N != v.size2())){
      return result;
    }

    InterpInfo xInfo = interpInfo(x, xi);

    if (xInfo.extrapolated){
      switch(extrapMethod){
        case NoneExtrap:
          // set all weights to zero
          xInfo.wa = 0.0; xInfo.wb = 0.0;
          break;
        case NearestExtrap:
          // pick closest point
          // no-op
          break;
      }
    }else{
      switch(interpMethod){
        case LinearInterp:
          // linear interpolation
          // no-op
          break;
        case NearestInterp:
          // pick closest point
          if(xInfo.wa > xInfo.wb){
            xInfo.wa = 1.0; xInfo.wb = 0.0;
          }else{
            xInfo.wa = 0.0; xInfo.wb = 1.0;
          }
          break;
        case HoldLastInterp:
          // set to previous value
          xInfo.wa = 1.0; xInfo.wb = 0.0;
          break;
        case HoldNextInterp:
          // set to next value
          xInfo.wa = 0.0; xInfo.wb = 1.0;
          break;
      }
    }

    InterpInfo yInfo = interpInfo(y, yi);

    if (yInfo.extrapolated){
      switch(extrapMethod){
        case NoneExtrap:
          // set all weights to zero
          yInfo.wa = 0.0; yInfo.wb = 0.0;
          break;
        case NearestExtrap:
          // pick closest point
          // no-op
          break;
      }
    }else{
      switch(interpMethod){
        case LinearInterp:
          // linear interpolation
          // no-op
          break;
        case NearestInterp:
          // pick closest point
          if(yInfo.wa > yInfo.wb){
            yInfo.wa = 1.0; yInfo.wb = 0.0;
          }else{
            yInfo.wa = 0.0; yInfo.wb = 1.0;
          }
          break;
        case HoldLastInterp:
          // set to previous value
          yInfo.wa = 1.0; yInfo.wb = 0.0;
          break;
        case HoldNextInterp:
          // set to next value
          yInfo.wa = 0.0; yInfo.wb = 1.0;
          break;
      }
    }

    // we have set weights appropriately so that here we can compute in the same way all the time
    result = xInfo.wa*yInfo.wa*v(xInfo.ia, yInfo.ia) +
             xInfo.wa*yInfo.wb*v(xInfo.ia, yInfo.ib) +
             xInfo.wb*yInfo.wa*v(xInfo.ib, yInfo.ia) +
             xInfo.wb*yInfo.wb*v(xInfo.ib, yInfo.ib);

    return result;
  }

  /// linear interpolation of the function v = f(x, y) at points xi, yi
  /// assumes that x and y are strictly increasing
  Vector interp(const Vector& x, const Vector& y, const Matrix& v, const Vector& xi, double yi, InterpMethod interpMethod, ExtrapMethod extrapMethod)
  {
    unsigned M = x.size();

    Vector result(M);

    if (M != v.size1()){
      return result;
    }

    for (unsigned i = 0; i < M; ++i){
      result(i) = interp(x, y, v, xi(i), yi, interpMethod, extrapMethod);
    }

    return result;
  }

  /// linear interpolation of the function v = f(x, y) at points xi, yi
  /// assumes that x and y are strictly increasing
  Vector interp(const Vector& x, const Vector& y, const Matrix& v, double xi, const Vector& yi, InterpMethod interpMethod, ExtrapMethod extrapMethod)
  {
    unsigned N = y.size();

    Vector result(N);

    if (N != v.size2()){
      return result;
    }

    for (unsigned j = 0; j < N; ++j){
      result(j) = interp(x, y, v, xi, yi(j), interpMethod, extrapMethod);
    }

    return result;
  }

  /// linear interpolation of the function v = f(x, y) at points xi, yi
  /// assumes that x and y are strictly increasing
  Matrix interp(const Vector& x, const Vector& y, const Matrix& v, const Vector& xi, const Vector& yi, InterpMethod interpMethod, ExtrapMethod extrapMethod)
  {
    unsigned M = x.size();
    unsigned N = y.size();

    Matrix result(M, N);

    if ((M != v.size1()) || (N != v.size2())){
      return result;
    }

    for (unsigned i = 0; i < M; ++i){
      for (unsigned j = 0; j < N; ++j){
        result(i,j) = interp(x, y, v, xi(i), yi(j), interpMethod, extrapMethod);
      }
    }

    return result;
  }

  /// matrix product
  Matrix prod(const Matrix& lop, const Matrix& rop)
  {
    return boost::numeric::ublas::prod(lop, rop);
  }

  /// vector product
  Vector prod(const Matrix& m, const Vector& v)
  {
    return boost::numeric::ublas::prod(m, v);
  }

  /// outer product
  Matrix outerProd(const Vector& lhs, const Vector& rhs)
  {
    return boost::numeric::ublas::outer_prod(lhs, rhs);
  }

  /// take the natural logarithm of a Matrix
  Matrix log(const Matrix& v)
  {
    unsigned M = v.size1();
    unsigned N = v.size2();
    Matrix result(M, N);
    for (unsigned i = 0; i < M; ++i){
      for (unsigned j = 0; j < N; ++j){
        result(i,j) = std::log(v(i,j));
      }
    }
    return result;
  }

  /// take the logarithm of a Matrix with certain base
  Matrix log(const Matrix& v, double base)
  {
    double logBase = std::log(base);
    unsigned M = v.size1();
    unsigned N = v.size2();
    Matrix result(M, N);
    for (unsigned i = 0; i < M; ++i){
      for (unsigned j = 0; j < N; ++j){
        result(i,j) = std::log(v(i,j)) / logBase;
      }
    }
    return result;
  }

  /// generates a Matrix of M*N points randomly drawn between and including a and b.
  Matrix randMatrix(double a, double b, unsigned M, unsigned N)
  {
    // ETH@20100120. What library does this come from? The user should be able to seed the
    // generator independently of this function.
    // seed random number generator
    static std::minstd_rand generator(42u);

    // define distribution
    boost::uniform_real<> dist(a,b);

    // create a generator
    boost::variate_generator<std::minstd_rand&, boost::uniform_real<> > uniformGenerator(generator, dist);

    // ETH@20120723 Started seeing this as DataFixture.Matrix_RandMatrix hanging on Windows 7,
    // with BoostPro installer.
    // handle degenerate case
    OptionalDouble singlePoint;
    if (equal(a,b)) {
      singlePoint = (a + b) / 2.0;
    }

    Matrix result(M, N);
    for (unsigned i = 0; i < M; ++i){
      for (unsigned j = 0; j < N; ++j){
        if (singlePoint) {
          result(i,j) = *singlePoint;
        }
        else {
          result(i,j) = uniformGenerator();
        }
      }
    }
    return result;
  }

  /// sum
  double sum(const Matrix& matrix)
  {
    double result = 0.0;
    for (unsigned i = 0; i < matrix.size1(); ++i){
      for (unsigned j = 0; j < matrix.size2(); ++j){
        result += matrix(i,j);
      }
    }
    return result;
  }

  /// maximum
  double maximum(const Matrix& matrix)
  {
    double max = 0;
    if ((matrix.size1() > 0) && (matrix.size2() > 0)){
      max = matrix(0,0);
      for (unsigned i = 0; i < matrix.size1(); ++i){
        for (unsigned j = 0; j < matrix.size2(); ++j){
          max = std::max(max, matrix(i,j));
        }
      }
    }
    return max;
  }

  /// minimum
  double minimum(const Matrix& matrix)
  {
    double min = 0;
    if ((matrix.size1() > 0) && (matrix.size2() > 0)){
      min = matrix(0,0);
      for (unsigned i = 0; i < matrix.size1(); ++i){
        for (unsigned j = 0; j < matrix.size2(); ++j){
          min = std::min(min, matrix(i,j));
        }
      }
    }
    return min;
  }

  /// mean
  double mean(const Matrix& matrix)
  {
    double avg = 0;
    unsigned N = matrix.size1()*matrix.size2();
    if (N > 0){
      avg = sum(matrix) / N;
    }
    return avg;
  }

  /// get the connected components from an NxN adjacency matrix (1.0 for i-j connected, 0.0 for i-j not connected)
  std::vector<std::vector<unsigned> > findConnectedComponents(const Matrix& matrix)
  {
    double tol = 0.001;

    std::vector<std::vector<unsigned> > result;

    unsigned N = matrix.size1();
    if (N != matrix.size2()){
      return result;
    }

    Matrix A(N, N, 0.0);
    for (unsigned i = 0; i < N; ++i){

      A(i,i) = 1.0; // must be self connected
      if ( std::abs(matrix(i,i) - 1.0) > tol){
        // warn
      }

      for (unsigned j = i+1; j < N; ++j){

        if (matrix(i,j) < 0){
          // warn
        }else if (matrix(i,j) > tol){
          A(i,j) = 1.0;
        }

        if (matrix(j,i) < 0){
          // warn
        }else if (matrix(j,i) > tol){
          A(j,i) = 1.0;
        }
      }
    }

    // raise A to the Nth power, maximum distance between two nodes
    for (unsigned i = 0; i < N; ++i){
      A = prod(A,A);
    }

    std::set<unsigned> added;
    for (unsigned i = 0; i < N; ++i){
      if (added.find(i) != added.end()){
        continue;
      }

      std::vector<unsigned> group;
      group.push_back(i);
      added.insert(i);

      for (unsigned j = i+1; j < N; ++j){
        if ((A(i,j) > 0) || (A(j,i) > 0)){
          group.push_back(j);
          added.insert(j);
        }
      }

      result.push_back(group);
    }

    return result;
  }

} // openstudio


#include <aslam/calibration/model/sensors/PositionSensor.h>

#include <memory>

#include <aslam/backend/TransformationExpression.hpp>
#include <glog/logging.h>
#include <boost/make_shared.hpp>

#include "aslam/calibration/calibrator/CalibratorI.h"
#include "aslam/calibration/data/PositionMeasurement.h"
#include "aslam/calibration/data/MeasurementsContainer.h"
#include "aslam/calibration/error-terms/ErrorTermPosition.h"
#include <aslam/calibration/tools/ErrorTermStatistics.h>
#include <aslam/calibration/tools/ErrorTermStatisticsWithProblemAndPredictor.h>
#include <aslam/calibration/model/Model.h>
#include <aslam/calibration/model/ModuleTools.h>
#include <aslam/calibration/model/Sensor.h>

namespace aslam {
namespace calibration {

//TODO B convert PositionSensor to external storage based

void PositionSensor::addMeasurement(const Timestamp t, const PositionMeasurement& position)
{
  measurements->emplace_back(t, position);
}

void PositionSensor::addMeasurement(const Timestamp t, const Eigen::Vector3d& p) {
  addMeasurement(t, PositionMeasurement{p});
}

void PositionSensor::addInputTo(Timestamp t, const PositionMeasurement& position, ModuleStorage& /*s*/) const {
  measurements->emplace_back(t, position);
}

void PositionSensor::writeConfig(std::ostream& out) const {
  Sensor::writeConfig(out);
  MODULE_WRITE_PARAM(covPosition);
}

PositionSensor::PositionSensor(Model& model, std::string name, sm::value_store::ValueStoreRef config) :
  Sensor(model, name, config),
  covPosition(getMyConfig().getChild("covPosition"), 3),
  targetFrame(getModel().getFrame(getMyConfig().getString("targetFrame")))
{
  if(isUsed()) {
    measurements = std::make_shared<PositionMeasurements>();
  }
}

PositionSensor::~PositionSensor() {
}

void PositionSensor::addMeasurementErrorTerms(CalibratorI& calib, const CalibrationConfI & /*ec*/, ErrorTermReceiver & problem, const bool observeOnly) const {
  const std::string errorTermGroupName = getName() + "Position";
  if(!measurements || measurements->empty()){
    LOG(WARNING) << "No measurements available for " << errorTermGroupName;
    return;
  }

  ErrorTermStatisticsWithProblemAndPredictor es(calib, errorTermGroupName, problem, observeOnly);

  ErrorTermGroupReference etgr(errorTermGroupName);

  auto interval = calib.getCurrentEffectiveBatchInterval();

  auto conditionalLowerBound = interval.start + getDelayUpperBound();
  auto conditionalUpperBound = interval.end + getDelayLowerBound();
  auto certainLowerBound = interval.start + getDelayLowerBound();
  auto certainUpperBound = interval.end + getDelayUpperBound();

  auto & delay = getDelayExpression();
  Timestamp currentDelay = delay.evaluate();
  if(currentDelay < getDelayLowerBound() || currentDelay > getDelayUpperBound()){
    throw std::runtime_error("Delay already out of bounds!");
  }

  for (auto & m : *measurements) {
    Timestamp timestamp = m.first;
    if(certainLowerBound > timestamp || certainUpperBound < timestamp){
      LOG(WARNING) << "Dropping out of bounds measurement for " << getName() << " at " << calib.secsSinceStart(timestamp) << "!";
      continue;
    }

    auto & positionMeasurement = m.second;

    aslam::backend::TransformationExpression T_target_s = getTransformationExpressionToAtMeasurementTimestamp(calib, timestamp, targetFrame, true);
    boost::shared_ptr<ErrorTermPosition> e_position(new ErrorTermPosition(T_target_s.inverse().toEuclideanExpression(), positionMeasurement, covPosition.getValue(), etgr));

    if(conditionalLowerBound > timestamp || conditionalUpperBound < timestamp){
      if(!hasDelay()){
        LOG(WARNING) << "Dropping out of bounds measurement for " << getName() << " at " << calib.secsSinceStart(timestamp) << "!";
        continue;
      }
      LOG(INFO) << "Adding conditional PositionErrorTerm for position measurement at " << calib.secsSinceStart(timestamp) << " because it could go out of bounds!";
      if(conditionalLowerBound > timestamp){
        e_position = addConditionShared<ErrorTermPosition>(*e_position, [=](){ return timestamp - delay.evaluate() >= interval.start; });
      }
      if(conditionalUpperBound < timestamp){
        e_position = addConditionShared<ErrorTermPosition>(*e_position, [=](){ return timestamp - delay.evaluate() <= interval.end; });
      }
    }

    es.add(timestamp, e_position, false);
  }
  es.printInto(LOG(INFO));
}

void PositionSensor::clearMeasurements() {
  measurements.reset();
}

} /* namespace calibration */
} /* namespace aslam */


#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/features/normal_3d.h>
#include <pcl/surface/gp3.h>
#include <pcl/io/vtk_io.h>
#include <pcl/io/ply_io.h>
#include <pcl/io/vtk_lib_io.h>
#include <stdio.h>
#include <pcl/visualization/cloud_viewer.h>
#include "tinyxml.h"
#include <Eigen/Dense>
#include <Eigen/Core>

using namespace Eigen;
using namespace std;
//#define SHOWORIENTATION_MEASURED
//#define SHOWORIENTATION_CARDSFLOW
//#define SHOW_EULER
//#define SHOW_SENSOR0_COMPONENTS
#define SHOWMAGNITUDE_SENSOR0
#define SHOWMAGNITUDE_SENSOR1
#define SHOWMAGNITUDE_SENSOR2
#define SHOWMAGNITUDE_SENSOR3
//#define SHOWSENSOR0
//#define SHOWSENSOR1
//#define SHOWSENSOR2
//#define SHOWSENSOR3
//#define GENERATESTL

boost::shared_ptr<pcl::visualization::PCLVisualizer> rgbVis (pcl::PointCloud<pcl::PointXYZRGB>::ConstPtr cloud)
{
    // --------------------------------------------
    // -----Open 3D viewer and add point cloud-----
    // --------------------------------------------
    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer (new pcl::visualization::PCLVisualizer ("3D Viewer"));
    viewer->setBackgroundColor (0, 0, 0);
    pcl::visualization::PointCloudColorHandlerRGBField<pcl::PointXYZRGB> rgb(cloud);
    viewer->addPointCloud<pcl::PointXYZRGB> (cloud, rgb, "sample cloud");
    viewer->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 1, "sample cloud");
    viewer->addCoordinateSystem (1.0);
    viewer->initCameraParameters ();
    return (viewer);
}

float sensor_location_scale = 5;
float sensor_scale = 2;

int main (int argc, char** argv)
{
    TiXmlDocument doc("/home/letrend/workspace/roboy3/src/robots/neck_magentic_field/cardsflow.xml");
    if (!doc.LoadFile()) {
        printf("Can't parse file\n");
        return false;
    }

    TiXmlElement *root = doc.RootElement();
    TiXmlElement *myoMuscle_it = NULL;
    vector<Vector3d> sensor_rel_locations;
    for (myoMuscle_it = root->FirstChildElement("myoMuscle"); myoMuscle_it;
         myoMuscle_it = myoMuscle_it->NextSiblingElement("myoMuscle")) {
        if (myoMuscle_it->Attribute("name")) {
            // myoMuscle joint acting on
            TiXmlElement *link_child_it = NULL;
            for (link_child_it = myoMuscle_it->FirstChildElement("link"); link_child_it;
                 link_child_it = link_child_it->NextSiblingElement("link")) {
                string link_name = link_child_it->Attribute("name");
                if (!link_name.empty()) {
                    TiXmlElement *viaPoint_child_it = NULL;
                    for (viaPoint_child_it = link_child_it->FirstChildElement("viaPoint"); viaPoint_child_it;
                         viaPoint_child_it = viaPoint_child_it->NextSiblingElement("viaPoint")) {
                        float x, y, z;
                        if (sscanf(viaPoint_child_it->GetText(), "%f %f %f", &x, &y, &z) != 3) {
                            printf("parser", "error reading [via point] (x y z)\n");
                            return false;
                        }
                        Vector3d local_coordinates(x, y, z);
                        sensor_rel_locations.push_back(local_coordinates);
                    }
                }
            }
        }
    }
    // Load input file into a PointCloud<T> with an appropriate type
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud (new pcl::PointCloud<pcl::PointXYZRGB>);
    FILE*       file = fopen("/home/letrend/workspace/roboy3/head_data0.log","r");

    if (NULL == file) {
        printf("Failed to open 'yourfile'");
        return -1;
    }
    fscanf(file, "%*[^\n]\n", NULL);
    float roll, pitch, yaw;
    float s[4][3];
    bool first = true;
//    Quaterniond quat_init;
    while(fscanf(file,"%f %f %f %f %f %f %f %f %f %f %f %f %f %f %f\n",
    &s[0][0],&s[0][1],&s[0][2],
    &s[1][0],&s[1][1],&s[1][2],
    &s[2][0],&s[2][1],&s[2][2],
    &s[3][0],&s[3][1],&s[3][2],&roll,&pitch,&yaw)==15) {
        Vector3d dir(0, 0, 1);
        Vector3d mag0(s[0][0], s[0][1], s[0][2]);
        Vector3d mag1(s[1][0], s[1][1], s[1][2]);
        Vector3d mag2(s[2][0], s[2][1], s[2][2]);
        Vector3d mag3(s[3][0], s[3][1], s[3][2]);

//        Quaterniond q(qw, qx, qy, qz);

//        Matrix3d rot = q.matrix();
        Matrix3d rot;
        rot = AngleAxisd(roll, Vector3d::UnitX())
        * AngleAxisd(pitch, Vector3d::UnitY())
        * AngleAxisd(yaw, Vector3d::UnitZ());

#ifdef SHOWORIENTATION_MEASURED
        dir = rot * dir;
//            mag = rot*mag*0.01;
//            for(int sensor=0;sensor<3;sensor++){
        pcl::PointXYZRGB p0;
        p0.x = dir[0];
        p0.y = dir[1];
        p0.z = dir[2];
        p0.r = 255;
        p0.g = 255;
        p0.b = 255;
        cloud->push_back(p0);
#endif
#ifdef SHOWORIENTATION_CARDSFLOW
        dir << 0,0,1;
        Quaterniond q2(q_top_w, q_top_x, q_top_y, q_top_z);
//        if(first){
//            quat_init = q2;
//            first = false;
//        }
//        q2 = q2*quat_init.inverse();
        Matrix3d rot2 = q2.matrix();
        dir = rot2 * dir;
        p0.x = dir[0];
        p0.y = dir[1];
        p0.z = dir[2];
        p0.r = 0;
        p0.g = 30;
        p0.b = 90;
        cloud->push_back(p0);
#endif
#ifdef SHOW_EULER
        if(abs(roll)<0.5 && abs(pitch)<0.5 && abs(yaw)<1.5) {
            pcl::PointXYZRGB p;
            p.x = roll;
            p.y = pitch;
            p.z = yaw;
            p.r = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOW_SENSOR0_COMPONENTS
        {
            dir << 0, 0, 1 + mag0[0] * sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 255;
            cloud->push_back(p);
        }
        {
            dir << 0, 0, 1 + mag0[1] * sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.g = 255;
            cloud->push_back(p);
        }
        {
            dir << 0, 0, 1 + mag0[2] * sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.b = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWMAGNITUDE_SENSOR0
        {
            double norm = mag0.norm();
            Vector3d dir = sensor_rel_locations[0]*norm*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWMAGNITUDE_SENSOR1
        {
            double norm = mag1.norm();
            Vector3d dir = sensor_rel_locations[1]*norm*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.g = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWMAGNITUDE_SENSOR2
        {
            double norm = mag2.norm();
            Vector3d dir = sensor_rel_locations[2]*norm*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.b = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWMAGNITUDE_SENSOR3
        {
            double norm = mag3.norm();
            Vector3d dir = sensor_rel_locations[3]*norm*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 235;
            p.g = 225;
            p.b = 52;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWSENSOR0
        {
            Vector3d dir = sensor_rel_locations[0]*sensor_location_scale + sensor_rel_locations[0]*mag0[0]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[0]*sensor_location_scale + sensor_rel_locations[0]*mag0[1]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.g = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[0]*sensor_location_scale + sensor_rel_locations[0]*mag0[2]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.b = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWSENSOR1
        {
            Vector3d dir = sensor_rel_locations[1]*sensor_location_scale + sensor_rel_locations[1]*mag1[0]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[1]*sensor_location_scale + sensor_rel_locations[1]*mag1[1]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.g = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[1]*sensor_location_scale + sensor_rel_locations[1]*mag1[2]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.b = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWSENSOR2
        {
            Vector3d dir = sensor_rel_locations[2]*sensor_location_scale + sensor_rel_locations[2]*mag2[0]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[2]*sensor_location_scale + sensor_rel_locations[2]*mag2[1]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.g = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[2]*sensor_location_scale + sensor_rel_locations[2]*mag2[2]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.b = 255;
            cloud->push_back(p);
        }
#endif
#ifdef SHOWSENSOR3
        {
            Vector3d dir = sensor_rel_locations[3]*sensor_location_scale + sensor_rel_locations[3]*mag3[0]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.r = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[3]*sensor_location_scale + sensor_rel_locations[3]*mag3[1]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.g = 255;
            cloud->push_back(p);
        }
        {
            Vector3d dir = sensor_rel_locations[3]*sensor_location_scale + sensor_rel_locations[3]*mag3[2]*sensor_scale;
            dir = rot * dir;
            pcl::PointXYZRGB p;
            p.x = dir[0];
            p.y = dir[1];
            p.z = dir[2];
            p.b = 255;
            cloud->push_back(p);
        }
#endif

    }

    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer = rgbVis(cloud);

    //use the following functions to get access to the underlying more advanced/powerful
    //PCLVisualizer

    while (!viewer->wasStopped ())
    {
        viewer->spinOnce (100);
        boost::this_thread::sleep (boost::posix_time::microseconds (100000));
    }
    //* the data should be available in cloud

#ifdef GENERATESTL
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_xyz(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::copyPointCloud(*cloud, *cloud_xyz);

    // Normal estimation*
    pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> n;
    pcl::PointCloud<pcl::Normal>::Ptr normals (new pcl::PointCloud<pcl::Normal>);
    pcl::search::KdTree<pcl::PointXYZ>::Ptr tree (new pcl::search::KdTree<pcl::PointXYZ>);
    tree->setInputCloud (cloud_xyz);
    n.setInputCloud (cloud_xyz);
    n.setSearchMethod (tree);
    n.setKSearch (20);
    n.compute (*normals);
    //* normals should not contain the point normals + surface curvatures

    // Concatenate the XYZ and normal fields*
    pcl::PointCloud<pcl::PointNormal>::Ptr cloud_with_normals (new pcl::PointCloud<pcl::PointNormal>);
    pcl::concatenateFields (*cloud_xyz, *normals, *cloud_with_normals);
    //* cloud_with_normals = cloud + normals

    // Create search tree*
    pcl::search::KdTree<pcl::PointNormal>::Ptr tree2 (new pcl::search::KdTree<pcl::PointNormal>);
    tree2->setInputCloud (cloud_with_normals);

    // Initialize objects
    pcl::GreedyProjectionTriangulation<pcl::PointNormal> gp3;
    pcl::PolygonMesh triangles;

    // Set the maximum distance between connected points (maximum edge length)
    gp3.setSearchRadius (0.5);

    // Set typical values for the parameters
    gp3.setMu (2.5);
    gp3.setMaximumNearestNeighbors (10000);
    gp3.setMaximumSurfaceAngle(M_PI/4); // 45 degrees
    gp3.setMinimumAngle(M_PI/18); // 10 degrees
    gp3.setMaximumAngle(2*M_PI/3); // 120 degrees
    gp3.setNormalConsistency(true);

    // Get result
    gp3.setInputCloud (cloud_with_normals);
    gp3.setSearchMethod (tree2);
    gp3.reconstruct (triangles);

    // Additional vertex information
    std::vector<int> parts = gp3.getPartIDs();
    std::vector<int> states = gp3.getPointStates();

    pcl::io::savePolygonFileSTL ("/home/letrend/workspace/roboy_control/src/ball_in_socket_estimator/mesh.stl", triangles);
#endif
    // Finish
    return (0);
}


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

#include <boost/program_options.hpp>
#include <chrono>
#include <cmath>
#include <ctime>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <sys/resource.h>

#define DEVELOP_BLOCKSPARSE

#include "poplibs_support/logging.hpp"
#include "poplibs_test/Util.hpp"
#include <poplar/IPUModel.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplin/MatMul.hpp>
#include <poplin/codelets.hpp>
#include <popnn/NonLinearity.hpp>
#include <popops/codelets.hpp>
#include <popsparse/experimental/BlockSparse.hpp>
#include <popsparse/experimental/BlockSparseMatMul.hpp>
#include <poputil/TileMapping.hpp>
#include <poputil/exceptions.hpp>

#include "poplibs_test/Pass.hpp"
#include "poplibs_test/Util.hpp"

using namespace poplar;
using namespace poplin;
using namespace poplar::program;
using namespace poplibs_test::util;
using namespace popsparse::experimental;
using namespace poplibs_support;
using namespace poplibs_test;

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

bool ReadMatrixMask(std::string &fileName, int &nonZeroBlock, int &row,
                    int &col, std::vector<unsigned char> &mask,
                    bool readTransposed) {
  std::ifstream cin(fileName);
  if (!cin.is_open()) {
    std::cerr << " Can not open matrix sparsity mask file " << fileName
              << std::endl;
    return false;
  }

  cin >> row >> col;
  if (row <= 0 || col <= 0) {
    std::cerr << " Incorrect matrix sparsity mask file format " << fileName
              << std::endl;
    return false;
  }
  mask.resize(row * col);
  unsigned char *data = mask.data();

  if (!data) {
    return false;
  }

  for (int i = 0; i < row; i++) {
    for (int j = 0; j < col; j++) {
      int c;
      cin >> c;
      if (!readTransposed) {
        data[i * col + j] = c ? 1 : 0;
      } else {
        data[j * row + i] = c ? 1 : 0;
      }
    }
  }

  nonZeroBlock = 0;
  for (int i = 0; i < row * col; i++) {
    if (data[i] == 1) {
      nonZeroBlock++;
    }
  }
  return true;
}

void populateMatrixData(const std::array<int, 3> &dim,
                        const std::array<int, 3> &blockSize,
                        boost::multi_array<float, 2> &lhsMatrix,
                        boost::multi_array<float, 2> &rhsMatrix,
                        const unsigned char *sparsity, bool rhsNeedTranspose) {

  std::mt19937 randomEngine;
  randomEngine.seed(102302);
  for (int i = 0; i < dim[0]; i++) {
    for (int j = 0; j < dim[1]; j++) {
      lhsMatrix[i][j] = static_cast<float>(randomEngine()) / INT_MAX;
    }
  }

  int blockRow;
  int blockCol;
  int nRowBlock;
  int nColBlock;
  if (!rhsNeedTranspose) {
    blockRow = blockSize[1];
    blockCol = blockSize[2];
    nRowBlock = dim[1] / blockSize[1];
    nColBlock = dim[2] / blockSize[2];
  } else {
    blockRow = blockSize[2];
    blockCol = blockSize[1];
    nRowBlock = dim[2] / blockSize[2];
    nColBlock = dim[1] / blockSize[1];
  }

  for (int i = 0; i < nRowBlock; i++) {
    for (int j = 0; j < nColBlock; j++) {
      int blockRowStart = i * blockRow;
      int blockColStart = j * blockCol;
      int blockRowEnd = blockRowStart + blockRow;
      int blockColEnd = blockColStart + blockCol;

      if (!sparsity || sparsity[i * nColBlock + j] == 1) {
        for (int r = blockRowStart; r < blockRowEnd; r++) {
          for (int c = blockColStart; c < blockColEnd; c++) {
            rhsMatrix[r][c] = static_cast<float>(randomEngine()) /
                              static_cast<float>(INT_MAX);
          }
        }
      }
    }
  }
}

void populateMatrixData(const std::array<int, 3> &dim,
                        const std::array<int, 3> &blockSize, int numGroups,
                        boost::multi_array<float, 3> &lhsMatrix,
                        boost::multi_array<float, 3> &rhsMatrix,
                        const unsigned char *sparsity, bool rhsNeedTranspose) {

  std::mt19937 randomEngine;
  randomEngine.seed(102302);
  for (int g = 0; g < numGroups; ++g) {
    for (int r = 0; r < dim[0]; r++) {
      for (int c = 0; c < dim[1]; c++) {
        lhsMatrix[g][r][c] = static_cast<float>(randomEngine()) / INT_MAX;
      }
    }
  }

  int blockRow;
  int blockCol;
  int nRowBlock;
  int nColBlock;
  if (!rhsNeedTranspose) {
    blockRow = blockSize[1];
    blockCol = blockSize[2];
    nRowBlock = dim[1] / blockSize[1];
    nColBlock = dim[2] / blockSize[2];
  } else {
    blockRow = blockSize[2];
    blockCol = blockSize[1];
    nRowBlock = dim[2] / blockSize[2];
    nColBlock = dim[1] / blockSize[1];
  }

  std::size_t sparsityDenseSize = nRowBlock * nColBlock;
  const unsigned char *itemSparsity = sparsity;
  for (int g = 0; g < numGroups; ++g) {
    for (int i = 0; i < nRowBlock; i++) {
      for (int j = 0; j < nColBlock; j++) {
        int blockRowStart = i * blockRow;
        int blockColStart = j * blockCol;
        int blockRowEnd = blockRowStart + blockRow;
        int blockColEnd = blockColStart + blockCol;
        if (!itemSparsity || itemSparsity[i * nColBlock + j] == 1) {
          for (int r = blockRowStart; r < blockRowEnd; r++) {
            for (int c = blockColStart; c < blockColEnd; c++) {
              rhsMatrix[g][r][c] = static_cast<float>(randomEngine()) /
                                   static_cast<float>(INT_MAX);
            }
          }
        }
      }
    }
    if (sparsity) {
      itemSparsity += sparsityDenseSize;
    }
  }
}

void populateMatrixData(const std::array<int, 2> &dim,
                        const std::array<int, 2> &blockSize,
                        boost::multi_array<float, 2> &matrix,
                        const unsigned char *sparsity) {
  std::mt19937 randomEngine;
  randomEngine.seed(102302);

  int blockRow = blockSize[0];
  int blockCol = blockSize[1];
  int nRowBlock = dim[0] / blockSize[0];
  int nColBlock = dim[1] / blockSize[1];

  for (int i = 0; i < nRowBlock; i++) {
    for (int j = 0; j < nColBlock; j++) {
      int blockRowStart = i * blockRow;
      int blockColStart = j * blockCol;
      int blockRowEnd = blockRowStart + blockRow;
      int blockColEnd = blockColStart + blockCol;

      float val = 0.0f;
      if (!sparsity || sparsity[i * nColBlock + j] == 1) {
        for (int r = blockRowStart; r < blockRowEnd; r++) {
          for (int c = blockColStart; c < blockColEnd; c++) {
            val = static_cast<float>(randomEngine()) /
                  static_cast<float>(INT_MAX);
            matrix[r][c] = val;
          }
        }
      }
    }
  }
}

void populateMatrixData(const std::array<int, 3> &dim,
                        boost::multi_array<float, 2> &lhsMatrix,
                        boost::multi_array<float, 2> &rhsMatrix) {

  std::mt19937 randomEngine;
  randomEngine.seed(102302);

  for (int i = 0; i < dim[0]; i++) {
    for (int j = 0; j < dim[1]; j++) {
      lhsMatrix[i][j] = static_cast<float>(randomEngine()) / INT_MAX;
    }
  }
  for (int i = 0; i < dim[1]; i++) {
    for (int j = 0; j < dim[2]; j++) {
      rhsMatrix[i][j] = static_cast<float>(randomEngine()) / INT_MAX;
    }
  }
}

void populateMatrixData(const std::array<int, 2> &dim,
                        boost::multi_array<float, 2> &matrix) {

  std::mt19937 randomEngine;
  randomEngine.seed(102302);

  for (int i = 0; i < dim[0]; i++) {
    for (int j = 0; j < dim[1]; j++) {
      matrix[i][j] = static_cast<float>(randomEngine()) / INT_MAX;
    }
  }
}

void printMatrix(std::string msg, float *matrix, int row, int col) {
#ifdef DEBUG
  std::cout << msg << std::endl;

  for (int i = 0; i < row; i++) {
    for (int j = 0; j < col; j++) {
      std::cout << matrix[i * col + j] << " ";
    }
    std::cout << std::endl;
  }
#endif
}

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

  unsigned sparseSize = nonZeroBlock * rowsInBlock * colsInBlock;
  sparseMat.resize(boost::extents[sparseSize]);

  unsigned index = 0;
  for (int br = 0, idxDense = 0; br < blockRows; br++) {
    for (int bc = 0; bc < blockCols; ++bc, ++idxDense) {
      if (sparsity[idxDense] == 1) {
        int rowStart = br * rowsInBlock;
        int colStart = bc * colsInBlock;
        int rowEnd = rowStart + rowsInBlock;
        int colEnd = colStart + colsInBlock;
        for (int r = rowStart; r < rowEnd; r++) {
          for (int c = colStart; c < colEnd; c++) {
            assert(index < sparseSize);
            sparseMat[index++] = denseMat[r][c];
          }
        }
      }
    }
  }
  assert(index == sparseSize);
}

void getSparseMatrixBlocks(int rows, int cols, int rowsInBlock, int colsInBlock,
                           int numGroups, unsigned char *sparsity,
                           unsigned nonZeroBlock,
                           const boost::multi_array<float, 3> &denseMat,
                           boost::multi_array<float, 1> &sparseMat) {
  const int blockRows = rows / rowsInBlock;
  const int blockCols = cols / colsInBlock;

  unsigned sparseSize = nonZeroBlock * rowsInBlock * colsInBlock;
  sparseMat.resize(boost::extents[sparseSize]);

  std::size_t sparsityDenseSize = blockRows * blockCols;
  const unsigned char *itemSparsity = sparsity;
  unsigned index = 0;
  for (int g = 0; g < numGroups; ++g) {
    int blockCount = 0;
    for (int br = 0; br < blockRows; br++) {
      for (int bc = 0; bc < blockCols; bc++) {
        if (itemSparsity[br * blockCols + bc] == 1) {
          int rowStart = br * rowsInBlock;
          int colStart = bc * colsInBlock;
          int rowEnd = rowStart + rowsInBlock;
          int colEnd = colStart + colsInBlock;
          for (int r = rowStart; r < rowEnd; r++) {
            for (int c = colStart; c < colEnd; c++) {
              assert(index < sparseSize);
              sparseMat[index++] = denseMat[g][r][c];
            }
          }
          blockCount++;
        }
      }
    }
    itemSparsity += sparsityDenseSize;
  }
  assert(index == sparseSize);
}

BSMatMulParams createBsMatMul(const std::array<int, 3> &dim,
                              const std::array<int, 3> &blockSize,
                              const std::vector<unsigned char> &sparsity,
                              poplar::Type dataType, bool isResSparse,
                              bool rhsNeedTranspose, SubBlockMask mask,
                              unsigned int numGroups) {
  if (isResSparse) {
    return BSMatMulParams(dim, blockSize, sparsity, dataType, dataType,
                          dataType, mask, numGroups);
  } else
    return BSMatMulParams(dim, blockSize, sparsity, rhsNeedTranspose, dataType,
                          dataType, dataType, numGroups);
}

enum class OperandsType {
  null,
  d,   // 1 operand: dense
  s,   // 1 operand: sparse
  ddd, // 3 operands: dense, dense, dense
  dsd, // 3 operands: dense, sparse, dense
  dds, // 3 operands: dense, dense, sparse
};

enum Scenario {
  ddd = 0x0,
  dsd = 0x1,
  dds = 0x2,
  smd = 0x3,
  sms = 0x4,
  fwdMask = 0x7,
  bwd = 0x8,
  wu = 0x10,
  all = bwd | wu,
  dsdBwd = dsd | bwd,
  dsdWu = dsd | wu,
  dsdAll = dsdBwd | dsdWu,
  ddsBwd = dds | bwd,
  ddsWu = dds | wu,
  ddsAll = ddsBwd | ddsWu,
  smdBwd = smd | bwd,
  smsBwd = sms | bwd,
};

int main(int argc, char **argv) {
  // Default options
  DeviceType deviceType = DeviceType::Hw;
  std::string sparsityFileName = "";
  std::string profileDir = ".";
  std::string subBlockMask = "None";
  std::string scenarioName = "dsd";
  std::string partitionMethod = "strip";
  int lhsBlockRowSize0 = 8;
  int lhsBlockColSize0 = 8;
  int rhsBlockSize0 = 8;
  int lhsRows0 = 0;
  int lhsCols0 = 0;
  int rhsCols0 = 0;
  int lhsBlockCols0 = 0;
  int rhsNeedTranspose0 = 0;
  boost::optional<unsigned> tilesPerIPU;
  Type dataType = FLOAT;
  float memoryCycleRatio = 1;
  int runs = 1;
  int nPass = 1;
  int numReps = 1;
  int numGroups = 1;

  namespace po = boost::program_options;

  po::options_description desc("Options");
  desc.add_options()
      // help
      ("help", "Produce help message")
      // scenario
      ("scenario",
       po::value<std::string>(&scenarioName)->default_value(scenarioName),
       "Scenario: "
       "ddd|ddd-sm|dsd|dds|dsd-bwd|dsd-wu|dsd-all|dds-bwd|dds-wu|dds-"
       "all|dds-sm|dds-sm-all|sm|smd|sms-bwd|smd-bwd\n"
       "ddd = dense x dense = dense\n"
       "dsd = dense x sparse = dense\n"
       "dds = dense x dense = sparse\n"
       "-bwd = also compute the gradient of the left matrix\n"
       "-wu = also compute the gradient of the right matrix\n"
       "-all = also compute the gradient of left and right matrices\n"
       "-sm = also compute softmax of the result matrix\n"
       "-all = also compute softmax gradient and then use it as an outer "
       "gradient and compute the gradient of left and right matrices\n"
       "sms = compute softmax only on a sparse matrix (LHS block dimensions "
       "used)\n"
       "smd = compute dense softmax on a dense matrix (LHS dimensions used)\n"
       "sms-bwd = compute softmax gradient only on a sparse matrix (LHS block "
       "dimensions used)\n"
       "smd-bwd = compute dense softmax gradient on a dense matrix (LHS "
       "dimensions used)\n")
      // compile-only
      ("compile-only", "Stop after compilation; don't run the program")
      // device-type
      ("device-type",
       po::value<DeviceType>(&deviceType)->default_value(deviceType),
       deviceTypeHelp)
      // tiles-per-ipu
      ("tiles-per-ipu", po::value(&tilesPerIPU), "Number of tiles per IPU")
      // data-type
      ("data-type", po::value<Type>(&dataType)->default_value(dataType),
       "matmul data type")
      // profile
      ("profile", "Output profiling report")
      // profile-execution
      ("profile-execution", "Output execution steps in the profiling report")
      // profile-vars
      ("profile-vars", "Output variables storage in the profiling report")
      // profile-dir
      ("profile-dir", po::value<std::string>(&profileDir),
       "The directory to output profiling report")
      // lhs-rows
      ("lhs-rows", po::value<int>(&lhsRows0),
       "The number of rows for the LHS matrix (for ddd, dsd-, smd- scenarios)")
      // lhs-cols
      ("lhs-cols", po::value<int>(&lhsCols0),
       "The number of cols for the LHS matrix (for ddd, smd- scenarios)")
      // rhs-cols
      ("rhs-cols", po::value<int>(&rhsCols0),
       "The number of cols for the RHS matrix (for ddd scenario)")
      // lhs-block-cols
      ("lhs-block-cols", po::value<int>(&lhsBlockCols0),
       "The number of blocks columns for the LHS matrix (for dds- scenarios)")
      // lhs-block-row
      ("lhs-block-row", po::value<int>(&lhsBlockRowSize0),
       "The block size for the row of rhe LHS matrix (for dsd-, dds, sms- "
       "scenarios)")
      // lhs-block-col
      ("lhs-block-col", po::value<int>(&lhsBlockColSize0),
       "The block size for the column of the LHS matrix (for dsd-, dds, sms- "
       "scenarios)")
      // rhs-block
      ("rhs-block", po::value<int>(&rhsBlockSize0),
       "The block col for the right matrix that does not need be transposed or "
       "the block row for the right matrix that needs be transposed (for dsd-, "
       "dds- scenarios)")
      // rhs-need-transpose
      ("rhs-need-transpose", po::value<int>(&rhsNeedTranspose0),
       "the right matrix need be transposed (for dsd- scenarios)")
      // sparsity-matrix
      ("sparsity-matrix",
       po::value<std::string>(&sparsityFileName)
           ->default_value(sparsityFileName),
       "The file name for the sparsity mask (for dsd-, dds-, sms- scenarios)")
      // memory-cycle-ratio
      ("memory-cycle-ratio", po::value<float>(&memoryCycleRatio),
       "the ratio between memory weight and cycle weight")
      // number-of-pass
      ("number-of-pass", po::value<int>(&nPass)->default_value(nPass),
       "number of pass")
      // sub-block-mask
      ("sub-block-mask",
       po::value<std::string>(&subBlockMask)->default_value(subBlockMask),
       "the mask inside a block: None, ZeroUpperTriangle, "
       "ZeroLowerTriangle")
      // partition-method
      ("partition-method",
       po::value<std::string>(&partitionMethod)->default_value(partitionMethod),
       "The method to generate the computation graph: block, strip")
      // runs
      ("runs", po::value<int>(&runs), "Number of calls to Engine::run")
      // number-or-reps
      ("number-of-reps", po::value<int>(&numReps), "Number of repetitions")
      // number-or-groups
      ("number-of-groups", po::value<int>(&numGroups), "Number of groups")
      // in-place
      ("in-place", "perform softmax operation in place if possible")
      // check-result
      ("check-result", "check if the result is correct");

  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, desc), vm);
    if (vm.count("help") || !vm.count("sparsity-matrix")) {
      std::cout << desc << std::endl;
      return -1;
    }
    po::notify(vm);
  } catch (std::exception &e) {
    std::cerr << "Command line parsing error: " << e.what() << std::endl;
    return -1;
  }

  Scenario scenario = Scenario::ddd;
  // Up to 3 steps
  unsigned numSteps = 1;
  std::array<OperandsType, 3> operandsType = {
      OperandsType::null, OperandsType::null, OperandsType::null};
  std::array<Pass, 3> runType = {Pass::FWD, Pass::FWD, Pass::FWD};
  bool computeMatmul = true;
  bool computeSoftmax = false;
  if (scenarioName == "ddd") {
    scenario = Scenario::ddd;
    operandsType[0] = OperandsType::ddd;
    runType[0] = Pass::FWD;
  } else if (scenarioName == "dsd") {
    scenario = Scenario::dsd;
    operandsType[0] = OperandsType::dsd;
    runType[0] = Pass::FWD;
  } else if (scenarioName == "dds") {
    scenario = Scenario::dds;
    operandsType[0] = OperandsType::dds;
    runType[0] = Pass::FWD;
  } else if (scenarioName == "ddd-sm") {
    scenario = Scenario::ddd;
    operandsType[0] = OperandsType::ddd;
    runType[0] = Pass::FWD;
    computeSoftmax = true;
  } else if (scenarioName == "dds-sm") {
    scenario = Scenario::dds;
    operandsType[0] = OperandsType::dds;
    runType[0] = Pass::FWD;
    computeSoftmax = true;
  } else if (scenarioName == "smd") {
    scenario = Scenario::smd;
    operandsType[0] = OperandsType::d;
    runType[0] = Pass::FWD;
    computeMatmul = false;
    computeSoftmax = true;
  } else if (scenarioName == "sms") {
    scenario = Scenario::sms;
    operandsType[0] = OperandsType::s;
    runType[0] = Pass::FWD;
    computeMatmul = false;
    computeSoftmax = true;
  } else if (scenarioName == "smd-bwd") {
    scenario = Scenario::smdBwd;
    operandsType[0] = OperandsType::d;
    runType[0] = Pass::FWD;
    computeMatmul = false;
    computeSoftmax = true;
  } else if (scenarioName == "sms-bwd") {
    scenario = Scenario::smsBwd;
    operandsType[0] = OperandsType::s;
    runType[0] = Pass::FWD;
    computeMatmul = false;
    computeSoftmax = true;
  } else if (scenarioName == "dsd-bwd") {
    numSteps = 2;
    scenario = Scenario::dsdBwd;
    operandsType[0] = OperandsType::dsd;
    operandsType[1] = OperandsType::dsd;
    runType[0] = Pass::FWD;
    runType[1] = Pass::BWD;
  } else if (scenarioName == "dsd-wu") {
    numSteps = 2;
    scenario = Scenario::dsdWu;
    operandsType[0] = OperandsType::dsd;
    operandsType[1] = OperandsType::dds;
    runType[0] = Pass::FWD;
    runType[1] = Pass::WU;
  } else if (scenarioName == "dsd-all") {
    numSteps = 3;
    scenario = Scenario::dsdAll;
    operandsType[0] = OperandsType::dsd;
    operandsType[1] = OperandsType::dsd;
    operandsType[2] = OperandsType::dds;
    runType[0] = Pass::FWD;
    runType[1] = Pass::BWD;
    runType[2] = Pass::WU;
  } else if (scenarioName == "dds-bwd") {
    numSteps = 2;
    scenario = Scenario::ddsBwd;
    operandsType[0] = OperandsType::dds;
    operandsType[1] = OperandsType::dsd;
    runType[0] = Pass::FWD;
    runType[1] = Pass::BWD;
  } else if (scenarioName == "dds-wu") {
    numSteps = 2;
    scenario = Scenario::ddsWu;
    operandsType[0] = OperandsType::dds;
    operandsType[1] = OperandsType::dsd;
    runType[0] = Pass::FWD;
    runType[1] = Pass::WU;
  } else if (scenarioName == "dds-all") {
    numSteps = 3;
    scenario = Scenario::ddsAll;
    operandsType[0] = OperandsType::dds;
    operandsType[1] = OperandsType::dsd;
    operandsType[2] = OperandsType::dsd;
    runType[0] = Pass::FWD;
    runType[1] = Pass::BWD;
    runType[2] = Pass::WU;
  } else if (scenarioName == "dds-sm-all") {
    numSteps = 3;
    scenario = Scenario::ddsAll;
    operandsType[0] = OperandsType::dds;
    operandsType[1] = OperandsType::dsd;
    operandsType[2] = OperandsType::dsd;
    runType[0] = Pass::FWD;
    runType[1] = Pass::BWD;
    runType[2] = Pass::WU;
    computeSoftmax = true;
  } else {
    std::cerr << "--scenario: Unrecognized scenario: " << scenarioName
              << std::endl;
    return -1;
  }
  std::cout << scenarioName << " scenario" << std::endl;
  if (numGroups > 1) {
    std::cout << "Number of groups: " << numGroups << std::endl;
  }

  std::map<std::string, std::pair<int *, std::vector<unsigned int>>>
      optionsUsage = {
          {"lhs-rows",
           {&lhsRows0, {0, 1, Scenario::ddd, Scenario::dsd, Scenario::smd}}},
          {"lhs-cols", {&lhsCols0, {0, 1, Scenario::ddd, Scenario::smd}}},
          {"rhs-cols", {&rhsCols0, {0, 1, Scenario::ddd}}},
          {"lhs-block-cols", {&lhsBlockCols0, {0, 1, Scenario::dds}}},
          {"lhs-block-row",
           {&lhsBlockRowSize0,
            {0, 1, Scenario::dsd, Scenario::dds, Scenario::sms}}},
          {"lhs-block-col",
           {&lhsBlockColSize0,
            {0, 1, Scenario::dsd, Scenario::dds, Scenario::sms}}},
          {"rhs-block", {&rhsBlockSize0, {0, 1, Scenario::dsd, Scenario::dds}}},
          {"rhs-need-transpose", {&rhsNeedTranspose0, {1, 0, Scenario::dsd}}},
          {"sparsity-matrix",
           {nullptr, {0, 0, Scenario::dsd, Scenario::dds, Scenario::sms}}},
          {"sub-block-mask",
           {nullptr, {1, 0, Scenario::dsd, Scenario::dds, Scenario::sms}}},
          {"num-groups", {&numGroups, {1, 1, Scenario::dsd, Scenario::dds}}},
      };

  for (auto iter = optionsUsage.begin(); iter != optionsUsage.end(); ++iter) {
    const std::string option = iter->first;
    const std::vector<unsigned int> &usedIn = iter->second.second;
    int *optionValue = iter->second.first;
    // First 2 elements of usedIn has special meanings:
    unsigned int hasDefault = usedIn[0];
    int minValue = static_cast<int>(usedIn[1]);
    bool used = false;
    for (unsigned int i = 2; i < usedIn.size(); ++i) {
      if ((scenario & Scenario::fwdMask) == usedIn[i]) {
        used = true;
        break;
      }
    }
    if (vm.count(option) > 0) {
      if (!used) {
        logging::popsparse::warn(
            "Option {} is not used in scenario {}. Ignored.", option.c_str(),
            scenarioName.c_str());
      } else {
        if (optionValue && *optionValue < minValue) {
          std::cerr << "Option " << option << " minimum value is "
                    << *optionValue << std::endl;
          return -1;
        }
      }
    } else {
      if (used && !hasDefault) {
        std::cerr << "Option " << option << " is missing." << std::endl;
        return -1;
      }
    }
  }
  bool compileIPUCode = false;
  if (isIpuModel(deviceType))
    compileIPUCode = true;
  const unsigned numIPUs = 1;
  auto device =
      tilesPerIPU
          ? createTestDevice(deviceType, numIPUs, *tilesPerIPU, compileIPUCode)
          : createTestDeviceFullSize(deviceType, numIPUs, compileIPUCode);

  bool inPlace = vm.count("in-place") > 0;
  if (!computeSoftmax) {
    logging::popsparse::warn(
        "Option in-place is used only with softmax scenarios. Ignored.");
  }

  bool checkResult = vm.count("check-result") > 0;

  if (checkResult && (((scenario & Scenario::bwd) == Scenario::bwd) ||
                      ((scenario & Scenario::wu) == Scenario::wu))) {
    std::cerr << "Checking results in not supported in scenario "
              << scenarioName << std::endl;
    return -1;
  }
  if (checkResult && inPlace &&
      (scenario == Scenario::smd || scenario == Scenario::sms)) {
    std::cerr << "Checking results in not supported in in-place scenario "
              << scenarioName << std::endl;
    return -1;
  }

  if (checkResult && numReps > 1) {
    std::cerr << "Checking results in not supported for many repetitions."
              << std::endl;
    return -1;
  }

  if (((scenario & Scenario::fwdMask) != Scenario::dsd) &&
      ((scenario & Scenario::fwdMask) != Scenario::dds) &&
      ((scenario & Scenario::fwdMask) != Scenario::sms)) {
    numGroups = 1;
  }

  bool hasSparsity = ((scenario & Scenario::fwdMask) == Scenario::dsd) ||
                     ((scenario & Scenario::fwdMask) == Scenario::dds) ||
                     ((scenario & Scenario::fwdMask) == Scenario::sms);

  SubBlockMask mask = SubBlockMask::None;
  if (hasSparsity) {
    if (subBlockMask == "ZeroUpperTriangle") {
      mask = SubBlockMask::ZeroUpperTriangle;
    } else if (subBlockMask == "ZeroLowerTriangle") {
      mask = SubBlockMask::ZeroLowerTriangle;
    } else if (subBlockMask != "None") {
      std::cerr << "Unrecognized sub-block mask parameter: " << subBlockMask
                << std::endl;
      return -1;
    }
  }

  bool doProfiling = vm.count("profile") > 0;
  bool doProfilingExecution = doProfiling && vm.count("profile-execution") > 0;
  bool doProfilingVars = doProfiling && vm.count("profile-vars") > 0;

  int nonZeroBlock = 0;
  int sparsityRows, sparsityCols;
  std::vector<unsigned char> sparsity;
  std::size_t sparsityDenseSize = 0;
  if (hasSparsity) {
    bool readTransposed =
        ((scenario & Scenario::fwdMask) == Scenario::dsd) && rhsNeedTranspose0;
    if (!ReadMatrixMask(sparsityFileName, nonZeroBlock, sparsityRows,
                        sparsityCols, sparsity, readTransposed)) {
      return -1;
    }

    if (numGroups == 1) {
      std::cout << "non zero block: " << nonZeroBlock << std::endl;
      std::cout << "sparsity: "
                << (1.0 - nonZeroBlock /
                              static_cast<float>(sparsityRows * sparsityCols))
                << std::endl;
    } else {
      if (sparsityRows % numGroups != 0) {
        std::cerr << "error: Number of rows " << sparsityRows
                  << " in concatenated group sparsity mask "
                     "is not divisible by the number of groups "
                  << numGroups << std::endl;
        return -1;
      }
      std::cout << "Average non zero block: "
                << static_cast<float>(nonZeroBlock) / numGroups << std::endl;
      std::cout << "Average sparsity: "
                << (1.0 - nonZeroBlock /
                              static_cast<float>(sparsityRows * sparsityCols))
                << std::endl;
      sparsityRows /= numGroups;
    }
    sparsityDenseSize = sparsityRows * sparsityCols;
  }

  int lhsBlockRows0 = 0;
  int rhsBlockCols0 = 0;
  bool isRhsMatrixSparse0 = false;
  bool isResMatrixSparse0 = false;

  int resBlockRows0 = 0;
  int resBlockCols0 = 0;
  int resBlockRowSize0 = 0;
  int resBlockColSize0 = 0;
  int resRows0;
  int resCols0;

  if (computeMatmul) {
    if (hasSparsity) {
      isRhsMatrixSparse0 = (operandsType[0] == OperandsType::dsd);
      isResMatrixSparse0 = (operandsType[0] == OperandsType::dds);

      if (isRhsMatrixSparse0) {
        if (lhsRows0 % lhsBlockRowSize0 != 0) {
          std::cerr << "error: Number of LHS rows " << lhsRows0
                    << " is not divisible "
                       "by the number of rows in a block for the LHS matrix "
                    << lhsBlockRowSize0 << std::endl;
          return -1;
        }
        lhsBlockRows0 = lhsRows0 / lhsBlockRowSize0;
        lhsBlockCols0 = sparsityRows;
      } else if (isResMatrixSparse0) {
        lhsBlockRows0 = sparsityRows;
        lhsRows0 = lhsBlockRows0 * lhsBlockRowSize0;
        rhsNeedTranspose0 = false;
      } else {
        assert(0);
      }

      rhsBlockCols0 = sparsityCols;
      lhsCols0 = lhsBlockCols0 * lhsBlockColSize0;

      resBlockRows0 = lhsBlockRows0;
      resBlockCols0 = rhsBlockCols0;
      resBlockRowSize0 = lhsBlockRowSize0;
      resBlockColSize0 = rhsBlockSize0;
      resRows0 = resBlockRows0 * resBlockRowSize0;
      resCols0 = resBlockCols0 * resBlockColSize0;
    } else {
      rhsNeedTranspose0 = false;
      resRows0 = lhsRows0;
      resCols0 = rhsCols0;
    }
  } else if (computeSoftmax) {
    if (hasSparsity) {
      resBlockRows0 = sparsityRows;
      resBlockCols0 = sparsityCols;
      resBlockRowSize0 = lhsBlockRowSize0;
      resBlockColSize0 = lhsBlockColSize0;
      resRows0 = resBlockRows0 * resBlockRowSize0;
      resCols0 = resBlockCols0 * resBlockColSize0;
      isResMatrixSparse0 = true;
    } else {
      resRows0 = lhsRows0;
      resCols0 = lhsCols0;
    }
  } else {
    assert(0);
  }

  const std::string debugPrefix = "bs_matmul";
  Sequence matSeq, uploadProg, downloadProg;

  std::vector<std::unique_ptr<char[]>> lhsRawHost(numSteps);
  std::vector<std::unique_ptr<char[]>> rhsRawHost(numSteps);
  std::vector<std::unique_ptr<char[]>> inputRawHost(numSteps);
  std::unique_ptr<char[]> outputRawHost;
  std::vector<std::pair<std::string, char *>> tmap;

  const auto &target = device.getTarget();
  Graph graph(target);

  popops::addCodelets(graph);
  poplin::addCodelets(graph);

  double flops = 0.0;

  boost::multi_array<float, 2> lhsHost;
  boost::multi_array<float, 2> rhsHost;
  boost::multi_array<float, 2> inputHost;

  boost::multi_array<float, 3> lhsHost3D;
  boost::multi_array<float, 3> rhsHost3D;
  boost::multi_array<float, 3> inputHost3D;

  poplar::Tensor aTensor;
  poplar::Tensor bTensor;
  poplar::Tensor dCTensor;
  poplar::Tensor dBTensor;

  poplar::OptionFlags options = {
      {"memoryCycleRatio", std::to_string(memoryCycleRatio)},
      {"numberOfPass", std::to_string(nPass)},
      {"partitionMethod", partitionMethod},
  };

  for (int nRep = 0; nRep < numReps; ++nRep) {
    for (unsigned nStep = 0; nStep < numSteps; ++nStep) {
      std::string stepSuffix =
          numSteps == 1 ? "" : "(" + std::to_string(nStep) + ")";
      if (computeMatmul) {
        int lhsBlockCols = lhsBlockCols0;
        int rhsBlockCols = rhsBlockCols0;
        int lhsRows = lhsRows0;
        int lhsCols = lhsCols0;
        int rhsCols = rhsCols0;
        int lhsBlockRowSize = lhsBlockRowSize0;
        int lhsBlockColSize = lhsBlockColSize0;
        int rhsBlockSize = rhsBlockSize0;
        bool rhsNeedTranspose = rhsNeedTranspose0;

        if ((scenario & Scenario::fwdMask) == Scenario::dsd) {
          if (runType[nStep] == Pass::BWD) {
            lhsBlockCols = rhsBlockCols0;
            rhsBlockCols = lhsBlockCols0;
            lhsRows = lhsRows0;
            lhsBlockColSize = rhsBlockSize0;
            rhsBlockSize = lhsBlockColSize0;
            rhsNeedTranspose = !rhsNeedTranspose;
          } else if (runType[nStep] == Pass::WU) {
            lhsBlockCols = lhsBlockRows0;
            rhsBlockCols = rhsBlockCols0;
            lhsRows = lhsCols0;
            lhsBlockRowSize = lhsBlockColSize0;
            lhsBlockColSize = lhsBlockRowSize0;
            rhsBlockSize = rhsBlockSize0;
            rhsNeedTranspose = false;
          }
        } else if ((scenario & Scenario::fwdMask) == Scenario::dds) {
          if (runType[nStep] == Pass::BWD) {
            lhsBlockCols = rhsBlockCols0;
            rhsBlockCols = lhsBlockRows0;
            lhsRows = lhsCols0;
            lhsBlockColSize = rhsBlockSize0;
            lhsBlockRowSize = lhsBlockColSize0;
            rhsBlockSize = lhsBlockRowSize0;
          } else if (runType[nStep] == Pass::WU) {
            lhsBlockCols = lhsBlockRows0;
            rhsBlockCols = rhsBlockCols0;
            lhsRows = lhsCols0;
            lhsBlockColSize = lhsBlockRowSize0;
            lhsBlockRowSize = lhsBlockColSize0;
            rhsBlockSize = rhsBlockSize0;
            rhsNeedTranspose = true;
          }
        }
        if (hasSparsity) {
          lhsCols = lhsBlockCols * lhsBlockColSize;
          rhsCols = rhsBlockCols * rhsBlockSize;
        }
        int lhsBlockRows = lhsRows / lhsBlockRowSize;

        int resBlockRows = lhsBlockRows;
        int resBlockCols = rhsBlockCols;
        int resBlockRowSize = lhsBlockRowSize;
        int resBlockColSize = rhsBlockSize;
        int resRows = resBlockRows * resBlockRowSize;
        int resCols = resBlockCols * resBlockColSize;

        bool isResMatrixSparse = (operandsType[nStep] == OperandsType::dds ||
                                  operandsType[nStep] == OperandsType::s);
        bool isRhsMatrixSparse = (operandsType[nStep] == OperandsType::dsd);

        std::array<int, 3> dim;
        dim[0] = lhsRows;
        dim[1] = lhsCols;
        dim[2] = rhsCols;

        std::array<int, 3> blockSize = {lhsBlockRowSize, lhsBlockColSize,
                                        rhsBlockSize};

        std::cout << "Step " << nStep << std::endl;
        std::cout << "dim: " << dim[0] << " x " << dim[1] << " x " << dim[2]
                  << std::endl;
        if (operandsType[nStep] != OperandsType::ddd) {
          std::cout << "block size: " << blockSize[0] << " x " << blockSize[1]
                    << " x " << blockSize[2] << std::endl;
        }
        if (rhsNeedTranspose) {
          logging::popsparse::info("Right matrix must be transposed");
        }
        logging::popsparse::info("LHS rows: {} ", lhsRows);

        float flopsOp;
        if (operandsType[nStep] == OperandsType::ddd) {
          flopsOp = 2.0 * dim[0] * dim[1] * dim[2];
        } else {
          if (!isResMatrixSparse) {
            flopsOp = 2.0 * nonZeroBlock * dim[0] * blockSize[1] * blockSize[2];
          } else {
            flopsOp = 2.0 * nonZeroBlock * dim[1] * blockSize[0] * blockSize[2];
          }
        }
        logging::popsparse::info("flops op: {:.1f} ", flopsOp);
        flops += flopsOp;

        if (numGroups == 1) {
          lhsHost.resize(boost::extents[dim[0]][dim[1]]);
          if (!rhsNeedTranspose) {
            rhsHost.resize(boost::extents[dim[1]][dim[2]]);
          } else {
            rhsHost.resize(boost::extents[dim[2]][dim[1]]);
          }
          if (operandsType[nStep] == OperandsType::ddd) {
            populateMatrixData(dim, lhsHost, rhsHost);
          } else {
            populateMatrixData(dim, blockSize, lhsHost, rhsHost,
                               (!isRhsMatrixSparse ? nullptr : sparsity.data()),
                               rhsNeedTranspose);
          }
        } else {
          lhsHost3D.resize(boost::extents[numGroups][dim[0]][dim[1]]);
          if (!rhsNeedTranspose) {
            rhsHost3D.resize(boost::extents[numGroups][dim[1]][dim[2]]);
          } else {
            rhsHost3D.resize(boost::extents[numGroups][dim[2]][dim[1]]);
          }
          populateMatrixData(dim, blockSize, numGroups, lhsHost3D, rhsHost3D,
                             (!isRhsMatrixSparse ? nullptr : sparsity.data()),
                             rhsNeedTranspose);
        }

        if (operandsType[nStep] == OperandsType::ddd) {
          poplin::matmul::PlanningCache cache;
          OptionFlags mmOpt;

          // LHS matrix
          poplar::Tensor lhsTensor = graph.addVariable(
              dataType, {(unsigned long)dim[0], (unsigned long)dim[1]},
              debugPrefix + "/matrix_lhs");
          poputil::mapTensorLinearly(graph, lhsTensor);

          if (numReps == 1) { // Copying data from host is not supported for
                              // multi-rep tests yet
            lhsRawHost[nStep] = allocateHostMemoryForTensor(
                lhsTensor, debugPrefix + "/matrix_lhs" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
            assert(lhsHost.num_elements() == lhsTensor.numElements());
            copy(target, lhsHost, dataType, lhsRawHost[nStep].get());
          }

          // RHS matrix
          poplar::Tensor rhsTensor = graph.addVariable(
              dataType, {(unsigned long)dim[1], (unsigned long)dim[2]},
              debugPrefix + "/matrix_rhs");
          poputil::mapTensorLinearly(graph, rhsTensor);

          if (numReps == 1) {
            rhsRawHost[nStep] = allocateHostMemoryForTensor(
                rhsTensor, debugPrefix + "/matrix_rhs" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
            assert(rhsHost.num_elements() == rhsTensor.numElements());
            copy(target, rhsHost, dataType, rhsRawHost[nStep].get());
          }

          // output matrix
          poplar::Tensor outTensor =
              matMul(graph, lhsTensor, rhsTensor, matSeq,
                     debugPrefix + "/lhsxrhs", mmOpt, &cache);
          if (runType[nStep] == Pass::FWD && computeSoftmax) {
            float nonLinearityScaling;
            if (!inPlace) {
              outTensor = popnn::nonLinearity(
                  graph, popnn::NonLinearityType::SOFTMAX_STABLE, outTensor,
                  nonLinearityScaling, matSeq, debugPrefix + "/softmax");
            } else {
              popnn::nonLinearity(
                  graph, popnn::NonLinearityType::SOFTMAX_STABLE, outTensor,
                  nonLinearityScaling, matSeq, debugPrefix + "/softmax");
            }
          }
          if (checkResult == 1) {
            outputRawHost = poplibs_test::util::allocateHostMemoryForTensor(
                outTensor, debugPrefix + "/matrix_output" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
          }
        } else {
          BSMatMulParams bsMatMulObj = createBsMatMul(
              dim, blockSize, sparsity, dataType, isResMatrixSparse,
              rhsNeedTranspose, mask, numGroups);

          // LHS dense matrix
          poplar::Tensor lhsTensor;
          poplar::Tensor rhsTensor;
          if (runType[nStep] == Pass::FWD) {
            lhsTensor = createBSMatMulInputLHS(
                graph, bsMatMulObj, debugPrefix + "/matrix_lhs", options);
            rhsTensor = createBSMatMulInputRHS(
                graph, bsMatMulObj, debugPrefix + "/matrix_rhs", options);
            if (((scenario & Scenario::bwd) == Scenario::bwd) ||
                ((scenario & Scenario::wu) == Scenario::wu)) {
              // Save A tensor
              aTensor = lhsTensor;
              // Save B tensor
              bTensor = rhsTensor;
            }
          } else if (runType[nStep] == Pass::BWD) {
            if ((scenario & Scenario::fwdMask) == Scenario::dsd) {
              lhsTensor = createBSMatMulInputLHS(
                  graph, bsMatMulObj, debugPrefix + "/matrix_lhs", options);
              rhsTensor = bTensor;
              if ((scenario & Scenario::wu) == Scenario::wu) {
                // Save dC tensor
                dCTensor = lhsTensor;
              }
            } else if ((scenario & Scenario::fwdMask) == Scenario::dds) {
              lhsTensor = bTensor;
              rhsTensor = createBSMatMulInputRHS(
                  graph, bsMatMulObj, debugPrefix + "/matrix_rhs", options);
              if ((scenario & Scenario::wu) == Scenario::wu) {
                // Save dB tensor
                dBTensor = rhsTensor;
              }
            } else {
              assert(0);
            }
          } else if (runType[nStep] == Pass::WU) {
            if ((scenario & Scenario::fwdMask) == Scenario::dsd) {
              assert(aTensor.valid());
              lhsTensor = aTensor.transpose();
              if ((scenario & Scenario::bwd) == Scenario::bwd) {
                // If tensor was saved
                assert(dCTensor.valid());
                rhsTensor = dCTensor;
              } else {
                rhsTensor = createBSMatMulInputRHS(
                    graph, bsMatMulObj, debugPrefix + "/matrix_rhs", options);
              }
            } else if ((scenario & Scenario::fwdMask) == Scenario::dds) {
              assert(aTensor.valid());
              lhsTensor = aTensor.transpose();
              if ((scenario & Scenario::bwd) == Scenario::bwd) {
                // If tensor was saved
                assert(dBTensor.valid());
                rhsTensor = dBTensor;
              } else {
                rhsTensor = createBSMatMulInputRHS(
                    graph, bsMatMulObj, debugPrefix + "/matrix_rhs", options);
              }
            } else {
              assert(0);
            }
          } else {
            assert(0);
          }

          // sparse matmul
          poplar::Tensor outTensor;
          try {
            outTensor =
                bsMatMul(graph, bsMatMulObj, matSeq, lhsTensor, rhsTensor,
                         options, debugPrefix + "/bs-matmul");
          } catch (const poputil::poplibs_error &e) {
            std::cerr << "bsMatMul() failed with the message: " << e.what()
                      << std::endl;
            return -1;
          }
          if (isResMatrixSparse) {
            assert((int)outTensor.dim(0) == nonZeroBlock);
          }
          if ((scenario & Scenario::fwdMask) == Scenario::dds) {
            if (runType[nStep] == Pass::BWD) {
              outTensor = outTensor.transpose();
            }
          }
          if (runType[nStep] == Pass::BWD) {
            assert(outTensor.shape() == aTensor.shape());
          } else if (runType[nStep] == Pass::WU) {
            assert(outTensor.shape() == bTensor.shape());
          }
          if (runType[nStep] == Pass::FWD && computeSoftmax) {
            std::array<int, 2> dimSm;
            std::array<int, 2> blockSizeSm;
            dimSm[0] = resRows;
            dimSm[1] = resCols;
            blockSizeSm[0] = resBlockRowSize;
            blockSizeSm[1] = resBlockColSize;

            if (!inPlace) {
              Tensor softmaxTensor = bsSoftmax(
                  graph, outTensor, dimSm, blockSizeSm, sparsity, mask,
                  numGroups, matSeq, debugPrefix + "/bs-softmax");
              outTensor = softmaxTensor;
            } else {
              bsSoftmaxInPlace(graph, outTensor, dimSm, blockSizeSm, sparsity,
                               mask, numGroups, matSeq,
                               debugPrefix + "/bs-softmax");
            }
            if (((scenario & Scenario::bwd) == Scenario::bwd) ||
                ((scenario & Scenario::wu) == Scenario::wu)) {
              // Create helper BSMatMulParams in order to generate a matrix in
              // block-sparse format w/out extra matmul overhead
              std::array<int, 3> dimHelper;
              std::array<int, 3> blockSizeHelper;
              dimHelper[0] = blockSize[0]; // Can be anything
              dimHelper[1] = resRows;
              dimHelper[2] = resCols;
              blockSizeHelper[0] = blockSize[0]; // Can be anything
              blockSizeHelper[1] = resBlockRowSize;
              blockSizeHelper[2] = resBlockColSize;

              BSMatMulParams bsMatMulObjHelper =
                  createBsMatMul(dimHelper, blockSizeHelper, sparsity, dataType,
                                 false, false, mask, numGroups);
              Tensor outerGrad =
                  createBSMatMulInputRHS(graph, bsMatMulObjHelper,
                                         debugPrefix + "/outer_grad", options);
              assert(outerGrad.shape() == outTensor.shape());
              Tensor softmaxGradTensor = bsSoftmaxGrad(
                  graph, outTensor, outerGrad, dimSm, blockSizeSm, sparsity,
                  matSeq, debugPrefix + "/bs-softmax-grad");
            }
          }

          if (numReps == 1) {
            lhsRawHost[nStep] = allocateHostMemoryForTensor(
                lhsTensor, debugPrefix + "/matrix_lhs" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
            if (numGroups == 1) {
              assert(lhsHost.num_elements() == lhsTensor.numElements());
              copy(target, lhsHost, dataType, lhsRawHost[nStep].get());
            } else {
              assert(lhsHost3D.num_elements() == lhsTensor.numElements());
              copy(target, lhsHost3D, dataType, lhsRawHost[nStep].get());
            }
          }

          if (isRhsMatrixSparse) {
            // RHS sparse matrix
            boost::multi_array<float, 1> rhsBlocksHost;
            if (numGroups == 1) {
              if (!rhsNeedTranspose) {
                getSparseMatrixBlocks(dim[1], dim[2], blockSize[1],
                                      blockSize[2], sparsity.data(),
                                      nonZeroBlock, rhsHost, rhsBlocksHost);
              } else {
                getSparseMatrixBlocks(dim[2], dim[1], blockSize[2],
                                      blockSize[1], sparsity.data(),
                                      nonZeroBlock, rhsHost, rhsBlocksHost);
              }
            } else {
              if (!rhsNeedTranspose) {
                getSparseMatrixBlocks(dim[1], dim[2], blockSize[1],
                                      blockSize[2], numGroups, sparsity.data(),
                                      nonZeroBlock, rhsHost3D, rhsBlocksHost);
              } else {
                getSparseMatrixBlocks(dim[2], dim[1], blockSize[2],
                                      blockSize[1], numGroups, sparsity.data(),
                                      nonZeroBlock, rhsHost3D, rhsBlocksHost);
              }
            }
            if (numReps == 1) {
              rhsRawHost[nStep] = allocateHostMemoryForTensor(
                  rhsTensor, debugPrefix + "/matrix_rhs" + stepSuffix, graph,
                  uploadProg, downloadProg, tmap);
              assert(rhsBlocksHost.num_elements() == rhsTensor.numElements());
              copy(target, rhsBlocksHost, dataType, rhsRawHost[nStep].get());
            }
          } else {
            // RHS dense matrix
            if (numReps == 1) {
              rhsRawHost[nStep] = allocateHostMemoryForTensor(
                  rhsTensor, debugPrefix + "/matrix_rhs" + stepSuffix, graph,
                  uploadProg, downloadProg, tmap);
              if (numGroups == 1) {
                assert(rhsHost.num_elements() == rhsTensor.numElements());
                copy(target, rhsHost, dataType, rhsRawHost[nStep].get());
              } else {
                assert(rhsHost3D.num_elements() == rhsTensor.numElements());
                copy(target, rhsHost3D, dataType, rhsRawHost[nStep].get());
              }
            }
          }
          if (checkResult == 1) {
            outputRawHost = poplibs_test::util::allocateHostMemoryForTensor(
                outTensor, debugPrefix + "/matrix_output" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
          }
        }
      } else if (computeSoftmax) {
        std::array<int, 2> dim = {resRows0, resCols0};

        std::cout << "dim: " << dim[0] << " x " << dim[1] << std::endl;

        if (hasSparsity) {
          std::array<int, 2> blockSize = {resBlockRowSize0, resBlockColSize0};

          // Create helper BSMatMulParams in order to generate a matrix in
          // block-sparse format w/out extra matmul overhead
          std::array<int, 3> dimHelper;
          std::array<int, 3> blockSizeHelper;
          dimHelper[0] = blockSize[0]; // Can be anything
          dimHelper[1] = resRows0;
          dimHelper[2] = resCols0;
          blockSizeHelper[0] = blockSize[0]; // Can be anything
          blockSizeHelper[1] = resBlockRowSize0;
          blockSizeHelper[2] = resBlockColSize0;

          BSMatMulParams bsMatMulObjHelper =
              createBsMatMul(dimHelper, blockSizeHelper, sparsity, dataType,
                             false, false, mask, numGroups);

          Tensor inputTensor = createBSMatMulInputRHS(
              graph, bsMatMulObjHelper, debugPrefix + "/input", options);
          // Ignore hypergraph mapping
          poputil::mapTensorLinearly(graph, inputTensor);

          Tensor outerGradTensor;
          if (scenario == Scenario::smsBwd) {
            outerGradTensor = createBSMatMulInputRHS(
                graph, bsMatMulObjHelper, debugPrefix + "/outer_grad", options);
            assert(outerGradTensor.shape() == inputTensor.shape());
            // Ignore hypergraph mapping
            poputil::mapTensorLinearly(graph, outerGradTensor);
          }

          poplar::Tensor outTensor;
          if (scenario == Scenario::sms) {
            if (!inPlace) {
              outTensor =
                  bsSoftmax(graph, inputTensor, dim, blockSize, sparsity, mask,
                            numGroups, matSeq, debugPrefix + "/bs-softmax");
            } else {
              bsSoftmaxInPlace(graph, inputTensor, dim, blockSize, sparsity,
                               mask, numGroups, matSeq,
                               debugPrefix + "/bs-softmax");
              outTensor = inputTensor;
            }
          } else if (scenario == Scenario::smsBwd) {
            outTensor = bsSoftmaxGrad(graph, inputTensor, outerGradTensor, dim,
                                      blockSize, sparsity, matSeq,
                                      debugPrefix + "/bs-softmax-grad");
          } else {
            assert(0);
          }

          inputHost.resize(boost::extents[dim[0]][dim[1]]);
          populateMatrixData(dim, blockSize, inputHost, sparsity.data());

          boost::multi_array<float, 1> inputBlocksHost;
          getSparseMatrixBlocks(dim[0], dim[1], blockSize[0], blockSize[1],
                                sparsity.data(), nonZeroBlock, inputHost,
                                inputBlocksHost);
          if (numReps == 1) {
            inputRawHost[nStep] = allocateHostMemoryForTensor(
                inputTensor, debugPrefix + "/input" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
            assert(inputBlocksHost.num_elements() == inputTensor.numElements());
            copy(target, inputBlocksHost, dataType, inputRawHost[nStep].get());
          }
          if (checkResult == 1) {
            outputRawHost = poplibs_test::util::allocateHostMemoryForTensor(
                outTensor, debugPrefix + "/output" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
          }
        } else {
          Tensor inputTensor = graph.addVariable(
              dataType, {(unsigned long)dim[0], (unsigned long)dim[1]},
              debugPrefix + "/input");
          poputil::mapTensorLinearly(graph, inputTensor);

          Tensor outerGradTensor;
          if (scenario == Scenario::smdBwd) {
            outerGradTensor = graph.addVariable(
                dataType, {(unsigned long)dim[0], (unsigned long)dim[1]},
                debugPrefix + "/outer_grad");
            assert(outerGradTensor.shape() == inputTensor.shape());
            poputil::mapTensorLinearly(graph, outerGradTensor);
          }

          poplar::Tensor outTensor;
          if (scenario == Scenario::smd) {
            float nonLinearityScaling;
            if (!inPlace) {
              outTensor = popnn::nonLinearity(
                  graph, popnn::NonLinearityType::SOFTMAX_STABLE, inputTensor,
                  nonLinearityScaling, matSeq, debugPrefix + "/softmax");
            } else {
              popnn::nonLinearity(
                  graph, popnn::NonLinearityType::SOFTMAX_STABLE, inputTensor,
                  nonLinearityScaling, matSeq, debugPrefix + "/softmax");
              outTensor = inputTensor;
            }
          } else if (scenario == Scenario::smdBwd) {
            outTensor = popnn::nonLinearityInputGradient(
                graph, popnn::NonLinearityType::SOFTMAX_STABLE, inputTensor,
                outerGradTensor, matSeq, debugPrefix + "/softmax-grad");
          } else {
            assert(0);
          }

          inputHost.resize(boost::extents[dim[0]][dim[1]]);
          populateMatrixData(dim, inputHost);

          if (numReps == 1) {
            inputRawHost[nStep] = allocateHostMemoryForTensor(
                inputTensor, debugPrefix + "/input" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
            assert(inputHost.num_elements() == inputTensor.numElements());
            copy(target, inputHost, dataType, inputRawHost[nStep].get());
          }
          if (checkResult == 1) {
            outputRawHost = poplibs_test::util::allocateHostMemoryForTensor(
                outTensor, debugPrefix + "/output" + stepSuffix, graph,
                uploadProg, downloadProg, tmap);
          }
        }
      } else {
        assert(0);
      }
    }
  }

  auto engineOptions = defaultEngineOptions;
  if (doProfiling) {
    engineOptions.set("debug.instrumentCompute", "true");
    if (doProfilingExecution) {
      engineOptions.set("debug.instrumentExternalExchange", "true");
    }
    if (doProfilingVars) {
      engineOptions.set("debug.loweredVarDumpFile", profileDir + "/vars.capnp");
    }
  }
  if (vm.count("profile-dir")) {
    engineOptions.set("autoReport.all", "true");
    engineOptions.set("autoReport.directory", profileDir);
  }
  engineOptions.set("exchange.multicastPolicy", "balanced");

  Sequence allSequence;
  // if run many times, we are doing benchmark, so ignore the host data copy
  if (checkResult == 1) {
    allSequence.add(uploadProg);
    allSequence.add(matSeq);
    allSequence.add(downloadProg);
  } else {
    allSequence.add(uploadProg);
    allSequence.add(Repeat(runs, matSeq));
    allSequence.add(downloadProg);
  }

  Engine engine(graph, allSequence, engineOptions);

  if (vm.count("compile-only"))
    std::exit(0);

  std::cout << "Start " << runs << " runs of " << numReps << " rep(s)..."
            << std::endl;
  poplibs_test::util::attachStreams(engine, tmap);
  device.bind([&](const Device &d) {
    engine.load(d);
    auto start = std::chrono::system_clock::now();
    engine.run(0);
    auto end = std::chrono::system_clock::now();
    std::chrono::duration<double> elapsedSeconds = end - start;
    std::cout << "elapsed time: " << std::fixed << std::setprecision(6)
              << elapsedSeconds.count() << std::endl;

    double avgTime = elapsedSeconds.count() / runs;
    std::cout << "average kernel run time: " << std::fixed
              << std::setprecision(1) << avgTime * 1000000 << " mcs"
              << std::endl;
    std::cout << "flop: " << std::fixed << std::setprecision(1) << flops
              << std::endl;
    std::cout << "GFLOPS: " << std::fixed << std::setprecision(6)
              << flops / 1000000000.0 / avgTime << std::endl;
  });

  if (deviceType != DeviceType::Cpu && doProfiling) {
    engine.printProfileSummary(
        std::cout,
        OptionFlags{
            {"showExecutionSteps", doProfilingExecution ? "true" : "false"},
            {"showVarStorage", doProfilingVars ? "true" : "false"}});
  }

  if (!checkResult) {
    return 0;
  }

  std::cout << "Checking the result..." << std::endl;
  boost::multi_array<float, 4> denseMatC(
      boost::extents[1][1][resRows0][resCols0]);
  boost::multi_array<float, 5> denseMatC3D(
      boost::extents[numGroups][1][1][resRows0][resCols0]);

  std::fill_n(denseMatC.data(), denseMatC.num_elements(), 0.0f);
  if (numGroups == 1) {
    if (!isResMatrixSparse0) {
      copy(target, dataType, outputRawHost.get(), denseMatC);
    } else {
      denseMatC.reshape(std::vector<int>(
          {resBlockRows0, resBlockCols0, resBlockRowSize0, resBlockColSize0}));
      boost::multi_array<float, 3> denseMatCTmp(
          boost::extents[nonZeroBlock][resBlockRowSize0][resBlockColSize0]);
      copy(target, dataType, outputRawHost.get(), denseMatCTmp);
      int blockCount = 0;
      for (int br = 0; br < resBlockRows0; ++br) {
        for (int bc = 0; bc < resBlockCols0; ++bc) {
          if (sparsity[br * resBlockCols0 + bc] > 0) {
            denseMatC[br][bc] = denseMatCTmp[blockCount++];
          }
        }
      }
    }
  } else {
    if (!isResMatrixSparse0) {
      copy(target, dataType, outputRawHost.get(), denseMatC3D);
    } else {
      denseMatC3D.reshape(std::vector<int>(
          {static_cast<int>(numGroups), resBlockRows0, resBlockCols0,
           resBlockRowSize0, resBlockColSize0}));
      const unsigned char *itemSparsity = sparsity.data();
      boost::multi_array<float, 3> denseMatCTmp(
          boost::extents[nonZeroBlock][resBlockRowSize0][resBlockColSize0]);
      copy(target, dataType, outputRawHost.get(), denseMatCTmp);
      int blockCount = 0;
      for (int g = 0; g < numGroups; ++g) {
        for (int br = 0; br < resBlockRows0; ++br) {
          for (int bc = 0; bc < resBlockCols0; ++bc) {
            if (itemSparsity[br * resBlockCols0 + bc] > 0) {
              denseMatC3D[g][br][bc] = denseMatCTmp[blockCount++];
            }
          }
        }
        itemSparsity += sparsityDenseSize;
      }
    }
  }

  std::vector<std::vector<float>> hostMatC(resRows0,
                                           std::vector<float>(resCols0, 0.0f));

  std::vector<std::vector<std::vector<float>>> hostMatC3D(
      numGroups, std::vector<std::vector<float>>(
                     resRows0, std::vector<float>(resCols0, 0.0f)));

  if (computeMatmul) {
    if (!isResMatrixSparse0) {
      for (int r = 0; r < resRows0; r++) {
        for (int c = 0; c < resCols0; c++) {
          if (numGroups == 1) {
            float sum = 0.0f;
            for (int cmul = 0; cmul < lhsCols0; cmul++) {
              sum +=
                  lhsHost[r][cmul] *
                  (rhsNeedTranspose0 ? (rhsHost[c][cmul]) : (rhsHost[cmul][c]));
            }
            hostMatC[r][c] = sum;
          } else {
            for (int g = 0; g < numGroups; ++g) {
              float sum = 0.0f;
              for (int cmul = 0; cmul < lhsCols0; cmul++) {
                sum += lhsHost3D[g][r][cmul] * (rhsNeedTranspose0
                                                    ? (rhsHost3D[g][c][cmul])
                                                    : (rhsHost3D[g][cmul][c]));
              }
              hostMatC3D[g][r][c] = sum;
            }
          }
        }
      }
    } else {
      for (int r = 0; r < resRows0; r++) {
        for (int c = 0; c < resCols0; c++) {
          if (numGroups == 1) {
            float sum = 0.0f;
            if (!((mask == SubBlockMask::ZeroUpperTriangle && r < c) ||
                  (mask == SubBlockMask::ZeroLowerTriangle && r > c))) {
              int br = r / lhsBlockRowSize0;
              int bc = c / rhsBlockSize0;
              if (sparsity[br * rhsBlockCols0 + bc] != 0) {
                for (int cmul = 0; cmul < lhsCols0; cmul++) {
                  sum += lhsHost[r][cmul] * (rhsNeedTranspose0
                                                 ? (rhsHost[c][cmul])
                                                 : (rhsHost[cmul][c]));
                }
              }
            }
            hostMatC[r][c] = sum;
          } else {
            const unsigned char *itemSparsity = sparsity.data();
            for (int g = 0; g < numGroups; ++g) {
              float sum = 0.0f;
              if (!((mask == SubBlockMask::ZeroUpperTriangle && r < c) ||
                    (mask == SubBlockMask::ZeroLowerTriangle && r > c))) {
                int br = r / lhsBlockRowSize0;
                int bc = c / rhsBlockSize0;
                if (itemSparsity[br * rhsBlockCols0 + bc] != 0) {
                  for (int cmul = 0; cmul < lhsCols0; cmul++) {
                    sum += lhsHost3D[g][r][cmul] *
                           (rhsNeedTranspose0 ? (rhsHost3D[g][c][cmul])
                                              : (rhsHost3D[g][cmul][c]));
                  }
                }
              }
              hostMatC3D[g][r][c] = sum;
              itemSparsity += sparsityDenseSize;
            }
          }
        }
      }
    }
  } else {
    for (int r = 0; r < resRows0; r++) {
      for (int c = 0; c < resCols0; c++) {
        hostMatC[r][c] = inputHost[r][c];
      }
    }
  }

  if (computeSoftmax) {
    if (hasSparsity) {
      std::vector<float> maxs(resRows0, -3.4028235e+38);
      for (int r = 0; r < resRows0; r++) {
        for (int c = 0; c < resCols0; c++) {
          if (!((mask == SubBlockMask::ZeroUpperTriangle && r < c) ||
                (mask == SubBlockMask::ZeroLowerTriangle && r > c))) {
            int br = r / resBlockRowSize0;
            int bc = c / resBlockColSize0;
            if (sparsity[br * resBlockCols0 + bc] != 0) {
              maxs[r] = std::max(maxs[r], hostMatC[r][c]);
            }
          }
        }
      }
      for (int r = 0; r < resRows0; r++) {
        std::vector<double> expVals(resCols0, 0.0f);
        double expSum = 0.0f;
        for (int c = 0; c < resCols0; c++) {
          if (!((mask == SubBlockMask::ZeroUpperTriangle && r < c) ||
                (mask == SubBlockMask::ZeroLowerTriangle && r > c))) {
            int br = r / resBlockRowSize0;
            int bc = c / resBlockColSize0;
            if (sparsity[br * resBlockCols0 + bc] != 0) {
              double expV =
                  std::exp(static_cast<double>(hostMatC[r][c] - maxs[r]));
              expVals[c] = expV;
              expSum += expV;
            }
          }
        }
        for (int c = 0; c < resCols0; c++) {
          hostMatC[r][c] = 0.0f;
          if (!((mask == SubBlockMask::ZeroUpperTriangle && r < c) ||
                (mask == SubBlockMask::ZeroLowerTriangle && r > c))) {
            int br = r / resBlockRowSize0;
            int bc = c / resBlockColSize0;
            if (sparsity[br * resBlockCols0 + bc] != 0) {
              hostMatC[r][c] = expVals[c] / expSum;
            }
          }
        }
      }
    } else {
      std::vector<float> maxs(resRows0, -3.4028235e+38);
      for (int r = 0; r < resRows0; r++) {
        for (int c = 0; c < resCols0; c++) {
          maxs[r] = std::max(maxs[r], hostMatC[r][c]);
        }
      }
      for (int r = 0; r < resRows0; r++) {
        std::vector<double> expVals(resCols0, 0.0f);
        double expSum = 0.0f;
        for (int c = 0; c < resCols0; c++) {
          double expV = std::exp(static_cast<double>(hostMatC[r][c] - maxs[r]));
          expVals[c] = expV;
          expSum += expV;
        }
        for (int c = 0; c < resCols0; c++) {
          hostMatC[r][c] = expVals[c] / expSum;
        }
      }
    }
  }

  const float threshold = dataType == poplar::FLOAT ? 1.0e-4 : 1.0e-2;
  float maxErr = 0.0f;
  int errRow = 0;
  int errCol = 0;
  float valTruth = 0.0f;
  float valTest = 0.0f;
  bool success = true;
  bool foundNan = false;
  for (int g = 0; g < numGroups; ++g && !foundNan) {
    for (int r = 0; r < resRows0; r++ && !foundNan) {
      for (int c = 0; c < resCols0; c++ && !foundNan) {
        float err;
        float curValTruth, curValTest;
        if (numGroups == 1) {
          curValTruth = hostMatC[r][c];
        } else {
          curValTruth = hostMatC3D[g][r][c];
        }
        if (!isResMatrixSparse0) {
          if (numGroups == 1) {
            curValTest = denseMatC[0][0][r][c];
          } else {
            curValTest = denseMatC3D[g][0][0][r][c];
          }
          err = fabs(curValTest - curValTruth);
        } else {
          int br = r / resBlockRowSize0;
          int bc = c / resBlockColSize0;
          int rb = r % resBlockRowSize0;
          int cb = c % resBlockColSize0;
          if (numGroups == 1) {
            curValTest = denseMatC[br][bc][rb][cb];
          } else {
            curValTest = denseMatC3D[g][br][bc][rb][cb];
          }
          err = fabs(curValTest - curValTruth);
        }
        if (std::isnan(curValTest)) {
          foundNan = true;
          success = false;
          logging::popsparse::debug("[{}][{}] True: {}, test = {}", r, c,
                                    curValTruth, curValTest);
        } else if (!checkIsClose(curValTruth, curValTest, threshold)) {
          success = false;
          logging::popsparse::debug("[{}][{}] True: {}, test = {}", r, c,
                                    curValTruth, curValTest);
          if (err > maxErr) {
            valTruth = curValTruth;
            valTest = curValTest;
            errRow = r;
            errCol = c;
            maxErr = err;
          }
        }
      }
    }
  }

  if (success) {
    std::cout << "Result check succeeded." << std::endl;
  } else {
    if (foundNan) {
      std::cout << "Result check failed." << std::endl
                << "Found NaN in the output" << std::endl;
    } else {
      std::cout << "Result check failed." << std::endl
                << " Maximum error = " << std::fixed << std::setprecision(5)
                << maxErr << std::endl
                << " test[" << errRow << "][" << errCol << "] = " << std::fixed
                << std::setprecision(5) << valTest << std::endl
                << " true[" << errRow << "][" << errCol << "] = " << std::fixed
                << std::setprecision(5) << valTruth << std::endl;
    }
  }

  return (success ? 0 : -1);
}


#pragma once

#include <deal.II/base/index_set.h>
#include <deal.II/base/utilities.h>

// own includes ----------------------------------------------------------
#include "matrix/dofs/dofindex_sets.hpp"
#include "spectral/basis/global_indexer.hpp"

namespace boltzmann {
class MomentBase
{
 public:
  template <typename VEC_IN, typename VEC_OUT, typename INDEXER>
  void compute(VEC_OUT& dst,
               const VEC_IN& src,
               const dealii::IndexSet& relevant_phys_dofs,
               const INDEXER& indexer) const;

 protected:
  typedef double value_t;
  typedef std::pair<unsigned int, value_t> entry_t;
  unsigned int n_velo_dofs;
  std::vector<entry_t> contributions;

 public:
  const std::vector<entry_t>& entries() const { return contributions; }
};

// ----------------------------------------------------------------------
template <typename VEC_IN, typename VEC_OUT, typename INDEXER>
void
MomentBase::compute(VEC_OUT& dst,
                    const VEC_IN& src,
                    const dealii::IndexSet& relevant_phys_dofs,
                    const INDEXER& indexer) const
{
#pragma omp parallel for
  for (unsigned int ix = 0; ix < relevant_phys_dofs.size(); ++ix) {
    if (relevant_phys_dofs.is_element(ix)) {
      double sum = 0;
      for (unsigned int j = 0; j < contributions.size(); ++j) {
        unsigned int jx = contributions[j].first;
        sum += contributions[j].second * src[indexer.to_global(ix, jx)];
      }
      dst[ix] = sum;
    }
  }
}

}  // end namespace boltzmann