// --------------------------------------------------------------------------
//                   OpenMS -- Open-Source Mass Spectrometry
// --------------------------------------------------------------------------
// Copyright The OpenMS Team -- Eberhard Karls University Tuebingen,
// ETH Zurich, and Freie Universitaet Berlin 2002-2018.
//
// This software is released under a three-clause BSD license:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of any author or any participating institution
//    may be used to endorse or promote products derived from this software
//    without specific prior written permission.
// For a full list of authors, refer to the file AUTHORS.
// --------------------------------------------------------------------------
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL ANY OF THE AUTHORS OR THE CONTRIBUTING
// INSTITUTIONS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// --------------------------------------------------------------------------
// $Maintainer: Timo Sachsenberg $
// $Authors: Stephan Aiche, Chris Bielow $
// --------------------------------------------------------------------------

#include <OpenMS/ANALYSIS/QUANTITATION/IsobaricIsotopeCorrector.h>
#include <OpenMS/ANALYSIS/QUANTITATION/IsobaricQuantitationMethod.h>
#include <OpenMS/ANALYSIS/QUANTITATION/IsobaricQuantifierStatistics.h>

#include <OpenMS/DATASTRUCTURES/Utils/MatrixUtils.h>
#include <OpenMS/KERNEL/ConsensusMap.h>

// NNLS isotope correction
#include <OpenMS/MATH/MISC/NonNegativeLeastSquaresSolver.h>

#include <Eigen/LU>

// #define ISOBARIC_QUANT_DEBUG

namespace OpenMS
{

  IsobaricQuantifierStatistics
  IsobaricIsotopeCorrector::correctIsotopicImpurities(
    const ConsensusMap& consensus_map_in, ConsensusMap& consensus_map_out,
    const IsobaricQuantitationMethod* quant_method)
  {
    OPENMS_PRECONDITION(consensus_map_in.size() == consensus_map_out.size(),
                        "The in- and output map need to have the same size.")

    // the stats object to fill while correcting
    IsobaricQuantifierStatistics stats;
    stats.number_ms2_total = consensus_map_out.size();
    stats.channel_count = quant_method->getNumberOfChannels();

    Matrix<double> correction_matrix = quant_method->getIsotopeCorrectionMatrix();

    if (matrixIsIdentityMatrix(correction_matrix))
    {
      throw Exception::InvalidParameter(__FILE__, __LINE__, OPENMS_PRETTY_FUNCTION,
                                        "IsobaricIsotopeCorrector: The given isotope correction matrix is an identity matrix leading to no correction. "
                                        "Please provide a valid isotope_correction matrix as it was provided with the sample kit!");
    }

    // convert to Eigen matrix
    EigenMatrixXdPtr m(convertOpenMSMatrix2EigenMatrixXd(correction_matrix));
    Eigen::FullPivLU<Eigen::MatrixXd> ludecomp(*m);
    Eigen::VectorXd b;
    b.resize(quant_method->getNumberOfChannels());
    b.setZero();
    std::vector<double> x(quant_method->getNumberOfChannels(), 0);

    if (!ludecomp.isInvertible())
    {
      // clean up before we leave
      throw Exception::InvalidParameter(__FILE__, __LINE__, OPENMS_PRETTY_FUNCTION, "IsobaricIsotopeCorrector: The given isotope correction matrix is not invertible!");
    }

    // data structures for NNLS
    Matrix<double> m_b(quant_method->getNumberOfChannels(), 1);
    Matrix<double> m_x(quant_method->getNumberOfChannels(), 1);

    // correct all consensus elements
    for (ConsensusMap::size_type i = 0; i < consensus_map_out.size(); ++i)
    {
#ifdef ISOBARIC_QUANT_DEBUG
      std::cout << "\nMAP element  #### " << i << " #### \n" << std::endl;
#endif
      // delete only the consensus handles from the output map
      consensus_map_out[i].clear();

      // fill b vector
      fillInputVector_(b, m_b, consensus_map_in[i], consensus_map_in);

      //solve
      Eigen::MatrixXd e_mx = ludecomp.solve(b);
      if (!((*m) * e_mx).isApprox(b))
      {
        throw Exception::InvalidParameter(__FILE__, __LINE__, OPENMS_PRETTY_FUNCTION, "IsobaricIsotopeCorrector: Cannot multiply!");
      }
      solveNNLS_(correction_matrix, m_b, m_x);

      // update the output consensus map with the corrected intensities
      float cf_intensity = updateOutpuMap_(consensus_map_in, consensus_map_out, i, m_x);

      // check consistency
      computeStats_(m_x, e_mx, cf_intensity, quant_method, stats);
    }

    return stats;
  }

  void
  IsobaricIsotopeCorrector::fillInputVector_(Eigen::VectorXd& b,
                                             Matrix<double>& m_b, const ConsensusFeature& cf, const ConsensusMap& cm)
  {
    for (ConsensusFeature::HandleSetType::const_iterator it_elements = cf.getFeatures().begin();
         it_elements != cf.getFeatures().end();
         ++it_elements)
    {
      //find channel_id of current element
      Int index = Int(cm.getColumnHeaders().find(it_elements->getMapIndex())->second.getMetaValue("channel_id"));
#ifdef ISOBARIC_QUANT_DEBUG
      std::cout << "  map_index " << it_elements->getMapIndex() << "-> id " << index << " with intensity " << it_elements->getIntensity() << "\n" << std::endl;
#endif
      // this is deprecated, but serves as quality measurement
      b(index) = it_elements->getIntensity();
      m_b(index, 0) = it_elements->getIntensity();
    }
  }

  void
  IsobaricIsotopeCorrector::solveNNLS_(const Matrix<double>& correction_matrix,
                                       const Matrix<double>& m_b, Matrix<double>& m_x)
  {
    Int status = NonNegativeLeastSquaresSolver::solve(correction_matrix, m_b, m_x);
    if (status != NonNegativeLeastSquaresSolver::SOLVED)
    {
      throw Exception::FailedAPICall(__FILE__, __LINE__, OPENMS_PRETTY_FUNCTION, "IsobaricIsotopeCorrector: Failed to find least-squares fit!");
    }
  }

  void
  IsobaricIsotopeCorrector::computeStats_(const Matrix<double>& m_x,
                                          const Eigen::MatrixXd& x, const float cf_intensity,
                                          const IsobaricQuantitationMethod* quant_method, IsobaricQuantifierStatistics& stats)
  {
    Size s_negative(0);
    Size s_different_count(0); // happens when naive solution is negative in other channels
    double s_different_intensity(0);

    // ISOTOPE CORRECTION: compare solutions of Matrix inversion vs. NNLS
    for (Size index = 0; index < quant_method->getNumberOfChannels(); ++index)
    {
      if (x(index) < 0.0)
      {
        ++s_negative;
      }
      else if ((((std::fabs(m_x(index, 0) - x(index)))/m_x(index, 0))*100) > 1)
      {
        ++s_different_count;
        s_different_intensity += std::fabs(m_x(index, 0) - x(index));
      }
    }

    if (s_negative == 0 && s_different_count > 0) //some solutions are inconsistent, despite being positive
    {
      OPENMS_LOG_WARN << "IsobaricIsotopeCorrector: Isotope correction values of alternative method differ!" << std::endl;
    }

    // update global stats
    stats.iso_number_reporter_negative += s_negative;
    stats.iso_number_reporter_different += s_different_count;
    stats.iso_solution_different_intensity += s_different_intensity;

    if (s_negative > 0)
    {
      ++stats.iso_number_ms2_negative;
      stats.iso_total_intensity_negative += cf_intensity;
    }
  }

  float
  IsobaricIsotopeCorrector::updateOutpuMap_(
    const ConsensusMap& consensus_map_in, ConsensusMap& consensus_map_out,
    ConsensusMap::size_type current_cf, const Matrix<double>& m_x)
  {
    float cf_intensity(0);
    for (ConsensusFeature::HandleSetType::const_iterator it_elements = consensus_map_in[current_cf].begin();
         it_elements != consensus_map_in[current_cf].end();
         ++it_elements)
    {
      FeatureHandle handle = *it_elements;
      //find channel_id of current element
      Int index = Int(consensus_map_out.getColumnHeaders()[it_elements->getMapIndex()].getMetaValue("channel_id"));
      handle.setIntensity(float(m_x(index, 0)));

      consensus_map_out[current_cf].insert(handle);
      cf_intensity += handle.getIntensity(); // sum up all channels for CF

#ifdef ISOBARIC_QUANT_DEBUG
      std::cout <<  it_elements->getIntensity() << " -> " << handle.getIntensity() << std::endl;
#endif
    }
    consensus_map_out[current_cf].setIntensity(cf_intensity); // set overall intensity of CF (sum of all channels)

    return cf_intensity;
  }

} // namespace


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

#ifndef METRO_LIKELIHOOD_INDEPENDENTOBSERVATION_LOGLIKELIHOOD_HPP
#define METRO_LIKELIHOOD_INDEPENDENTOBSERVATION_LOGLIKELIHOOD_HPP

#include <memory>
#include <boost/noncopyable.hpp>
#include <Eigen/Core>
#include "metro/LogLikelihood.hpp"

namespace metro {
	template< typename Scalar, typename Vector, typename Matrix >
	struct IndependentObservationLogLikelihood: public LogLikelihood< Scalar, Vector, Matrix > {
		typedef std::auto_ptr< IndependentObservationLogLikelihood > UniquePtr ;
		typedef typename Eigen::Ref< Matrix > MatrixRef ;

		// A function which returns the underlying terms of the log-likelihood.
		virtual void get_terms_of_function( MatrixRef result ) const = 0 ;

		virtual void set_data( Matrix const& ) = 0 ;
	} ;
}

#endif


//==================================================================================================
/**
  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_X86_XOP_SIMD_FUNCTION_IS_LESS_HPP_INCLUDED
#define BOOST_SIMD_ARCH_X86_XOP_SIMD_FUNCTION_IS_LESS_HPP_INCLUDED

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/meta/as_logical.hpp>

#if !defined(_MM_PCOMCTRL_LT)
#define _MM_PCOMCTRL_LT 0
#define BOOST_SIMD_MISSING_MM_PCOMCTRL_LT
#endif

#if BOOST_HW_SIMD_X86_AMD_XOP
namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::int8_<A0>, bs::sse_>
                          , bs::pack_<bd::int8_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epi8(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epi8(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::int16_<A0>, bs::sse_>
                          , bs::pack_<bd::int16_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epi16(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epi16(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::int32_<A0>, bs::sse_>
                          , bs::pack_<bd::int32_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epi32(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epi32(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::int64_<A0>, bs::sse_>
                          , bs::pack_<bd::int64_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epi64(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epi64(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::uint8_<A0>, bs::sse_>
                          , bs::pack_<bd::uint8_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epu8(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epu8(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::uint16_<A0>, bs::sse_>
                          , bs::pack_<bd::uint16_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epu16(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epu16(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::uint32_<A0>, bs::sse_>
                          , bs::pack_<bd::uint32_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epu32(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epu32(a0,a1);
      #endif
    }
  };

  BOOST_DISPATCH_OVERLOAD ( is_less_
                          , (typename A0)
                          , bs::avx_
                          , bs::pack_<bd::uint64_<A0>, bs::sse_>
                          , bs::pack_<bd::uint64_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE
    bs::as_logical_t<A0> operator()(const A0& a0, const A0& a1) const BOOST_NOEXCEPT
    {
      #if defined(__clang__)
      return _mm_com_epu64(a0,a1,_MM_PCOMCTRL_LT);
      #else
      return _mm_comlt_epu64(a0,a1);
      #endif
    }
  };
} } }
#endif

#if defined(BOOST_SIMD_MISSING_MM_PCOMCTRL_LT)
#undef _MM_PCOMCTRL_LT
#undef BOOST_SIMD_MISSING_MM_PCOMCTRL_LT
#endif

#endif


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_SIMD_SIGN_INCLUDED
#define BOOST_SIMD_FUNCTION_SIMD_SIGN_INCLUDED

#include <boost/simd/function/scalar/sign.hpp>
#include <boost/simd/arch/common/generic/function/autodispatcher.hpp>
#include <boost/simd/arch/common/simd/function/sign.hpp>

#if defined(BOOST_HW_SIMD_X86)

#  if BOOST_HW_SIMD_X86 >= BOOST_HW_SIMD_X86_SSSE3_VERSION
#    include <boost/simd/arch/x86/ssse3/simd/function/sign.hpp>
#  endif

#endif

#endif


// Copyright (c) 2016 Graphcore Ltd. All rights reserved.
#include "poputil/exceptions.hpp"
#include <algorithm>
#include <boost/multi_array.hpp>
#include <boost/optional.hpp>
#include <boost/optional/optional_io.hpp>
#include <boost/program_options.hpp>
#include <boost/random.hpp>
#include <chrono>
#include <exception>
#include <fstream>
#include <istream>
#include <ostream>
#include <poplar/CycleCount.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/IPUModel.hpp>
#include <poplibs_support/Algorithm.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/GeneralMatrixMultiply.hpp>
#include <poplibs_test/Pass.hpp>
#include <poplibs_test/SparseMatrix.hpp>
#include <poplibs_test/Util.hpp>
#include <poplin/codelets.hpp>
#include <popops/codelets.hpp>
#include <popsparse/SparsePartitioner.hpp>
#include <popsparse/codelets.hpp>
#include <queue>
#include <random>
#include <ratio>

#include "../lib/popsparse/FullyConnectedOnTile.hpp"
#include "../lib/popsparse/FullyConnectedOptions.hpp"
#include "../lib/popsparse/FullyConnectedPlan.hpp"
#include "../lib/popsparse/SparsePartitionerImpl.hpp"
#include "../lib/popsparse/SparseStorageInternal.hpp"
#include "poplibs_support/VectorUtils.hpp"
#include "poplibs_support/logging.hpp"
#include "popsparse/FullyConnected.hpp"
#include "popsparse/FullyConnectedParams.hpp"

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

using namespace popsparse;
using namespace popsparse::dynamic;

// Tolerances used when data is not ignored
#define FLOAT_REL_TOL 0.01
#define HALF_REL_TOL 0.1

// TODO: This could be a program option if at all needed
using EType = float;

template <typename T>
static void logBucketStatistics(std::vector<PNBucket> &buckets,
                                const CSRMatrix<T> &csrMatrix) {
  if (buckets.empty()) {
    std::cerr << "   - No buckets found"
              << "\n";
    return;
  }

  std::size_t pnUsed = buckets.size();

  std::size_t maxNzElements = 0, maxMetaInfo = 0;
  std::size_t totalNzElements = 0, totalMetaInfo = 0;
  std::for_each(buckets.begin(), buckets.end(), [&](const PNBucket &b) {
    maxNzElements = std::max(maxNzElements, b.numNzElements);
    maxMetaInfo = std::max(maxMetaInfo, b.metaInfoElements);
    totalNzElements += b.numNzElements;
    totalMetaInfo += b.metaInfoElements;
  });

  std::cerr << "   - NZ entries " << csrMatrix.nzValues.size() << "\n";
  std::cerr << "   - NZ elements/PN max : " << maxNzElements;
  std::cerr << " avg : " << static_cast<double>(totalNzElements) / pnUsed;
  std::cerr << "\n";
  std::cerr << "   - Meta info elements/PN max : " << maxMetaInfo;
  std::cerr << " avg : " << static_cast<double>(totalMetaInfo) / pnUsed;
  std::cerr << "\n";
}

int main(int argc, char **argv) try {
  namespace po = boost::program_options;

  DeviceType deviceType = DeviceType::IpuModel2;
  unsigned numGroups = 1;
  unsigned inputSize;
  unsigned outputSize;
  unsigned batchSize;
  bool reportPlan;
  boost::optional<std::string> profileDir;
  Type dataType;
  Type partialsType;
  unsigned numIPUs = 1;
  boost::optional<unsigned> tilesPerIPU;
  Pass pass = Pass::ALL;
  std::string matmulOptionsString;
  std::string planConstraintsString;
  double sparsityFactor;
  ShapeOption<std::size_t> weightedAreaBegin;
  ShapeOption<std::size_t> weightedAreaEnd;
  ShapeOption<std::size_t> blockSize;
  weightedAreaBegin.val = weightedAreaEnd.val = {0, 0};
  double weightedAreaWeighting = 1.0;
  bool denseGradWSerialSplits = false;

  po::options_description desc("Options");
  // clang-format off
  desc.add_options()
    ("help", "Produce help message")
    ("compile-only", "Stop after compilation; don't run the program")
    ("device-type",
     po::value<DeviceType>(&deviceType)->default_value(deviceType),
     deviceTypeHelp)
    ("input-size", po::value<unsigned>(&inputSize)->required(),
     "Number of inputs")
    ("output-size", po::value<unsigned>(&outputSize)->required(),
     "Number of output channels")
    ("sparsity-factor", po::value<double>(&sparsityFactor)->required(),
     "Sparsity factor (ratio of number of non-zero values to total weight "
     "values")
    ("data-type",
     po::value<Type>(&dataType)->default_value(HALF),
     "Type of the input and output data")
    ("partials-type",
     po::value<Type>(&partialsType)->default_value(FLOAT),
     "Type of partials used during the operation")
    ("tiles-per-ipu",
     po::value(&tilesPerIPU),
     "Number of tiles per IPU")
    ("batch-size",
     po::value<unsigned>(&batchSize)->default_value(1),
     "Batch size")
    ("block-size",
     po::value<ShapeOption<std::size_t>>(&blockSize)->default_value(1),
     "Block size as rows and columns (only square blocks are supported)")
    ("single-phase",
     po::value<Pass>(&pass)->default_value(pass),
     "Run phase all | fwd | bwd | wu")
    ("ignore-data", "When set, no upload/download or verification of "
     "results is performed")
    ("plan-only", "Whether to perform planning only and skip creation "
     "and running of the program")
    ("profile", "Enable profiling and print profiling report")
    ("profile-dir",
     po::value<decltype(profileDir)>(&profileDir)
      ->default_value(boost::none),
     "Write profile files to the specified directory.")
    ("report-plan", po::value<bool>(&reportPlan)->default_value(false),
     "Display plan")
    ("report-total-cycle-counts", "Report total cycle count ignoring upload/download for "
     "each pass. Note not compatible with 'profile' option")
    ("variable-seed", "Use a variable seed based on clock, rather than a "
     "single fixed seed that does not change between runs of this tool")
    ("weighted-area-begin",
     po::value<ShapeOption<std::size_t>>(&weightedAreaBegin)->default_value(weightedAreaBegin),
     "Starting indices of an area of the sparse operand with a different "
     "level of sparsity to the rest")
    ("weighted-area-end",
     po::value<ShapeOption<std::size_t>>(&weightedAreaEnd)->default_value(weightedAreaEnd),
     "Ending indices of an area of the sparse operand with a different "
     "level of sparsity to the rest")
    ("weighted-area-weighting",
     po::value<double>(&weightedAreaWeighting)->default_value(weightedAreaWeighting),
     "Weighting for probability that a sparse element resides within the "
     "specified area")
    ("matmul-options", po::value<std::string>(&matmulOptionsString),
     "Options to use for the matrix multiplication, specified as a JSON "
     "string, e.g. {\"key\":\"value\"}")
    ("plan-constraints", po::value<std::string>(&planConstraintsString),
     "Plan constraints to use for the matrix multiplication, specified as "
     "a JSON string")
    ("report-dense-gradw-serial-splits",
      po::value<bool>(&denseGradWSerialSplits)->
        default_value(denseGradWSerialSplits),
     "Report dense GradW splits when GradW pass is enabled")
  ;
  // clang-format on
  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, desc), vm);
    if (vm.count("help")) {
      std::cout << desc << "\n";
      return 1;
    }
    po::notify(vm);
  } catch (std::exception &e) {
    std::cerr << "error: " << e.what() << "\n";
    return 1;
  }

  bool profile = vm.count("profile");
  bool profilingEnabled = profile || profileDir;
  bool reportTotalCycleCounts =
      vm.count("report-total-cycle-counts") && deviceType == DeviceType::Hw;
  bool ignoreData = vm.count("ignore-data");
  bool planOnly = vm.count("plan-only");
  bool variableSeed = vm.count("variable-seed");

  if (reportTotalCycleCounts && profilingEnabled) {
    throw poputil::poplibs_error(
        "--report-total-cycle-counts and --profile or --profile-dr specified "
        "at the same time. This is not allowed as one affects the other");
  }

  if (blockSize.val.size() > 2) {
    throw poputil::poplibs_error("Block size must be of dimension 2");
  }

  const std::size_t blockRows = blockSize[0];
  const std::size_t blockCols =
      blockSize.val.size() == 1 ? blockRows : blockSize[1];
  const auto blockArea = blockRows * blockCols;

  if (inputSize % blockRows) {
    throw poputil::poplibs_error("Input size must be an integer multiple of "
                                 "rows in a block");
  }

  if (outputSize % blockCols) {
    throw poputil::poplibs_error("output size must be an integer multiple of "
                                 "columns in a block");
  }

  if (weightedAreaBegin[0] > weightedAreaEnd[0] ||
      weightedAreaBegin[1] > weightedAreaEnd[1]) {
    std::stringstream ss;
    ss << "Invalid weighted area specified: " << weightedAreaBegin.val << ","
       << weightedAreaEnd.val;
    throw poputil::poplibs_error(ss.str());
  }

  if (weightedAreaEnd[0] > outputSize || weightedAreaEnd[1] > inputSize) {
    std::stringstream ss;
    ss << "Specified weighted area is out of bounds: Weighted area="
       << weightedAreaBegin.val << "," << weightedAreaEnd.val
       << " out of bounds {" << outputSize << "," << inputSize << "}";
    throw poputil::poplibs_error(ss.str());
  }

  // align weighted area to a block size grid
  weightedAreaBegin.val[0] = roundDown(weightedAreaBegin.val[0], blockRows);
  weightedAreaBegin.val[1] = roundDown(weightedAreaBegin.val[1], blockCols);
  weightedAreaEnd.val[0] = roundDown(weightedAreaEnd.val[0], blockRows);
  weightedAreaEnd.val[1] = roundDown(weightedAreaEnd.val[1], blockCols);

  PlanningCache cache;

  poplar::OptionFlags options;
  bool doBwdPass = pass == Pass::BWD || pass == Pass::ALL;
  bool doWuPass = pass == Pass::WU || pass == Pass::ALL;
  options.set("availableMemoryProportion", "1.0");
  options.set("doGradAPass", doBwdPass ? "true" : "false");
  options.set("doGradWPass", doWuPass ? "true" : "false");
  options.set("partialsType", partialsType.toString());

  // User options specified via --matmul-options override defaults
  if (!matmulOptionsString.empty()) {
    poplar::readJSON(matmulOptionsString, options);
  }
  if (!planConstraintsString.empty()) {
    options.set("planConstraints", planConstraintsString);
  }

  auto device = tilesPerIPU
                    ? createTestDevice(deviceType, numIPUs, *tilesPerIPU, true)
                    : createTestDeviceFullSize(deviceType, numIPUs, true);
  const auto &target = device.getTarget();

  const auto sparsityType =
      blockArea == 1 ? SparsityType::Element : SparsityType::Block;

  SparsityParams sparsityParams(sparsityType, SparsityStructure::Unstructured,
                                {blockRows, blockCols});

  const auto params = FullyConnectedParams::createWithNzRatio(
      std::move(sparsityParams), sparsityFactor, batchSize, numGroups,
      inputSize, outputSize);

  popsparse::fullyconnected::Plan plan;
  popsparse::fullyconnected::Cost planCost;

  // Always do forward
  std::tie(plan, planCost) = popsparse::fullyconnected::getPlan(
      target, dataType, params, options, &cache);

  if (reportPlan) {
    std::string str;
    if (doBwdPass || doWuPass) {
      str += "Joint : ";
      if (doBwdPass) {
        str += " GradA + ";
      }
      if (doWuPass) {
        str += " GradW + ";
      }
    }
    str += "Fwd Plan \n";
    std::cerr << str << plan << "\n" << planCost << "\n";
  }

  std::size_t fwdMetaInfoBucketSize = plan.fwdMetaInfoElemsPerBucket;
  std::size_t gradAMetaInfoBucketSize = plan.gradAMetaInfoElemsPerBucket;
  std::size_t nzElementBucketSize = plan.nzElemsPerBucket;

  std::cerr << "Using bucket sizes:"
            << "\n  meta info (forward): " << fwdMetaInfoBucketSize
            << "\n  meta info (grad-a): " << gradAMetaInfoBucketSize
            << "\n  nz element : " << nzElementBucketSize << "\n";

  Partitioner<EType> partitioner(params, dataType, target, options, &cache);
  std::mt19937 randomEngine;
  if (variableSeed) {
    using namespace std::chrono;
    using SeedDurationType = duration<std::mt19937::result_type, std::nano>;
    const auto now = high_resolution_clock::now();
    const auto seed =
        duration_cast<SeedDurationType>(now.time_since_epoch()).count();
    randomEngine.seed(seed);
  }

  const bool floatingPointCouldRepresentMaxAccum = [&] {
    const auto maxVal = maxContiguousInteger(dataType);

    double weightedThreshold, remainingThreshold;
    std::tie(weightedThreshold, remainingThreshold) =
        poplibs_test::sparse::calculateWeightedVsRemainingSparsityFactor(
            {outputSize / blockRows, inputSize / blockCols}, sparsityFactor,
            {weightedAreaBegin[0] / blockRows,
             weightedAreaBegin[1] / blockCols},
            {weightedAreaEnd[0] / blockRows, weightedAreaEnd[1] / blockCols},
            weightedAreaWeighting);
    const auto numWeightedInputChannels =
        (weightedAreaEnd[1] - weightedAreaBegin[1]);
    const auto numWeightedOutputChannels =
        (weightedAreaEnd[0] - weightedAreaBegin[0]);
    std::size_t maxInputChannels =
        numWeightedInputChannels * weightedThreshold +
        (params.getInputChannelsPerGroup() - numWeightedInputChannels) *
            remainingThreshold;
    std::size_t maxOutputChannels =
        numWeightedOutputChannels * weightedThreshold +
        (params.getOutputChannelsPerGroup() - numWeightedOutputChannels) *
            remainingThreshold;
    maxInputChannels = roundDown(maxInputChannels, blockCols);
    maxOutputChannels = roundDown(maxOutputChannels, blockRows);
    const auto getOpsPerOutputElementEstimate = [&](const Pass &pass) -> int {
      const auto numAccumulations = pass == Pass::FWD   ? maxInputChannels
                                    : pass == Pass::BWD ? maxOutputChannels
                                                        : params.getBatchSize();
      return numAccumulations;
    };
    // We use a modifier to account for the unlikeliness of picking all positive
    // or negative 1s which would actually get us to the max precisely
    // represented integer.
    constexpr int modifier = 10;
    // We use another modifier to account for the chance that sparsity is not
    // perfectly evenly spread in this instant.
    constexpr double wiggleRoom = 1.3;
    if (wiggleRoom * getOpsPerOutputElementEstimate(Pass::FWD) >
        maxVal * modifier) {
      return false;
    }
    if (doBwdPass && wiggleRoom * getOpsPerOutputElementEstimate(Pass::BWD) >
                         maxVal * modifier) {
      return false;
    }
    if (doWuPass && wiggleRoom * getOpsPerOutputElementEstimate(Pass::WU) >
                        maxVal * modifier) {
      return false;
    }
    return true;
  }();

  // create CSR matrix for the given sparsity factor
  const bool useBipolarDistribution = floatingPointCouldRepresentMaxAccum;
  std::array<std::size_t, 2> blockDims = {blockRows, blockCols};
  CSRMatrix<EType> csrMatrix(blockDims);
  std::tie(csrMatrix.nzValues, csrMatrix.columnIndices, csrMatrix.rowIndices) =
      poplibs_test::sparse::buildCSRMatrix<EType, std::size_t>(
          randomEngine, {outputSize, inputSize}, {blockRows, blockCols},
          sparsityFactor, weightedAreaBegin, weightedAreaEnd,
          weightedAreaWeighting, useBipolarDistribution);

  // Forward
  boost::multi_array<double, 2> hostInput(boost::extents[batchSize][inputSize]);
  if (useBipolarDistribution) {
    writeRandomBinaryValues(target, dataType, hostInput, -1.0, 1.0,
                            randomEngine);
  } else {
    writeRandomValues(target, dataType, hostInput, -3.0, 3.0, randomEngine);
  }
  boost::multi_array<double, 2> hostOutputActs(
      boost::extents[batchSize][outputSize]);

  // GradA
  boost::multi_array<double, 2> hostOutputGrad(
      boost::extents[batchSize][outputSize]);
  if (useBipolarDistribution) {
    writeRandomBinaryValues(target, dataType, hostOutputGrad, -1.0, 1.0,
                            randomEngine);
  } else {
    writeRandomValues(target, dataType, hostOutputGrad, -3.0, 3.0,
                      randomEngine);
  }
  boost::multi_array<double, 2> hostInputGrad(
      boost::extents[batchSize][inputSize]);

  if (!plan.useDense) {
    std::cerr << "Logging Forward pass bucket statistics: ";
    auto pnBucketsImpl = partitioner.getImpl().createBuckets(csrMatrix);

    logBucketStatistics(pnBucketsImpl.pnBuckets, csrMatrix);
  }

  if (planOnly) {
    return 0;
  }

  Graph graph(target);
  popops::addCodelets(graph);
  poplin::addCodelets(graph);
  popsparse::addCodelets(graph);
  Sequence fwdProg, bwdProg, wuProg, uploadProg, downloadProg;

  // Build the graph
  std::cerr << "Constructing graph...\n";
  const SparseTensor weights = createFullyConnectedWeights(
      graph, dataType, params, "weights", options, &cache);
  const Tensor input = createFullyConnectedInput(graph, dataType, params,
                                                 "input", options, &cache);
  const Tensor outputActs = fullyConnectedFwd(graph, weights, input, params,
                                              fwdProg, "fwd", options, &cache);

  // GradW
  boost::multi_array<double, 1> hostWeightGrad(
      boost::extents[weights.getNzValuesTensor().numElements()]);

  Tensor outputGrad;
  if (doBwdPass || doWuPass) {
    outputGrad = graph.clone(outputActs, "outputGrad");
  }
  Tensor inputGrad;
  if (doBwdPass) {
    inputGrad = fullyConnectedGradA(graph, weights, outputGrad, params, bwdProg,
                                    "grada", options, &cache);
  }

  Tensor weightGrad;
  if (doWuPass) {
    weightGrad = fullyConnectedSparseGradW(graph, weights.getMetaInfoTensor(),
                                           outputGrad, input, params, wuProg,
                                           "wu", options, &cache);
  }
  std::cerr << "Done\n";

  std::vector<std::pair<std::string, char *>> tmap;
  auto rawMetaInfo =
      allocateHostMemoryForTensor(weights.getMetaInfoTensor(), "weights.meta",
                                  graph, uploadProg, downloadProg, tmap);

  auto rawNzInfo =
      allocateHostMemoryForTensor(weights.getNzValuesTensor(), "weights.nz",
                                  graph, uploadProg, downloadProg, tmap);
  auto rawInput = allocateHostMemoryForTensor(input, "input", graph, uploadProg,
                                              downloadProg, tmap);
  auto rawOutputActs = allocateHostMemoryForTensor(
      outputActs, "outputActs", graph, uploadProg, downloadProg, tmap);

  std::unique_ptr<char[]> rawOutputGrad;
  if (doBwdPass || doWuPass) {
    rawOutputGrad = allocateHostMemoryForTensor(outputGrad, "outputGrad", graph,
                                                uploadProg, downloadProg, tmap);
  }

  std::unique_ptr<char[]> rawInputGrad;
  if (!ignoreData && doBwdPass) {
    rawInputGrad = allocateHostMemoryForTensor(inputGrad, "inputGrad", graph,
                                               uploadProg, downloadProg, tmap);
  }

  std::unique_ptr<char[]> rawWeightGrad;
  if (!ignoreData && doWuPass) {
    rawWeightGrad = allocateHostMemoryForTensor(weightGrad, "weightGrad", graph,
                                                uploadProg, downloadProg, tmap);
  }

  Tensor fwdCycles, bwdCycles, wuCycles;
  if (reportTotalCycleCounts) {
    fwdCycles = cycleCount(graph, fwdProg, 0, SyncType::INTERNAL, "fwdCycles");
    graph.createHostRead("fwdCycles", fwdCycles);
    if (doBwdPass) {
      bwdCycles =
          cycleCount(graph, bwdProg, 0, SyncType::INTERNAL, "bwdCycles");
      graph.createHostRead("bwdCycles", bwdCycles);
    }
    if (doWuPass) {
      wuCycles = cycleCount(graph, wuProg, 0, SyncType::INTERNAL, "wuCycles");
      graph.createHostRead("wuCycles", wuCycles);
    }
  }
  Sequence controlProg({std::move(uploadProg), std::move(fwdProg),
                        std::move(bwdProg), std::move(wuProg)});
  if (!ignoreData) {
    controlProg.add(std::move(downloadProg));
  }

  std::cerr << "Creating engine...\n";
  OptionFlags engineOptions;
  if (profilingEnabled) {
    engineOptions.set("debug.instrument", "true");
    if (profileDir) {
      engineOptions.set("autoReport.all", "true");
      engineOptions.set("autoReport.directory", *profileDir);
    }
  }
  Engine engine(graph, std::move(controlProg), engineOptions);

  if (vm.count("compile-only"))
    return 0;

  attachStreams(engine, tmap);

  std::cerr << "Done\n";

  std::cerr << "Running...\n";

  // Actual bucket info use by device graph
  auto buckets = partitioner.createSparsityDataImpl(csrMatrix);

  const auto &metaInfoFlat = buckets.metaInfo;
  const auto &nzValuesFlat = buckets.nzValues;
  // Overflow info is the same for all passes at time of writing.
  if (!plan.useDense) {
    std::cerr << "overflowInfo = {" << metaInfoFlat.at(0) << ","
              << metaInfoFlat.at(1) << "," << metaInfoFlat.at(2) << "}\n";
  }

  copy(target, hostInput, dataType, rawInput.get());
  if (rawMetaInfo) {
    copy(target, metaInfoFlat, UNSIGNED_SHORT, rawMetaInfo.get());
  }
  copy(target, nzValuesFlat, dataType, rawNzInfo.get());

  if (!ignoreData && (doBwdPass || doWuPass)) {
    copy(target, hostOutputGrad, outputGrad.elementType(), rawOutputGrad.get());
  }

  device.bind([&](const Device &d) {
    engine.loadAndRun(d);
    if (reportTotalCycleCounts) {
      std::uint64_t cyclesBuffer;
      engine.readTensor("fwdCycles", &cyclesBuffer, &cyclesBuffer + 1);
      std::cerr << "  Forward pass cycles: " << cyclesBuffer << "\n";
      if (doBwdPass) {
        engine.readTensor("bwdCycles", &cyclesBuffer, &cyclesBuffer + 1);
        std::cerr << "  GradA pass cycles: " << cyclesBuffer << "\n";
      }
      if (doWuPass) {
        engine.readTensor("wuCycles", &cyclesBuffer, &cyclesBuffer + 1);
        std::cerr << "  GradW pass cycles: " << cyclesBuffer << "\n";
      }
    }
  });

  bool matchesModel = true;
  if (!ignoreData) {
    const double relTolerance = dataType == HALF ? HALF_REL_TOL : FLOAT_REL_TOL;
    copy(target, outputActs.elementType(), rawOutputActs.get(), hostOutputActs);
    boost::multi_array<double, 2> hostDenseWeights(
        boost::extents[outputSize][inputSize]);
    boost::multi_array<double, 2> modelOutputActs(
        boost::extents[batchSize][outputSize]);
    hostDenseWeights = poplibs_test::sparse::csrToDenseMatrix(
        csrMatrix.nzValues.data(), csrMatrix.columnIndices.data(),
        csrMatrix.rowIndices.data(), csrMatrix.nzValues.size(), outputSize,
        inputSize, blockRows, blockCols);

    poplibs_test::gemm::generalMatrixMultiply(hostInput, hostDenseWeights,
                                              modelOutputActs, false, true);
    matchesModel &= checkIsClose("outputActs", hostOutputActs, modelOutputActs,
                                 relTolerance);
    if (doBwdPass) {
      copy(target, inputGrad.elementType(), rawInputGrad.get(), hostInputGrad);
      boost::multi_array<double, 2> modelInputGrad(
          boost::extents[batchSize][inputSize]);
      poplibs_test::gemm::generalMatrixMultiply(
          hostOutputGrad, hostDenseWeights, modelInputGrad, false, false);
      matchesModel &= checkIsClose("inputGrad", hostInputGrad, modelInputGrad,
                                   relTolerance);
    }
    if (doWuPass) {
      copy(target, weightGrad.elementType(), rawWeightGrad.get(),
           hostWeightGrad);
      boost::multi_array<double, 2> modelWeightGrad(
          boost::extents[outputSize][inputSize]);
      poplibs_test::gemm::generalMatrixMultiply(hostOutputGrad, hostInput,
                                                modelWeightGrad, true, false);

      std::vector<EType> modelNzValuesCSR;
      auto columnIdxIt = csrMatrix.columnIndices.begin();
      for (auto rowIt = std::next(csrMatrix.rowIndices.begin());
           rowIt != csrMatrix.rowIndices.end(); ++rowIt) {
        const auto nnzThisRow = *rowIt - *std::prev(rowIt);
        const auto rowIdx =
            (std::distance(csrMatrix.rowIndices.begin(), rowIt) - 1) *
            blockRows;
        for (std::size_t i = 0; i < nnzThisRow / blockArea; ++i) {
          const auto columnIdx = *columnIdxIt++;
          for (std::size_t r = 0; r != blockRows; ++r) {
            for (std::size_t c = 0; c != blockCols; ++c) {
              // row major order within a block
              modelNzValuesCSR.emplace_back(
                  modelWeightGrad[rowIdx + r][columnIdx + c]);
            }
          }
        }
      }

      assert(modelNzValuesCSR.size() ==
             csrMatrix.columnIndices.size() * blockArea);

      SparsityDataImpl<EType> actualBuckets;
      std::vector<EType> actualWeightGrads;
      actualWeightGrads.reserve(hostWeightGrad.size());
      std::copy(hostWeightGrad.begin(), hostWeightGrad.end(),
                std::back_inserter(actualWeightGrads));

      actualBuckets.nzValues = std::move(actualWeightGrads);
      actualBuckets.metaInfo = std::move(buckets.metaInfo);
      auto actualCSR = partitioner.sparsityDataImplToCSRMatrix(actualBuckets);

      auto ait = actualCSR.nzValues.begin();
      auto mit = modelNzValuesCSR.begin();
      for (; ait != actualCSR.nzValues.end(); ++ait, ++mit) {
        bool elemMatch = checkIsClose(*mit, *ait, relTolerance);
        if (!elemMatch) {
          std::cerr << "mismatch at  WeightsGrad.nz[";
          std::cerr << std::distance(actualCSR.nzValues.begin(), ait);
          std::cerr << "]:" << *mit << "!=" << *ait << "\n";
        }
        matchesModel &= elemMatch;
      }
      auto columnsMatch = std::equal(
          actualCSR.columnIndices.begin(), actualCSR.columnIndices.end(),
          csrMatrix.columnIndices.begin(), csrMatrix.columnIndices.end());
      if (!columnsMatch) {
        std::cerr << "CSR columns indices do not match\n";
      }
      matchesModel &= columnsMatch;
    }
  }

  if (denseGradWSerialSplits && doWuPass) {
    auto serialSplits =
        fullyConnectedDenseGradWSerialSplits(graph, dataType, params, options);
    std::cerr << "Dense GradW serial splits : "
              << "   groups " << std::get<0>(serialSplits)
              << "   input channel " << std::get<1>(serialSplits)
              << "   output channel " << std::get<2>(serialSplits) << "\n";
  }

  std::cerr << "Done\n";

  if (profile) {
    engine.printProfileSummary(std::cout, {{"showExecutionSteps", "true"}});
  }

  if (!matchesModel) {
    std::cerr << "Validation failed\n";
    return 1;
  }

  return 0;
} catch (const poplar::graph_memory_allocation_error &e) {
  if (!e.profilePath.empty()) {
    poplar::printGraphSummary(std::cerr, e.profilePath,
                              {{"showVarStorage", "true"}});
  }
  throw;
}