// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#include "poplibs_support/VectorUtils.hpp"
#include "poplibs_support/print.hpp"
#include <boost/multi_array.hpp>
#include <boost/optional.hpp>
#include <boost/optional/optional_io.hpp>
#include <boost/program_options.hpp>
#include <boost/test/tools/floating_point_comparison.hpp>
#include <cassert>
#include <exception>
#include <fstream>
#include <istream>
#include <ostream>
#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplibs_support/Compiler.hpp>
#include <poplibs_support/MultiArray.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/Convolution.hpp>
#include <poplibs_test/GeneralMatrixAdd.hpp>
#include <poplibs_test/Multirate.hpp>
#include <poplibs_test/Pass.hpp>
#include <poplibs_test/Util.hpp>
#include <poplin/ConvUtil.hpp>
#include <poplin/MultiConvolution.hpp>
#include <poplin/codelets.hpp>
#include <popnn/Pooling.hpp>
#include <popnn/codelets.hpp>
#include <popops/ElementWise.hpp>
#include <popops/ScaledAdd.hpp>
#include <popops/codelets.hpp>
#include <poputil/exceptions.hpp>
#include <random>

// Default tolerances used in tests
#define FLOAT_REL_TOL 0.01
#define HALF_REL_TOL 0.1
#define FLOAT_ABS_TOL 1e-6
#define HALF_ABS_TOL 1e-5

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

using Array1d = boost::multi_array<double, 1>;
using Array3d = boost::multi_array<double, 3>;
using Array4d = boost::multi_array<double, 4>;
using Array1dRef = boost::multi_array_ref<double, 1>;
using Array3dRef = boost::multi_array_ref<double, 3>;

struct OctConvUserParams {
  Type inputType;
  Type outputType;
  std::size_t batchSize;
  std::vector<std::size_t> inputFieldShape;
  std::vector<std::size_t> kernelShape;
  std::size_t inChansPerConvGroup;
  std::size_t outChansPerConvGroup;
  std::size_t numConvGroups;
};

// OctConv input and output shapes
struct OctConvData {
  std::string name;
  std::vector<std::size_t> shape;
  std::size_t numFieldElements;
  std::size_t chansPerConvGroup;
  Tensor tensor;
  std::shared_ptr<Array3d> host;
  std::shared_ptr<Array3d> model; // only applicable for output

  OctConvData(std::string name, std::vector<std::size_t> shape,
              std::size_t chansPerConvGroup)
      : name(name), shape(shape), numFieldElements(product(shape)),
        chansPerConvGroup(chansPerConvGroup) {}
};

// Parameters for each component convolution
struct ConvParameters {
  std::string name;
  poplin::ConvParams fwdParams;
  poplin::ConvParams bwdParams;
  unsigned preDownsamplingRate;
  unsigned postUpsamplingRate;

  // OctConv input data (i.e., before pre-downsampling)
  std::shared_ptr<OctConvData> inputData;

  // OctConv output data (i.e., after post-upsampling)
  std::shared_ptr<OctConvData> outputData;

  ConvParameters(std::string name, poplin::ConvParams fwdParams,
                 poplin::ConvParams bwdParams, unsigned preDnsampling,
                 unsigned postUpsampling,
                 std::shared_ptr<OctConvData> inputData,
                 std::shared_ptr<OctConvData> outputData)
      : name(name), fwdParams(fwdParams), bwdParams(bwdParams),
        preDownsamplingRate(preDnsampling), postUpsamplingRate(postUpsampling),
        inputData(inputData), outputData(outputData) {}
};

// Input shapes for a particular frequency
struct InputFreqParameters {
  std::string name;
  std::size_t numChans;
  std::vector<std::size_t> fShape;
};

// Output shapes for a particular frequency
struct OutputFreqParameters {
  std::string name;
  std::size_t numChans;
  std::shared_ptr<OctConvData> data;
};

static std::vector<poplar::Tensor>
createAllWeights(poplar::Graph &graph,
                 const std::vector<poplin::multiconv::CreateTensorArgs> &args,
                 const poplar::OptionFlags &multiConvOptions,
                 poplin::PlanningCache *cache) {
  std::vector<Tensor> weights;
  for (size_t i = 0; i < args.size(); i++) {
    weights.push_back(poplin::multiconv::createWeights(
        graph, args, i, multiConvOptions, cache));
  }
  return weights;
}

static std::vector<poplar::Tensor>
createAllInputs(poplar::Graph &graph,
                const std::vector<poplin::multiconv::CreateTensorArgs> &args,
                const poplar::OptionFlags &multiConvOptions,
                poplin::PlanningCache *cache) {
  std::vector<Tensor> inputs;
  for (size_t i = 0; i < args.size(); i++) {
    inputs.push_back(poplin::multiconv::createInput(graph, args, i,
                                                    multiConvOptions, cache));
  }
  return inputs;
}

// Upsample or downsample a tensor shape of arbitrary number of dimensions
// by an integer factor
static std::vector<std::size_t>
TensorShapeRateChange(const std::vector<std::size_t> &inShape,
                      const unsigned factor, const bool up) {
  if (!up &&
      std::any_of(inShape.begin(), inShape.end(),
                  [factor](const std::size_t v) { return v % factor; })) {
    throw poputil::poplibs_error("At least one field dimension is not "
                                 "divisible by downsampling factor " +
                                 std::to_string(factor));
  }
  std::vector<std::size_t> outShape(inShape);
  std::transform(inShape.begin(), inShape.end(), outShape.begin(),
                 [factor, up](const std::size_t &v) {
                   return up ? (v * factor) : (v / factor);
                 });
  return outShape;
}

//
// The component convolutions are listed in the following table.
//
//    name  |pool2x|num-in-channels|field-size|num-out-channels|upsample2x
//  --------+------+---------------+----------+----------------+-----------
//   FreqLL |  -   |    a_i  * n_i |   f_lo   |    a_o  * n_o  |    -
//   FreqLH |  -   |    a_i  * n_i |   f_lo   | (1-a_o) * n_o  |    X
//   FreqHL |  X   | (1-a_i) * n_i |   f_lo   |    a_o  * n_o  |    -
//   FreqHH |  -   | (1-a_i) * n_i |   f_hi   | (1-a_o) * n_o  |    -
//
// where
//  - a_i and a_o denote alpha_in and alpha_out resply.
//  - n_i and n_o denote number of input channels and output channels resply.
//  - f_lo and f_hi denote field sizes for lo and hi frequencies resply.
//
// As shown above, there are two special cases:
//  - FreqLH convolution generates a smaller output field size (f_lo) than
//    the natural low frequency field size. Hence upsampling by 2x must be
//    applied on the output before it can be added to FreqHH output.
//  - FreqHL convolution uses a smaller input field size (f_lo) than the
//    natural high frequency field size. Hence average-pooling 2x is required
//    prior to the convolution.
//
static std::tuple<std::vector<ConvParameters>,
                  std::vector<std::shared_ptr<OctConvData>>,
                  std::vector<std::shared_ptr<OctConvData>>>
splitConvParamsByFrequency(
    const OctConvUserParams &userParam,
    const poplin::ConvParams::InputTransform &inputTransform,
    const poplin::ConvParams::InputTransform &kernelTransform,
    const poplin::ConvParams::OutputTransform &outputTransform,
    const double alphaIn, const double alphaOut) {
  auto bSize = userParam.batchSize;
  auto nCGroups = userParam.numConvGroups;
  unsigned samplingRate = 2;
  std::vector<ConvParameters> convParams;
  std::vector<std::shared_ptr<OctConvData>> octInputData;
  std::vector<std::shared_ptr<OctConvData>> octOutputData;

  // Calculate Low Frequency parameters
  auto pooledFieldShape =
      TensorShapeRateChange(userParam.inputFieldShape, samplingRate, false);
  std::size_t inChansLo =
      static_cast<std::size_t>(alphaIn * userParam.inChansPerConvGroup);
  std::size_t outChansLo =
      static_cast<std::size_t>(alphaOut * userParam.outChansPerConvGroup);

  // Enumerate shapes for both low and high frequency
  enum FreqType { LO, HI };
  std::map<FreqType, InputFreqParameters> octInputFreq = {
      {LO, {"Lo", inChansLo, pooledFieldShape}},
      {HI,
       {"Hi", userParam.inChansPerConvGroup - inChansLo,
        userParam.inputFieldShape}}};
  std::map<FreqType, OutputFreqParameters> octOutputFreq = {
      {LO, {"Lo", outChansLo, nullptr}},
      {HI, {"Hi", userParam.outChansPerConvGroup - outChansLo, nullptr}}};

  // Iterate over Low and High frequency input shapes
  for (auto &inData : octInputFreq) {
    auto inType = inData.first;
    auto &in = inData.second;
    if (in.numChans > 0) {
      // Allocate input oct-conv data. Model data is not required for input
      auto allocOctData =
          [nCGroups, bSize](const std::string name,
                            const std::size_t nChansPGrp,
                            std::vector<std::size_t> &shape,
                            std::vector<std::shared_ptr<OctConvData>> &storage,
                            const bool allocModel) {
            auto data = std::make_shared<OctConvData>(name, shape, nChansPGrp);
            storage.push_back(data);
            auto nChans = nChansPGrp * nCGroups;
            auto fSize = product(shape);
            data->host =
                std::make_shared<Array3d>(boost::extents[bSize][nChans][fSize]);
            if (allocModel) {
              data->model = std::make_shared<Array3d>(
                  boost::extents[bSize][nChans][fSize]);
            }
            return data;
          };
      allocOctData("freq" + in.name, in.numChans, in.fShape, octInputData,
                   false);

      // Iterate over Low and High frequency output shapes
      for (auto &outData : octOutputFreq) {
        auto outType = outData.first;
        auto &out = outData.second;
        if (out.numChans > 0) {
          auto postUpsampling =
              (inType == LO) && (outType == HI) ? samplingRate : 1;
          auto preDnsampling =
              (inType == HI) && (outType == LO) ? samplingRate : 1;
          auto fShape = (inType == HI) && (outType == HI)
                            ? octInputFreq[HI].fShape
                            : octInputFreq[LO].fShape;
          auto fwdParams = poplin::ConvParams{userParam.inputType,
                                              userParam.outputType,
                                              bSize,
                                              fShape,
                                              userParam.kernelShape,
                                              in.numChans,
                                              out.numChans,
                                              nCGroups,
                                              inputTransform,
                                              kernelTransform,
                                              outputTransform};
          auto bwdParams = getGradientParams(fwdParams);
          auto outFieldShape = fwdParams.getOutputFieldShape();
          if (postUpsampling > 1) {
            outFieldShape =
                TensorShapeRateChange(outFieldShape, postUpsampling, true);
          }
          if (out.data == nullptr) {
            // allocate input oct-conv data
            out.data = allocOctData("freq" + out.name, out.numChans,
                                    outFieldShape, octOutputData, true);
          } else if (product(outFieldShape) != out.data->numFieldElements) {
            throw poputil::poplibs_error(
                "Output field shape for convolution-" + in.name + out.name +
                " is not compatible with OctConv output-" + out.name);
          }
          convParams.emplace_back("freq" + in.name + out.name, fwdParams,
                                  bwdParams, preDnsampling, postUpsampling,
                                  octInputData.back(), out.data);
        }
      }
    }
  }
  return {convParams, octInputData, octOutputData};
}

static void
writeRandomValues(const poplar::Target &target, const poplar::Type type,
                  std::vector<std::shared_ptr<OctConvData>> &octConv,
                  const double min, const double max,
                  std::mt19937 &randomEngine) {
  for (unsigned i = 0; i < octConv.size(); ++i) {
    writeRandomValues(target, type, *(octConv[i]->host), min, max,
                      randomEngine);
  }
}

static void copy(const poplar::Target &target, const poplar::Type &type,
                 const std::vector<std::shared_ptr<OctConvData>> &octConv,
                 const std::vector<std::unique_ptr<char[]>> &dst) {
  for (unsigned i = 0; i < octConv.size(); ++i) {
    copy(target, *(octConv[i]->host), type, dst[i].get());
  }
}

static void copy(const poplar::Target &target, const poplar::Type &type,
                 const std::vector<std::unique_ptr<char[]>> &src,
                 const std::vector<std::shared_ptr<OctConvData>> &octConv) {
  for (unsigned i = 0; i < octConv.size(); ++i) {
    copy(target, type, src[i].get(), *(octConv[i]->host));
  }
}

static std::vector<Array3d>
createMultiArrayInput(const std::vector<ConvParameters> &octParams,
                      const bool bwdFlag) {
  auto numConvolutions = octParams.size();
  std::vector<Array3d> multiArray;
  for (unsigned i = 0; i < numConvolutions; ++i) {
    auto params = bwdFlag ? octParams[i].bwdParams : octParams[i].fwdParams;
    auto numConvGroups = params.getNumConvGroups();
    auto numInChans = params.getNumInputChansPerConvGroup() * numConvGroups;
    auto inputFieldShape = params.getInputFieldShape();
    auto batchSize = params.getBatchSize();
    auto numInFieldElems = product(inputFieldShape);
    multiArray.push_back(
        Array3d(boost::extents[batchSize][numInChans][numInFieldElems]));
  }
  return multiArray;
}

static std::vector<Array3d>
createMultiArrayOutput(const std::vector<ConvParameters> &params) {
  auto numConvolutions = params.size();
  std::vector<Array3d> multiArray;
  for (unsigned i = 0; i < numConvolutions; ++i) {
    auto numConvGroups = params[i].fwdParams.getNumConvGroups();
    auto numOutChans =
        params[i].fwdParams.getNumOutputChansPerConvGroup() * numConvGroups;
    auto batchSize = params[i].fwdParams.getBatchSize();
    auto numOutFieldElems = product(params[i].fwdParams.getOutputFieldShape());
    multiArray.push_back(
        Array3d(boost::extents[batchSize][numOutChans][numOutFieldElems]));
  }
  return multiArray;
}

static std::vector<Array4d>
createMultiArrayWeights(const std::vector<ConvParameters> &params,
                        const std::vector<std::size_t> kernelShape) {
  auto numConvolutions = params.size();
  std::vector<Array4d> multiArray;
  for (unsigned i = 0; i < numConvolutions; ++i) {
    auto numConvGroups = params[i].fwdParams.getNumConvGroups();
    auto numInChans = params[i].fwdParams.getNumInputChansPerConvGroup();
    auto numOutChans = params[i].fwdParams.getNumOutputChansPerConvGroup();
    auto numKernelElems = product(kernelShape);
    multiArray.push_back(Array4d(boost::extents[numConvGroups][numOutChans]
                                               [numInChans][numKernelElems]));
  }
  return multiArray;
}

// Create OctConv input tensors for each frequency
static void
createOctConvTensors(poplar::Graph &graph, poplin::PlanningCache &cache,
                     const OctConvUserParams &userParam,
                     const poplin::ConvParams::InputTransform &inputTransform,
                     const poplin::ConvParams::InputTransform &kernelTransform,
                     const poplin::ConvParams::OutputTransform &outputTransform,
                     const OptionFlags &fwdOptions,
                     const OptionFlags &bwdOptions,
                     std::vector<std::shared_ptr<OctConvData>> &octData,
                     const poplar::OptionFlags &multiConvOptions) {
  auto numConvGroups = userParam.numConvGroups;
  auto batchSize = userParam.batchSize;

  // Create up to two Oct-convolution inputs for Low and High Frequencies
  // respectively.
  using CreateTensorArgs = poplin::multiconv::CreateTensorArgs;
  std::vector<CreateTensorArgs> prevActArgs(octData.size());
  for (unsigned i = 0; i < prevActArgs.size(); ++i) {
    prevActArgs[i].params = poplin::ConvParams{userParam.inputType,
                                               userParam.outputType,
                                               batchSize,
                                               octData[i]->shape,
                                               userParam.kernelShape,
                                               octData[i]->chansPerConvGroup,
                                               userParam.outChansPerConvGroup,
                                               numConvGroups,
                                               inputTransform,
                                               kernelTransform,
                                               outputTransform};
    prevActArgs[i].options = fwdOptions;
    prevActArgs[i].name = octData[i]->name + "_prevAct";
  }

  // Create input tensors and model multiarray
  auto prevAct = createAllInputs(graph, prevActArgs, multiConvOptions, &cache);
  for (unsigned i = 0; i < octData.size(); ++i) {
    octData[i]->tensor = prevAct[i];
  }
}

// Create tensors for weights and zDeltas for forward, backward and weight
// update passes.
static std::tuple<std::vector<Tensor>, std::vector<Tensor>>
createConvInputTensors(poplar::Graph &graph, poplin::PlanningCache &cache,
                       const std::vector<ConvParameters> &params,
                       const OptionFlags &fwdOptions,
                       const OptionFlags &bwdOptions, const bool createZDeltas,
                       const poplar::OptionFlags &multiConvOptions) {
  // Create the required number of convolution inputs
  auto numConvolutions = params.size();
  using CreateTensorArgs = poplin::multiconv::CreateTensorArgs;
  std::vector<CreateTensorArgs> weightsArgs(numConvolutions);
  std::vector<CreateTensorArgs> zDeltasArgs(numConvolutions);
  for (unsigned i = 0; i < numConvolutions; ++i) {
    weightsArgs[i].name = params[i].name + "_weights";
    weightsArgs[i].params = params[i].fwdParams;
    weightsArgs[i].options = fwdOptions;
    zDeltasArgs[i].name = params[i].name + "_zDeltas";
    zDeltasArgs[i].params = params[i].bwdParams;
    zDeltasArgs[i].options = bwdOptions;
  }
  auto weights = createAllWeights(graph, weightsArgs, multiConvOptions, &cache);
  std::vector<Tensor> zDeltas;
  if (createZDeltas) {
    zDeltas = createAllInputs(graph, zDeltasArgs, multiConvOptions, &cache);
  }
  return {weights, zDeltas};
}

static std::vector<poplin::multiconv::ConvolutionArgs>
getConvolutionArguments(const std::vector<Tensor> &inA,
                        const std::vector<Tensor> &inB, const bool bwdFlag,
                        const std::vector<ConvParameters> &octParams,
                        const OptionFlags &options) {
  std::vector<poplin::multiconv::ConvolutionArgs> convArgs;
  auto numConvolutions = octParams.size();
  for (unsigned i = 0; i < numConvolutions; ++i) {
    poplin::ConvParams params =
        bwdFlag ? octParams[i].bwdParams : octParams[i].fwdParams;
    convArgs.push_back({inA[i], inB[i], params, options});
  }
  return convArgs;
}

static std::vector<poplin::multiconv::ConvWeightUpdateArgs>
getConvolutionWeightUpdateArguments(
    const std::vector<Tensor> &prevAct, const std::vector<Tensor> &zDeltas,
    const std::vector<Tensor> &weights, const Tensor &scale,
    const std::vector<ConvParameters> &octParams, const OptionFlags &options) {
  std::vector<poplin::multiconv::ConvWeightUpdateArgs> wuArgs;
  auto numConvolutions = octParams.size();
  for (unsigned i = 0; i < numConvolutions; ++i) {
    wuArgs.push_back({zDeltas[i], weights[i], prevAct[i], scale,
                      octParams[i].fwdParams, options});
  }
  return wuArgs;
}

static std::vector<std::unique_ptr<char[]>>
allocateOctHostMemory(Graph &graph, const std::string name,
                      const std::vector<ConvParameters> &params,
                      const std::vector<Tensor> &tensors,
                      boost::optional<Sequence &> uploadProg,
                      boost::optional<Sequence &> downloadProg,
                      std::vector<std::pair<std::string, char *>> &tmap) {
  std::vector<std::unique_ptr<char[]>> rawHostMem(tensors.size());
  for (unsigned i = 0; i < tensors.size(); ++i) {
    auto prefix = params[i].name;
    rawHostMem[i] = allocateHostMemoryForTensor(
        tensors[i], prefix + name, graph, uploadProg, downloadProg, tmap);
  }
  return rawHostMem;
}

// Upsample tensor by a given factor along every dimension
static Tensor upsampleTensor(Graph &graph, Tensor &input,
                             const unsigned samplingRate,
                             const std::size_t numFieldDims) {
  auto inputShape = input.shape();
  std::vector<std::size_t> upSamplingPartial(inputShape);
  auto size = product<std::size_t>(inputShape);
  std::size_t elemsToDupl = 1;
  std::size_t numDupl = size / elemsToDupl;
  Tensor inShuffled = input;
  int currDim = upSamplingPartial.size() - 1;
  int firstFieldDim = upSamplingPartial.size() - numFieldDims;
  while (currDim >= firstFieldDim) {
    upSamplingPartial[currDim] *= samplingRate;
    inShuffled = inShuffled.reshape({numDupl, elemsToDupl})
                     .broadcast(samplingRate, 1)
                     .reshape(upSamplingPartial);
    elemsToDupl *= upSamplingPartial[currDim];
    numDupl = product<std::size_t>(upSamplingPartial) / elemsToDupl;
    currDim--;
  }
  return inShuffled;
}

// Use Average-Pooling to sub-sample tensor by a given factor along every
// dimension.
static Tensor downsampleTensor(Graph &graph, const Tensor &input,
                               const unsigned samplingRate,
                               const std::size_t numFieldDims, Sequence &prog,
                               const OptionFlags &poolOptions = {}) {
  auto inputShape = input.shape();
  popnn::PoolingType poolingType = popnn::PoolingType::AVG;
  std::vector<std::size_t> inputFieldShape;
  std::vector<std::size_t> kernelShape;
  std::vector<unsigned> stride;
  std::vector<int> paddingLower;
  std::vector<int> paddingUpper;
  unsigned firstFieldDim = inputShape.size() - numFieldDims;
  for (unsigned i = firstFieldDim; i < inputShape.size(); ++i) {
    inputFieldShape.push_back(inputShape[i]);
    kernelShape.push_back(samplingRate);
    stride.push_back(samplingRate);
    paddingLower.push_back(0);
    paddingUpper.push_back(0);
  }
  const auto poolParams = popnn::pooling::PoolParams(
      poolingType, inputFieldShape, kernelShape, stride, paddingLower,
      paddingUpper, inputShape[1], inputShape[0], input.elementType());
  return popnn::pooling::pool(graph, poolParams, input, prog, "", poolOptions);
}

static std::vector<Tensor> preProcess(Graph &graph,
                                      const std::vector<ConvParameters> &params,
                                      Sequence &prog,
                                      const OptionFlags poolOptions = {}) {
  std::vector<Tensor> output(params.size());
  for (unsigned i = 0; i < params.size(); ++i) {
    auto &input = params[i].inputData->tensor;
    auto samplingRate = params[i].preDownsamplingRate;
    if (samplingRate > 1) {
      auto numFieldDims = params[i].fwdParams.getNumFieldDims();
      output[i] = downsampleTensor(graph, input, samplingRate, numFieldDims,
                                   prog, poolOptions);
    } else {
      output[i] = input;
    }
  }
  return output;
}

static void postProcess(Graph &graph, std::vector<ConvParameters> &params,
                        std::vector<Tensor> &input, Sequence &prog) {
  for (unsigned i = 0; i < params.size(); ++i) {
    Tensor tensor = input[i];
    if (params[i].postUpsamplingRate > 1) {
      auto name = params[i].name + "_upsampled";
      tensor = upsampleTensor(graph, input[i], params[i].postUpsamplingRate,
                              params[i].fwdParams.getNumFieldDims());
    }

    // Accumulate outputs of individual convolution paths which are of the same
    // shape.
    if (!params[i].outputData->tensor.valid()) {
      params[i].outputData->tensor = tensor;
    } else {
      params[i].outputData->tensor =
          popops::add(graph, params[i].outputData->tensor, tensor, prog,
                      params[i].outputData->name);
    }
  }
  return;
}

static void convolutionForwardModel(const std::vector<ConvParameters> &params,
                                    std::vector<Array3d> &prevAct,
                                    std::vector<Array4d> &weights,
                                    std::vector<Array3d> &nextAct,
                                    Array1dRef &bias) {
  for (unsigned i = 0; i < nextAct.size(); ++i) {
    auto fieldShape = params[i].fwdParams.getInputFieldShape();
    auto kernelShape = params[i].fwdParams.getKernelShape();
    auto &inTransf = params[i].fwdParams.inputTransform;
    auto &kTransf = params[i].fwdParams.kernelTransform;
    auto &outTransf = params[i].fwdParams.outputTransform;
    poplibs_test::conv::convolution(
        vectorConvert<unsigned>(fieldShape), inTransf.truncationLower,
        inTransf.truncationUpper, inTransf.dilation, inTransf.paddingLower,
        inTransf.paddingUpper, inTransf.flip,
        vectorConvert<unsigned>(kernelShape), kTransf.truncationLower,
        kTransf.truncationUpper, kTransf.dilation, kTransf.paddingLower,
        kTransf.paddingUpper, kTransf.flip, outTransf.truncationLower,
        outTransf.truncationUpper, outTransf.stride, outTransf.paddingLower,
        outTransf.paddingUpper, prevAct[i], weights[i], bias, nextAct[i]);
  }
}

static void convolutionBackwardModel(const std::vector<ConvParameters> &params,
                                     std::vector<Array3d> &zDeltas,
                                     std::vector<Array4d> &weights,
                                     std::vector<Array3d> &prevDeltas) {
  for (unsigned i = 0; i < prevDeltas.size(); ++i) {
    auto fieldShape = params[i].fwdParams.getInputFieldShape();
    auto kernelShape = params[i].fwdParams.getKernelShape();
    auto &inTransf = params[i].fwdParams.inputTransform;
    auto &kTransf = params[i].fwdParams.kernelTransform;
    auto &outTransf = params[i].fwdParams.outputTransform;
    poplibs_test::conv::convolutionBackward(
        vectorConvert<unsigned>(fieldShape), inTransf.truncationLower,
        inTransf.truncationUpper, inTransf.dilation, inTransf.paddingLower,
        inTransf.paddingUpper, inTransf.flip,
        vectorConvert<unsigned>(kernelShape), kTransf.truncationLower,
        kTransf.truncationUpper, kTransf.dilation, kTransf.paddingLower,
        kTransf.paddingUpper, kTransf.flip, outTransf.truncationLower,
        outTransf.truncationUpper, outTransf.stride, outTransf.paddingLower,
        outTransf.paddingUpper, zDeltas[i], weights[i], prevDeltas[i]);
  }
}

static void weightUpdateModel(const std::vector<ConvParameters> &params,
                              const std::vector<Array3d> &prevAct,
                              const std::vector<Array3d> &zDeltas,
                              const std::vector<Array4d> &weights,
                              const double learningRate,
                              const Array1dRef &bias) {
  for (unsigned i = 0; i < weights.size(); ++i) {
    auto fieldShape = params[i].fwdParams.getInputFieldShape();
    auto kernelShape = params[i].fwdParams.getKernelShape();
    auto &inTransf = params[i].fwdParams.inputTransform;
    auto &kTransf = params[i].fwdParams.kernelTransform;
    auto &outTransf = params[i].fwdParams.outputTransform;
    poplibs_test::conv::weightUpdate(
        vectorConvert<unsigned>(fieldShape), inTransf.truncationLower,
        inTransf.truncationUpper, inTransf.dilation, inTransf.paddingLower,
        inTransf.paddingUpper, inTransf.flip,
        vectorConvert<unsigned>(kernelShape), kTransf.truncationLower,
        kTransf.truncationUpper, kTransf.dilation, kTransf.paddingLower,
        kTransf.paddingUpper, kTransf.flip, outTransf.truncationLower,
        outTransf.truncationUpper, outTransf.stride, outTransf.paddingLower,
        outTransf.paddingUpper, learningRate, prevAct[i], zDeltas[i],
        weights[i], bias);
  }
}

static void PreProcessModel(const std::vector<ConvParameters> &params,
                            std::vector<Array3d> &outputs) {
  for (unsigned i = 0; i < params.size(); ++i) {
    auto &input = params[i].inputData->host;
    if (params[i].preDownsamplingRate > 1) {
      auto inShape = params[i].fwdParams.getInputFieldShape();
      poplibs_test::downsample(inShape, params[i].preDownsamplingRate, *input,
                               outputs[i]);
    } else {
      outputs[i] = *input;
    }
  }
}

static void PostProcessModel(std::vector<ConvParameters> &params,
                             const std::size_t numConvGroups,
                             const std::size_t batchSize,
                             const std::vector<Array3d> &inputs) {
  for (unsigned i = 0; i < params.size(); ++i) {
    auto &output = params[i].outputData->model;
    auto numOutElems = params[i].outputData->numFieldElements;
    auto numOutChans =
        params[i].fwdParams.getNumOutputChansPerConvGroup() * numConvGroups;
    Array3d tensor(boost::extents[batchSize][numOutChans][numOutElems]);
    if (params[i].postUpsamplingRate > 1) {
      auto outShape = params[i].fwdParams.getOutputFieldShape();
      poplibs_test::upsample(outShape, params[i].postUpsamplingRate, inputs[i],
                             tensor);
    } else {
      tensor = inputs[i];
    }
    // Accumulate outputs of individual convolution paths which are of the same
    // shape.
    if (output->size() == 0) {
      *output = tensor;
    } else {
      auto outputRef = Array3dRef(
          output->data(), boost::extents[batchSize][numOutChans][numOutElems]);
      auto tensorRef = Array3dRef(
          tensor.data(), boost::extents[batchSize][numOutChans][numOutElems]);
      poplibs_test::axpby::add(outputRef, tensorRef, outputRef);
    }
  }
}

static void printTensors(const std::string name,
                         const std::vector<ConvParameters> &params,
                         const std::vector<Tensor> &tensors, Sequence &prog) {
  for (unsigned i = 0; i < tensors.size(); ++i) {
    prog.add(PrintTensor(name + params[i].name, tensors[i]));
  }
}

template <unsigned N>
static bool
multipleCheckIsClose(const std::vector<ConvParameters> params,
                     const std::vector<boost::multi_array<double, N>> &inA,
                     const std::vector<boost::multi_array<double, N>> &inB,
                     const double relativeTol, const double absoluteTol) {
  for (unsigned i = 0; i < inA.size(); ++i) {
    if (!checkIsClose(params[i].name, inA[i], inB[i], relativeTol, absoluteTol))
      return false;
  }
  return true;
}

static bool
multipleCheckIsClose(const std::vector<std::shared_ptr<OctConvData>> &data,
                     const double relativeTol, const double absoluteTol) {
  for (unsigned i = 0; i < data.size(); ++i) {
    if (!checkIsClose(data[i]->name, *(data[i]->host), *(data[i]->model),
                      relativeTol, absoluteTol))
      return false;
  }
  return true;
}

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

  DeviceType deviceType = DeviceType::IpuModel2;
  bool doPrintTensors = false;
  unsigned fwdInChansPerConvGroup;
  unsigned fwdOutChansPerConvGroup;
  ShapeOption<std::size_t> inputFieldShapeOption;
  ShapeOption<std::size_t> kernelShapeOption;
  unsigned numConvGroups = 1;
  ShapeOption<unsigned> truncationLowerOption, truncationUpperOption,
      truncationOption;
  ShapeOption<unsigned> inDilationOption;
  ShapeOption<unsigned> paddingLowerOption, paddingUpperOption, paddingOption;
  ShapeOption<bool> flipInputOption;
  ShapeOption<unsigned> kernelTruncationLowerOption,
      kernelTruncationUpperOption, kernelTruncationOption;
  ShapeOption<unsigned> kernelDilationOption;
  ShapeOption<unsigned> kernelPaddingLowerOption, kernelPaddingUpperOption,
      kernelPaddingOption;
  ShapeOption<bool> flipKernelOption;
  ShapeOption<unsigned> outputTruncationOption, outputTruncationLowerOption,
      outputTruncationUpperOption;
  ShapeOption<unsigned> strideOption;
  ShapeOption<unsigned> outputPaddingOption, outputPaddingLowerOption,
      outputPaddingUpperOption;
  unsigned batchSize;
  Type inputType;
  Type outputType;
  double absoluteTolerance, relativeTolerance;
  double alpha = 1.0;
  double alphaIn = 1.0;
  double alphaOut = 1.0;
  unsigned numIPUs = 1;
  boost::optional<unsigned> tilesPerIPU;
  bool reportPlan;
  bool reportVarStorage;
  bool remapOutputTensor;
  bool poolOptimizeForSpeed;
  std::string multiConvOptionsString;

  Pass pass = Pass::ALL;
  poplin::PlanningCache cache;

  boost::optional<std::string> profileDir;

  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),
     "Device type")
    ("profile", "Output profiling report")
    ("profile-dir",
     po::value<decltype(profileDir)>(&profileDir)
      ->default_value(boost::none),
     "Write profile files to the specified directory.")
    ("print",
     po::value<bool>(&doPrintTensors)->default_value(doPrintTensors),
     "Print the tensors")
    ("ignore-data", "Don't upload and download the results from the device. "
     "Note that this means the result is not validated against the model.")
    ("input-channels", po::value<unsigned>(&fwdInChansPerConvGroup)->required(),
     "Number of input channels per grouped convolution")
    ("output-channels",
     po::value<unsigned>(&fwdOutChansPerConvGroup)->required(),
     "Number of output channels per grouped convolution")
    ("field",
     po::value<ShapeOption<std::size_t>>(&inputFieldShapeOption)->required(),
      "Field size")
    ("kernel-size",
      po::value<ShapeOption<std::size_t>>(&kernelShapeOption)->default_value(1),
     "Size of square kernel. If set, it is an error to also set either "
     "kernel-height and/or kernel-width")
    ("data-type",
     po::value<Type>(&inputType)->default_value(HALF),
     "Type of the input and output data")
    ("input-type",
     po::value<Type>(&inputType),
     "Type of the input data")
    ("output-type",
     po::value<Type>(&outputType),
     "Type of the output data and the parameters")
    ("truncation",
     po::value<ShapeOption<unsigned>>(&truncationOption)->default_value(0),
     "Amount to truncate the start and end of each dimension of the input")
    ("truncation-upper",
     po::value<ShapeOption<unsigned>>(&truncationUpperOption)->default_value(0),
     "Amount to truncate the end of each dimension of the input")
    ("truncation-lower",
     po::value<ShapeOption<unsigned>>(&truncationLowerOption)->default_value(0),
     "Amount to truncate the start of each dimension of the input")
    ("in-dilation",
     po::value<ShapeOption<unsigned>>(&inDilationOption)->default_value(1),
     "Input dilation")
    ("padding",
     po::value<ShapeOption<unsigned>>(&paddingOption)->default_value(0),
     "Amount of zero padding to add to the start and end of each dimension")
    ("padding-upper",
     po::value<ShapeOption<unsigned>>(&paddingUpperOption)->default_value(0),
     "Amount of zero padding to add at the end of each dimension")
    ("padding-lower",
     po::value<ShapeOption<unsigned>>(&paddingLowerOption)->default_value(0),
     "Amount of zero padding to add at the start of each dimension")
    ("flip-input",
     po::value<ShapeOption<bool>>(&flipInputOption)->default_value(false),
     "Whether to flip each input spatial field")
    ("kernel-truncation",
     po::value<ShapeOption<unsigned>>(&kernelTruncationOption)
         ->default_value(0),
     "Amount to truncate the start and end of each dimension of the kernel")
    ("kernel-truncation-upper",
     po::value<ShapeOption<unsigned>>(&kernelTruncationUpperOption)
         ->default_value(0),
     "Amount to truncate the end of each dimension of the kernel")
    ("kernel-truncation-lower",
     po::value<ShapeOption<unsigned>>(&kernelTruncationLowerOption)
         ->default_value(0),
     "Amount to truncate the start of each dimension of the kernel")
    ("kernel-dilation",
     po::value<ShapeOption<unsigned>>(&kernelDilationOption)
         ->default_value(1),
     "Kernel dilation")
    ("kernel-padding",
     po::value<ShapeOption<unsigned>>(&kernelPaddingOption)
         ->default_value(0),
     "Amount of zero kernel padding to add at the start and end of each"
     "dimension")
    ("kernel-padding-upper",
     po::value<ShapeOption<unsigned>>(&kernelPaddingUpperOption)
         ->default_value(0),
     "Amount of zero kernel padding to add at the start of each dimension")
    ("kernel-padding-lower",
     po::value<ShapeOption<unsigned>>(&kernelPaddingLowerOption)
         ->default_value(0),
     "Amount of zero kernel padding to add at the end of each dimension")
    ("flip-kernel",
     po::value<ShapeOption<bool>>(&flipKernelOption)->default_value(false),
     "Whether to flip each kernel spatial field")
    ("output-truncation",
     po::value<ShapeOption<unsigned>>(&outputTruncationOption)
         ->default_value(0),
     "Number of output elements to truncate")
    ("output-truncation-upper",
     po::value<ShapeOption<unsigned>>(&outputTruncationUpperOption)
         ->default_value(0),
     "Number of output elements to truncate at the end of each dimension")
    ("output-truncation-lower",
     po::value<ShapeOption<unsigned>>(&outputTruncationLowerOption)
         ->default_value(0),
     "Number of output elements to truncate at the start of each dimension")
    ("stride",
     po::value<ShapeOption<unsigned>>(&strideOption)->default_value(1),
     "Stride")
    ("output-padding",
     po::value<ShapeOption<unsigned>>(&outputPaddingOption)->default_value(0),
     "Number of output elements to truncate")
    ("output-padding-upper",
     po::value<ShapeOption<unsigned>>(&outputPaddingUpperOption)
         ->default_value(0),
     "Number of output elements to truncate at the end of each dimension")
    ("output-padding-lower",
     po::value<ShapeOption<unsigned>>(&outputPaddingLowerOption)
         ->default_value(0),
     "Number of output elements to truncate at the start of each dimension")
    ("alpha",
     po::value<double>(&alpha),
     "Fraction of input channels to use for Lower Frequency")
    ("alpha-in",
     po::value<double>(&alphaIn),
     "Fraction of input channels to use for Lower Frequency")
    ("alpha-out",
     po::value<double>(&alphaOut),
     "Fraction of output channels to use for Lower Frequency")
    ("single-phase",
     po::value<Pass>(&pass)->default_value(pass),
     "Run phase all | fwd | bwd | wu")
    ("plan-only", "Only plan the requested passes, don't build or run a graph")
    ("inference-only", "Benchmark inference only")
    ("tolerance", po::value<double>(&relativeTolerance),
     "Relative tolerance to use when validating results against the reference "
     "model")
    ("absolute-tolerance",
     po::value<double>(&absoluteTolerance),
     "Absolute tolerance to use when validating results against the reference "
     "model")
    ("ipus",
     po::value<unsigned>(&numIPUs)->default_value(numIPUs),
     "Number of IPUs")
    ("tiles-per-ipu", po::value(&tilesPerIPU),
     "Number of tiles per IPU")
    ("workers-per-tile",
     po::value<unsigned>(),
     "Number of worker contexts per tile")
    ("batch-size",
     po::value<unsigned>(&batchSize)->default_value(1),
     "Batch size")
    ("conv-groups",
     po::value<unsigned>(&numConvGroups)->default_value(1),
     "Number of convolution groups in grouped convolution")
    ("report-plan", po::value<bool>(&reportPlan)->default_value(false),
     "Display plan")
    ("report-var-storage",
     po::value<bool>(&reportVarStorage)->default_value(false),
     "Report variable storage information")
    ("remap-output-tensor",
     po::value<bool>(&remapOutputTensor)->default_value(false),
     "Remap output tensor if layout is detected to be poor")
    ("pool-optimize-for-speed",
     po::value<bool>(&poolOptimizeForSpeed)->default_value(false),
     "Optimize any pooling operation for speed, not memory")
    ("options", po::value<std::string>(&multiConvOptionsString),
    "Options to use for the multi-convolution, specified as a JSON string, "
    "e.g. {\"key\":\"value\"}")
    ;

  // 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\n";
      std::cout << "A multi-dimensional shape can be specified using a brace "
                   "enclosed comma\n"
                   "separated list, for example --stride={1,2}. You may also "
                   "specify a single\n"
                   "number without braces in which case that value is used for "
                   "each dimension,\n"
                   "for example --stride=2\n";
      return 1;
    }
    po::notify(vm);
  } catch (std::exception &e) {
    std::cerr << "error: " << e.what() << "\n";
    return 1;
  }
  OptionFlags poolOptions = {
      {"optimizeForSpeed", poolOptimizeForSpeed ? "true" : "false"}};
  if (alpha != 1.0) {
    alphaIn = alphaOut = alpha;
  }

  // alpha parameter range checking
  if ((alphaIn < 0.0) || (alphaOut < 0.0) || (alphaIn > 1.0) ||
      (alphaOut > 1.0)) {
    std::cerr << "error: alphaIn (" << alphaIn << ") and alphaOut (" << alphaOut
              << ") must be within the interval [0, 1]\n";
    return 1;
  }
  auto &inputFieldShape = inputFieldShapeOption.val;
  const auto numFieldDims = inputFieldShape.size();

  kernelShapeOption.broadcast(numFieldDims);
  auto &kernelShape = kernelShapeOption.val;

  struct UpperLowerOption {
    ShapeOption<unsigned> &lowerOption;
    ShapeOption<unsigned> &upperOption;
    std::string name;
  } upperLowerOptionTriples[] = {
      {paddingLowerOption, paddingUpperOption, "padding"},
      {truncationLowerOption, truncationUpperOption, "truncation"},
      {kernelTruncationLowerOption, kernelTruncationUpperOption,
       "kernel-truncation"},
      {kernelPaddingLowerOption, kernelPaddingUpperOption, "kernel-padding"},
      {outputTruncationLowerOption, outputTruncationUpperOption,
       "output-truncation"},
      {outputPaddingLowerOption, outputPaddingUpperOption, "output-padding"}};
  for (const auto &entry : upperLowerOptionTriples) {
    if (!vm[entry.name].defaulted()) {
      std::string conflictingOptions[] = {entry.name + "-lower",
                                          entry.name + "-upper"};
      for (auto option : conflictingOptions) {
        if (!vm[option].defaulted()) {
          std::cerr << "--" << entry.name << " as well as --";
          std::cerr << option << " set\n";
          return 1;
        }
      }
      entry.lowerOption = vm[entry.name].as<ShapeOption<unsigned>>();
      entry.upperOption = vm[entry.name].as<ShapeOption<unsigned>>();
    }
    entry.lowerOption.broadcast(numFieldDims);
    entry.upperOption.broadcast(numFieldDims);
  }
  auto &truncationLower = truncationLowerOption.val;
  auto &truncationUpper = truncationUpperOption.val;
  auto &paddingLower = paddingLowerOption.val;
  auto &paddingUpper = paddingUpperOption.val;

  auto &outputTruncationLower = outputTruncationLowerOption.val;
  auto &outputTruncationUpper = outputTruncationUpperOption.val;
  auto &outputPaddingLower = outputPaddingLowerOption.val;
  auto &outputPaddingUpper = outputPaddingUpperOption.val;

  auto &kernelTruncationLower = kernelTruncationLowerOption.val;
  auto &kernelTruncationUpper = kernelTruncationUpperOption.val;
  auto &kernelPaddingLower = kernelPaddingLowerOption.val;
  auto &kernelPaddingUpper = kernelPaddingUpperOption.val;

  inDilationOption.broadcast(numFieldDims);
  auto &inDilation = inDilationOption.val;
  flipInputOption.broadcast(numFieldDims);
  auto &flipInput = flipInputOption.val;

  kernelDilationOption.broadcast(numFieldDims);
  auto &kernelDilation = kernelDilationOption.val;
  flipKernelOption.broadcast(numFieldDims);
  auto &flipKernel = flipKernelOption.val;

  strideOption.broadcast(numFieldDims);
  auto &stride = strideOption.val;
  const auto fwdOutChans = fwdOutChansPerConvGroup * numConvGroups;

  const bool planOnly = vm.count("plan-only");
  const bool inferenceOnly = vm.count("inference-only");
  const bool ignoreData = vm.count("ignore-data");

  bool doFwdPass = pass == Pass::ALL || pass == Pass::FWD;
  bool doBwdPass = !inferenceOnly && (pass == Pass::ALL || pass == Pass::BWD);
  bool doWuPass = !inferenceOnly && (pass == Pass::ALL || pass == Pass::WU);

  if ((vm["output-type"].empty() != vm["input-type"].empty()) ||
      (!vm["data-type"].defaulted() && !vm["output-type"].empty())) {
    throw poputil::poplibs_error("Please specify either --data-type OR "
                                 "(--input-type AND --output-type), not both.");
  }
  if (vm["output-type"].empty()) {
    outputType = inputType;
  }

  if (vm["tolerance"].empty()) {
    if (outputType == FLOAT) {
      relativeTolerance = FLOAT_REL_TOL;
    } else {
      relativeTolerance = HALF_REL_TOL;
    }
  }
  if (vm["absolute-tolerance"].empty()) {
    if (outputType == FLOAT) {
      absoluteTolerance = FLOAT_ABS_TOL;
    } else {
      absoluteTolerance = HALF_ABS_TOL;
    }
  }

  OptionFlags multiConvOptions;
  if (!multiConvOptionsString.empty()) {
    poplar::readJSON(multiConvOptionsString, multiConvOptions);
  }

  auto dev = [&]() -> TestDevice {
    if (isIpuModel(deviceType)) {
      // When running on the IPU model we apply global exchange constraints,
      // which is why we create the device from the model here and not using
      // the normal createTestDevice factory function.
      IPUModel ipuModel(deviceTypeToIPUName(deviceType));
      ipuModel.numIPUs = numIPUs;
      if (vm.count("profile") || profileDir) {
        ipuModel.compileIPUCode = true;
      }
      if (vm.count("workers-per-tile"))
        ipuModel.numWorkerContexts = vm["workers-per-tile"].as<unsigned>();
      if (tilesPerIPU)
        ipuModel.tilesPerIPU = *tilesPerIPU;
      addGlobalExchangeConstraints(ipuModel);
      setGlobalSyncLatency(ipuModel);
      return ipuModel.createDevice();
    } else {
      if (tilesPerIPU)
        return createTestDevice(deviceType, numIPUs, *tilesPerIPU);
      else
        return createTestDeviceFullSize(deviceType, numIPUs);
    }
  }();

  Graph graph(dev.getTarget());
  popops::addCodelets(graph);
  poplin::addCodelets(graph);
  popnn::addCodelets(graph);
  const poplin::ConvParams::InputTransform inputTransform{
      truncationLower, truncationUpper, inDilation,
      paddingLower,    paddingUpper,    flipInput};
  const poplin::ConvParams::InputTransform kernelTransform{
      kernelTruncationLower, kernelTruncationUpper, kernelDilation,
      kernelPaddingLower,    kernelPaddingUpper,    flipKernel};
  const poplin::ConvParams::OutputTransform outputTransform{
      outputTruncationLower, outputTruncationUpper, stride, outputPaddingLower,
      outputPaddingUpper};
  const OctConvUserParams userParams{inputType,
                                     outputType,
                                     batchSize,
                                     inputFieldShape,
                                     kernelShape,
                                     fwdInChansPerConvGroup,
                                     fwdOutChansPerConvGroup,
                                     numConvGroups};
  OptionFlags convOptions;
  convOptions.set(
      {{"remapOutputTensor", remapOutputTensor ? "true" : "false"}});
  auto fwdOptions = convOptions;
  fwdOptions.set("pass", inferenceOnly ? "INFERENCE_FWD" : "TRAINING_FWD");
  auto bwdOptions = convOptions;
  bwdOptions.set("pass", "TRAINING_BWD");
  auto wuOptions = convOptions;
  wuOptions.set("pass", "TRAINING_WU");
  std::vector<ConvParameters> octConvParams;
  std::vector<std::shared_ptr<OctConvData>> octInData;
  std::vector<std::shared_ptr<OctConvData>> octOutData;
  std::tie(octConvParams, octInData, octOutData) =
      splitConvParamsByFrequency(userParams, inputTransform, kernelTransform,
                                 outputTransform, alphaIn, alphaOut);
  if (reportPlan) {
    for (unsigned i = 0; i < octConvParams.size(); ++i) {
      const auto outFieldShape =
          octConvParams[i].fwdParams.getOutputFieldShape();
      std::cout << "Convolution parameters for " << i
                << "'th :\n"
                   " Batch size: "
                << octConvParams[i].fwdParams.batchSize
                << "\n"
                   " Kernel:"
                << octConvParams[i].fwdParams.kernelShape
                << "\n"
                   " Stride:"
                << octConvParams[i].fwdParams.outputTransform.stride
                << "\n"
                   " Padding Lower: "
                << octConvParams[i].fwdParams.inputTransform.paddingLower
                << "\n"
                   " Padding Upper: "
                << octConvParams[i].fwdParams.inputTransform.paddingUpper
                << "\n"
                   " Group size: "
                << octConvParams[i].fwdParams.numConvGroups
                << "\n"
                   " Input: "
                << octConvParams[i].fwdParams.inputChannelsPerConvGroup << "x"
                << octConvParams[i].fwdParams.inputFieldShape
                << "\n"
                   " Output: "
                << octConvParams[i].fwdParams.outputChannelsPerConvGroup << "x"
                << outFieldShape << "\n";

      if (doFwdPass) {
        std::cout << "Forward plan:\n";
        poplin::reportPlanInfo(std::cout, graph, octConvParams[i].fwdParams,
                               fwdOptions, &cache);
        std::cout << "Forward FLOPs: "
                  << getFwdFlops(octConvParams[i].fwdParams) << "\n";
      }

      if (doBwdPass) {
        std::cout << "Backward plan:\n";
        poplin::reportPlanInfo(std::cout, graph, octConvParams[i].bwdParams,
                               bwdOptions, &cache);
        std::cout << "Backward FLOPs: "
                  << getBwdFlops(octConvParams[i].bwdParams) << "\n";
      }

      if (doWuPass) {
        std::cout << "WU plan:\n";
        poplin::reportWeightUpdatePlanInfo(
            std::cout, graph, octConvParams[i].fwdParams, wuOptions, &cache);
        std::cout << "WU FLOPs: " << getWuFlops(octConvParams[i].fwdParams)
                  << "\n";
      }
    }
  }
  if (planOnly) {
    return 0;
  }
  // Create OctConv inputs and outputs for target as well as model
  createOctConvTensors(graph, cache, userParams, inputTransform,
                       kernelTransform, outputTransform, fwdOptions, bwdOptions,
                       octInData, multiConvOptions);
  std::vector<Tensor> octConvInput;
  for (auto octConvIn : octInData) {
    octConvInput.push_back(octConvIn->tensor);
  }

  // Create and individual convolution inputs
  std::vector<Tensor> weights, zDeltas;
  std::tie(weights, zDeltas) = createConvInputTensors(
      graph, cache, octConvParams, fwdOptions, bwdOptions,
      (doBwdPass || doWuPass), multiConvOptions);
  auto fwdProg = Sequence();

  // Convert OctConv inputs to individual convolution prevAct tensors.
  auto prevAct = preProcess(graph, octConvParams, fwdProg, poolOptions);
  auto fwdArgs = getConvolutionArguments(prevAct, weights, false, octConvParams,
                                         fwdOptions);
  auto nextAct = poplin::multiconv::convolution(
      graph, fwdArgs, false, fwdProg, "fwd", multiConvOptions, &cache);
  std::vector<Tensor> octConvOutput;
  auto revProg = Sequence();
  if (doFwdPass) {
    // Generate OctConv outputs
    postProcess(graph, octConvParams, nextAct, fwdProg);
    for (auto octConvOut : octOutData) {
      octConvOutput.push_back(octConvOut->tensor);
    }
  } else {
    fwdProg = Sequence();
    // Generate convolution inputs
    prevAct = preProcess(graph, octConvParams, revProg, poolOptions);
  }
  const auto learningRate = 0.05;
  std::vector<Tensor> prevDeltas;
  if (doBwdPass) {
    auto bwdArgs = getConvolutionArguments(zDeltas, weights, true,
                                           octConvParams, bwdOptions);
    prevDeltas = poplin::multiconv::convolution(
        graph, bwdArgs, true, revProg, "bwd", multiConvOptions, &cache);
  }
  if (doWuPass) {
    auto scale = graph.addConstant(weights[0].elementType(), {}, -learningRate);
    graph.setTileMapping(scale, 0);
    auto wuArgs = getConvolutionWeightUpdateArguments(
        prevAct, zDeltas, weights, scale, octConvParams, wuOptions);
    poplin::multiconv::convolutionWeightUpdate(graph, wuArgs, revProg, "wu",
                                               multiConvOptions, &cache);
  }
  if (doPrintTensors) {
    if (doFwdPass || doWuPass) {
      printTensors("octConvInput", octConvParams, octConvInput, fwdProg);
      printTensors("prevAct", octConvParams, prevAct, fwdProg);
    }
    printTensors("weights", octConvParams, weights, fwdProg);
    if (doFwdPass) {
      printTensors("nextAct", octConvParams, nextAct, fwdProg);
      printTensors("octConvOutput", octConvParams, octConvOutput, fwdProg);
    }
    if (doBwdPass || doWuPass) {
      printTensors("zDeltas", octConvParams, zDeltas, revProg);
    }
    if (doBwdPass) {
      printTensors("prevDeltas", octConvParams, prevDeltas, revProg);
    }
  }

  Sequence uploadProg, downloadProg;
  std::vector<std::pair<std::string, char *>> tmap;
  auto rawHostOctConvInput =
      allocateOctHostMemory(graph, "octConvInput", octConvParams, octConvInput,
                            uploadProg, boost::none, tmap);
  auto rawHostWeights = allocateOctHostMemory(
      graph, "weights", octConvParams, weights, uploadProg, downloadProg, tmap);
  auto rawHostOctConvOutput =
      allocateOctHostMemory(graph, "octConvOutput", octConvParams,
                            octConvOutput, boost::none, downloadProg, tmap);
  std::vector<std::unique_ptr<char[]>> rawHostZDeltas;
  std::vector<std::unique_ptr<char[]>> rawHostPrevDeltas;
  if (doBwdPass || doWuPass) {
    rawHostZDeltas =
        allocateOctHostMemory(graph, "zDeltas", octConvParams, zDeltas,
                              uploadProg, downloadProg, tmap);
  }
  if (doBwdPass) {
    rawHostPrevDeltas =
        allocateOctHostMemory(graph, "prevDeltas", octConvParams, prevDeltas,
                              boost::none, downloadProg, tmap);
  }

  std::vector<Program> programs;
  const auto fwdProgIndex = programs.size(); // 0
  programs.push_back(std::move(fwdProg));
  const auto revProgIndex = programs.size(); // 1
  programs.push_back(std::move(revProg));
  const auto uploadProgIndex = programs.size(); // 2
  programs.push_back(std::move(uploadProg));
  const auto downloadProgIndex = programs.size(); // 3
  programs.push_back(std::move(downloadProg));

  OptionFlags engineOptions;
  if (vm.count("profile") || profileDir) {
    engineOptions.set("debug.instrumentCompute", "true");
    if (profileDir) {
      engineOptions.set("autoReport.all", "true");
      engineOptions.set("autoReport.directory", *profileDir);
    }
  }

  Engine engine(graph, std::move(programs), engineOptions);

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

  attachStreams(engine, tmap);

  auto hostPrevAct = createMultiArrayInput(octConvParams, false);
  auto hostWeights = createMultiArrayWeights(octConvParams, kernelShape);
  auto hostZDeltas = createMultiArrayInput(octConvParams, true);
  auto hostPrevDeltas = createMultiArrayInput(octConvParams, false);
  Array1d hostBiases(boost::extents[fwdOutChans]);
  std::fill(hostBiases.begin(), hostBiases.end(), 0);
  std::mt19937 randomEngine;
  auto target = graph.getTarget();
  writeRandomValues(target, inputType, octInData, -1.0, +5.0, randomEngine);
  copy(target, inputType, octInData, rawHostOctConvInput);
  writeRandomValues<double, 4>(target, inputType, hostWeights, -1.0, +7.0,
                               randomEngine);
  copy<double, 4>(target, inputType, hostWeights, rawHostWeights);
  if (doBwdPass || doWuPass) {
    writeRandomValues<double, 3>(target, inputType, hostZDeltas, -3.0, +7.0,
                                 randomEngine);
    copy<double, 3>(target, inputType, hostZDeltas, rawHostZDeltas);
  }

  // Run all configured convolution passes and pre/post processings.
  dev.bind([&](const Device &d) {
    engine.load(d);
    if (!ignoreData) {
      engine.run(uploadProgIndex);
    }
    if (doFwdPass) {
      engine.run(fwdProgIndex);
    }
    if (doBwdPass || doWuPass) {
      engine.run(revProgIndex);
    }
    if (!ignoreData) {
      engine.run(downloadProgIndex);
    }
  });

  // Validate against a reference model.
  bool matchesModel = true;
  if (!ignoreData) {
    PreProcessModel(octConvParams, hostPrevAct);
    auto modelNextAct = createMultiArrayOutput(octConvParams);
    convolutionForwardModel(octConvParams, hostPrevAct, hostWeights,
                            modelNextAct, hostBiases);
    if (doFwdPass) {
      copy(target, outputType, rawHostOctConvOutput, octOutData);
      PostProcessModel(octConvParams, numConvGroups, batchSize, modelNextAct);
      matchesModel &= multipleCheckIsClose(octOutData, relativeTolerance,
                                           absoluteTolerance);
    }
  }

  if (doBwdPass || doWuPass) {
    // Save biases and weights before they are overwritten with device contents.
    auto modelBiases = hostBiases;
    auto modelWeights = hostWeights;
    copy<double, 3>(target, outputType, rawHostZDeltas, hostZDeltas);
    if (doBwdPass) {
      // Copy result of Backward Pass from device
      copy<double, 3>(target, outputType, rawHostPrevDeltas, hostPrevDeltas);
    }
    if (!ignoreData) {
      // Validate against a reference model.
      if (doBwdPass) {
        auto modelPrevDeltas = createMultiArrayInput(octConvParams, false);
        convolutionBackwardModel(octConvParams, hostZDeltas, modelWeights,
                                 modelPrevDeltas);
        matchesModel &= multipleCheckIsClose<3>(
            octConvParams, hostPrevDeltas, modelPrevDeltas, relativeTolerance,
            absoluteTolerance);
      }
      if (doWuPass) {
        // Copy result of Weight Update Pass from device
        copy<double, 4>(target, inputType, rawHostWeights, hostWeights);
        weightUpdateModel(octConvParams, hostPrevAct, hostZDeltas, modelWeights,
                          learningRate, modelBiases);
        matchesModel &=
            multipleCheckIsClose<4>(octConvParams, hostWeights, modelWeights,
                                    relativeTolerance, absoluteTolerance);
      }
    }
  }
  if (!ignoreData && !matchesModel) {
    std::cerr << "Validation failed\n";
    return 1;
  }
  if (vm.count("profile")) {
    auto reportOptions = OptionFlags{{"showExecutionSteps", "true"}};
    if (reportVarStorage) {
      reportOptions.set("showVarStorage", "true");
    }
    engine.printProfileSummary(std::cout, reportOptions);
  }
  return 0;
} catch (const poplar::graph_memory_allocation_error &e) {
  std::cerr << e.what() << std::endl;

  // this exit code has been marked as a "skip" for ctest.
  return 77;
}


//   Copyright 2015 Patrick Putnam
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
#ifndef CLOTHO_FITNESS_SELECTION_GENERATOR_HPP_
#define CLOTHO_FITNESS_SELECTION_GENERATOR_HPP_

#include <boost/property_tree/ptree.hpp>

#include <boost/random/discrete_distribution.hpp>
#include <boost/random/uniform_int_distribution.hpp>

#include <vector>

#include "clotho/data_spaces/selection/selection_generator.hpp"

namespace clotho {
namespace genetics {

template < class GeneticSpaceType >
struct fitness_selection {};

template < class RNG, class GeneticSpaceType >
class SelectionGenerator< RNG, fitness_selection< GeneticSpaceType > > {
public:
    typedef RNG                                                                 random_engine_type;
    typedef GeneticSpaceType                                                    genetic_space_type;
    typedef typename genetic_space_type::individual_id_type                     individual_id_type;

    typedef std::vector< std::pair< individual_id_type, individual_id_type > >  mate_pair_vector;
    typedef typename mate_pair_vector::iterator                                 parent_iterator;
    typedef typename mate_pair_vector::const_iterator                           const_parent_iterator;

    typedef typename genetic_space_type::fitness_score_type                     fitness_score_type;
    typedef typename genetic_space_type::fitness_scores                         fitness_scores;

    typedef boost::random::discrete_distribution< individual_id_type, fitness_score_type >  discrete_distribution_type;
    typedef boost::random::uniform_int_distribution< individual_id_type >       uniform_distribution_type;

    SelectionGenerator( random_engine_type * rng, boost::property_tree::ptree & config ) :
        m_rand( rng )
    { }

    void update( genetic_space_type * parents, unsigned int count ) {

        // pre-scan fitness
        typedef typename genetic_space_type::fitness_iterator iterator;

        iterator first = parents->fitness_begin(), last = parents->fitness_end();
        bool constant_fitness = true;
        unsigned int N = 0;
        if( first != last ) {
            typename genetic_space_type::fitness_score_type prev = *first++;
            ++N;
            while( constant_fitness && first != last ) {
                constant_fitness = (prev == *first++);
                ++N;
            }
        }

        if( constant_fitness ) {
            uniform_distribution_type uni( 0, N - 1 );
            generate( uni, count );
        } else {
            discrete_distribution_type disc( parents->fitness_begin(), parents->fitness_end() );
            generate( disc, count );
        }
    }

    template < class DistributionType >
    void generate( DistributionType & dist, size_t count ) {
        m_pairs.clear();

        while( count-- ) {
            individual_id_type  id0 = dist( *m_rand );
            individual_id_type  id1 = dist( *m_rand );

            m_pairs.push_back( std::make_pair( id0, id1 ) );
        }
    }

    parent_iterator begin() {
        return m_pairs.begin();
    }

    parent_iterator end() {
        return m_pairs.end();
    }

    const_parent_iterator begin() const {
        return m_pairs.begin();
    }

    const_parent_iterator end() const {
        return m_pairs.end();
    }

    mate_pair_vector & getMatePairs() {
        return m_pairs;
    }

    size_t size() const {
        return m_pairs.size();
    }

    virtual ~SelectionGenerator() {}

protected:
    random_engine_type  * m_rand;
    mate_pair_vector    m_pairs;
};

}   // namespace genetics
}   // namespace clotho

#endif  // CLOTHO_FITNESS_SELECTION_GENERATOR_HPP_


/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center,
Division of Medical and Biological Informatics.
All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#include <mitkImageCast.h>
#include <mitkOdfImage.h>
#include <mitkTensorImage.h>
#include <mitkFiberBundle.h>
#include <itkGibbsTrackingFilter.h>
#include <itkDiffusionTensor3D.h>
#include <itkShCoefficientImageImporter.h>
#include <mitkImageToItk.h>
#include <mitkIOUtil.h>
#include "mitkCommandLineParser.h"
#include <boost/algorithm/string.hpp>
#include <itkFlipImageFilter.h>
#include <mitkCoreObjectFactory.h>

template<int shOrder>
typename itk::ShCoefficientImageImporter< float, shOrder >::OdfImageType::Pointer TemplatedConvertShCoeffs(mitk::Image* mitkImg, int toolkit, bool noFlip = false)
{
    typedef itk::ShCoefficientImageImporter< float, shOrder > FilterType;
    typedef mitk::ImageToItk< itk::Image< float, 4 > > CasterType;
    CasterType::Pointer caster = CasterType::New();
    caster->SetInput(mitkImg);
    caster->Update();
    itk::Image< float, 4 >::Pointer itkImage = caster->GetOutput();
    typename FilterType::Pointer filter = FilterType::New();

    if (noFlip)
    {
        filter->SetInputImage(itkImage);
    }
    else
    {
        std::cout << "Flipping image";
        itk::FixedArray<bool, 4> flipAxes;
        flipAxes[0] = true;
        flipAxes[1] = true;
        flipAxes[2] = false;
        flipAxes[3] = false;
        itk::FlipImageFilter< itk::Image< float, 4 > >::Pointer flipper = itk::FlipImageFilter< itk::Image< float, 4 > >::New();
        flipper->SetInput(itkImage);
        flipper->SetFlipAxes(flipAxes);
        flipper->Update();
        itk::Image< float, 4 >::Pointer flipped = flipper->GetOutput();
        itk::Matrix< double,4,4 > m = itkImage->GetDirection(); m[0][0] *= -1; m[1][1] *= -1;
        flipped->SetDirection(m);

        itk::Point< float, 4 > o = itkImage->GetOrigin();
        o[0] -= (flipped->GetLargestPossibleRegion().GetSize(0)-1);
        o[1] -= (flipped->GetLargestPossibleRegion().GetSize(1)-1);
        flipped->SetOrigin(o);
        filter->SetInputImage(flipped);
    }

    switch (toolkit)
    {
    case 0:
        filter->SetToolkit(FilterType::FSL);
        break;
    case 1:
        filter->SetToolkit(FilterType::MRTRIX);
        break;
    default:
        filter->SetToolkit(FilterType::FSL);
    }
    filter->GenerateData();
    return filter->GetOdfImage();
}

/*!
\brief Perform global fiber tractography (Gibbs tractography)
*/
int main(int argc, char* argv[])
{
    mitkCommandLineParser parser;

    parser.setTitle("Gibbs Tracking");
    parser.setCategory("Fiber Tracking and Processing Methods");
    parser.setDescription("Perform global fiber tractography (Gibbs tractography)");
    parser.setContributor("MIC");

    parser.setArgumentPrefix("--", "-");
    parser.addArgument("input", "i", mitkCommandLineParser::InputFile, "Input:", "input image (tensor, ODF or FSL/MRTrix SH-coefficient image)", us::Any(), false);
    parser.addArgument("parameters", "p", mitkCommandLineParser::InputFile, "Parameters:", "parameter file (.gtp)", us::Any(), false);
    parser.addArgument("mask", "m", mitkCommandLineParser::InputFile, "Mask:", "binary mask image");
    parser.addArgument("shConvention", "s", mitkCommandLineParser::String, "SH coefficient:", "sh coefficient convention (FSL, MRtrix)", std::string("FSL"), true);
    parser.addArgument("outFile", "o", mitkCommandLineParser::OutputFile, "Output:", "output fiber bundle (.fib)", us::Any(), false);
    parser.addArgument("noFlip", "f", mitkCommandLineParser::Bool, "No flip:", "do not flip input image to match MITK coordinate convention");

    std::map<std::string, us::Any> parsedArgs = parser.parseArguments(argc, argv);
    if (parsedArgs.size()==0)
        return EXIT_FAILURE;

    std::string inFileName = us::any_cast<std::string>(parsedArgs["input"]);
    std::string paramFileName = us::any_cast<std::string>(parsedArgs["parameters"]);
    std::string outFileName = us::any_cast<std::string>(parsedArgs["outFile"]);

    bool noFlip = false;
    if (parsedArgs.count("noFlip"))
        noFlip = us::any_cast<bool>(parsedArgs["noFlip"]);

    try
    {
        // instantiate gibbs tracker
        typedef itk::Vector<float, ODF_SAMPLING_SIZE>   OdfVectorType;
        typedef itk::Image<OdfVectorType,3>         ItkOdfImageType;
        typedef itk::GibbsTrackingFilter<ItkOdfImageType> GibbsTrackingFilterType;
        GibbsTrackingFilterType::Pointer gibbsTracker = GibbsTrackingFilterType::New();

        // load input image
        mitk::Image::Pointer mitkImage = dynamic_cast<mitk::Image*>(mitk::IOUtil::Load(inFileName)[0].GetPointer());

        // try to cast to Odf image
        if( boost::algorithm::ends_with(inFileName, ".odf") || boost::algorithm::ends_with(inFileName, ".qbi") )
        {
            std::cout << "Loading Odf image ...";
            mitk::OdfImage::Pointer mitkOdfImage = dynamic_cast<mitk::OdfImage*>(mitkImage.GetPointer());
            ItkOdfImageType::Pointer itk_odf = ItkOdfImageType::New();
            mitk::CastToItkImage(mitkOdfImage, itk_odf);
            gibbsTracker->SetOdfImage(itk_odf.GetPointer());
        }
        else if( boost::algorithm::ends_with(inFileName, ".dti") )
        {
            std::cout << "Loading tensor image ...";
            typedef itk::Image< itk::DiffusionTensor3D<float>, 3 >    ItkTensorImage;
            mitk::TensorImage::Pointer mitkTensorImage = dynamic_cast<mitk::TensorImage*>(mitkImage.GetPointer());
            ItkTensorImage::Pointer itk_dti = ItkTensorImage::New();
            mitk::CastToItkImage(mitkTensorImage, itk_dti);
            gibbsTracker->SetTensorImage(itk_dti);
        }
        else if ( boost::algorithm::ends_with(inFileName, ".nii") )
        {
            std::cout << "Loading sh-coefficient image ...";
            int nrCoeffs = mitkImage->GetLargestPossibleRegion().GetSize()[3];
            int c=3, d=2-2*nrCoeffs;
            double D = c*c-4*d;
            int shOrder = 0;
            if (D>0)
            {
                shOrder = (-c+sqrt(D))/2.0;
                if (shOrder<0)
                    shOrder = (-c-sqrt(D))/2.0;
            }
            else if (D==0)
                shOrder = -c/2.0;

            std::cout << "using SH-order " << shOrder;

            int toolkitConvention = 0;

            if (parsedArgs.count("shConvention"))
            {
                std::string convention = us::any_cast<std::string>(parsedArgs["shConvention"]).c_str();

                if ( boost::algorithm::equals(convention, "MRtrix") )
                {
                    toolkitConvention = 1;
                    std::cout << "Using MRtrix style sh-coefficient convention";
                }
                else
                    std::cout << "Using FSL style sh-coefficient convention";
            }
            else
                std::cout << "Using FSL style sh-coefficient convention";

            switch (shOrder)
            {
            case 4:
                gibbsTracker->SetOdfImage(TemplatedConvertShCoeffs<4>(mitkImage, toolkitConvention, noFlip));
                break;
            case 6:
                gibbsTracker->SetOdfImage(TemplatedConvertShCoeffs<6>(mitkImage, toolkitConvention, noFlip));
                break;
            case 8:
                gibbsTracker->SetOdfImage(TemplatedConvertShCoeffs<8>(mitkImage, toolkitConvention, noFlip));
                break;
            case 10:
                gibbsTracker->SetOdfImage(TemplatedConvertShCoeffs<10>(mitkImage, toolkitConvention, noFlip));
                break;
            case 12:
                gibbsTracker->SetOdfImage(TemplatedConvertShCoeffs<12>(mitkImage, toolkitConvention, noFlip));
                break;
            default:
                std::cout << "SH-order " << shOrder << " not supported";
            }
        }
        else
            return EXIT_FAILURE;

        // global tracking
        if (parsedArgs.count("mask"))
        {
            typedef itk::Image<float,3> MaskImgType;
            mitk::Image::Pointer mitkMaskImage = dynamic_cast<mitk::Image*>(mitk::IOUtil::Load(us::any_cast<std::string>(parsedArgs["mask"]))[0].GetPointer());
            MaskImgType::Pointer itk_mask = MaskImgType::New();
            mitk::CastToItkImage(mitkMaskImage, itk_mask);
            gibbsTracker->SetMaskImage(itk_mask);
        }

        gibbsTracker->SetDuplicateImage(false);
        gibbsTracker->SetLoadParameterFile( paramFileName );
//        gibbsTracker->SetLutPath( "" );
        gibbsTracker->Update();

        mitk::FiberBundle::Pointer mitkFiberBundle = mitk::FiberBundle::New(gibbsTracker->GetFiberBundle());
        mitkFiberBundle->SetReferenceGeometry(mitkImage->GetGeometry());

        mitk::IOUtil::Save(mitkFiberBundle, outFileName );
    }
    catch (itk::ExceptionObject e)
    {
        std::cout << e;
        return EXIT_FAILURE;
    }
    catch (std::exception e)
    {
        std::cout << e.what();
        return EXIT_FAILURE;
    }
    catch (...)
    {
        std::cout << "ERROR!?!";
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}


/*
* Copyright 2019 Virtru Corporation
*
* SPDX - License Identifier: BSD-3-Clause-Clear
*
*/
//
//  TDF SDK
//
//  Created by Sujan Reddy on 2019/04/10.
//

#include <openssl/pem.h>
#include <openssl/bio.h>
#include <boost/scope_exit.hpp>
#include <boost/algorithm/string.hpp>

#include <openssl/err.h>
#include <openssl/x509.h>

#include "crypto_utils.h"
#include "tdf_exception.h"
#include "asym_encryption.h"

namespace virtru::crypto {

    constexpr auto kX509CertTag = "BEGIN CERTIFICATE";

    /// Constructor
    AsymEncryption::AsymEncryption(RSA_free_ptr publicKey, size_t keySize)
    : m_publicKey(std::move(publicKey)), m_rsaSize {keySize} { }

    /// Creates an instance of AsymEncryption.
    std::unique_ptr<AsymEncryption> AsymEncryption::create(const std::string& publicKey) {
        RSA_free_ptr publicKeyPtr;
        BIO_free_ptr publicKeyBuffer { BIO_new_mem_buf(publicKey.data(), publicKey.size()) };

        if (!publicKeyBuffer) {
            ThrowOpensslException("Failed to allocate memory for public key.");
        }

        if (boost::contains(publicKey, kX509CertTag)) {

            X509_free_ptr x509Ptr{ PEM_read_bio_X509(publicKeyBuffer.get(), NULL, NULL, NULL) };
            if (!x509Ptr) {
                ThrowOpensslException("Failed to create X509 cert struct.");
            }

            EVP_PKEY_free_ptr evppkeyPtr { X509_get_pubkey(x509Ptr.get()) };
            if (!evppkeyPtr) {
                ThrowOpensslException("Failed to create EVP_PKEY.");
            }

            // Store the public key into RSA struct
            publicKeyPtr.reset(EVP_PKEY_get1_RSA(evppkeyPtr.get()));

        } else {
            // Store the public key into RSA struct
            publicKeyPtr.reset(PEM_read_bio_RSA_PUBKEY(publicKeyBuffer.get(), nullptr, nullptr, nullptr));
        }

        if (!publicKeyPtr) {
            ThrowOpensslException("Failed to create a public key.");
        }

        size_t keySize = RSA_size(publicKeyPtr.get());
        return std::unique_ptr<AsymEncryption>(new AsymEncryption(std::move(publicKeyPtr), keySize));
    }

    /// Performs encryption of the given data.
    void AsymEncryption::encrypt(Bytes data, WriteableBytes& encryptedData) const {

        if (data.size() > std::numeric_limits<int>::max()) {
            ThrowException("Asymmetric encoding input buffer is too big");
        }

        auto size = 0;
        BOOST_SCOPE_EXIT(&size, &encryptedData) {
            encryptedData = encryptedData.first(std::max(0, size));
        }
        BOOST_SCOPE_EXIT_END

        if (static_cast<size_t>(encryptedData.size()) < m_rsaSize) {
            ThrowException("Asymmetric encoding output buffer is too small");
        }

        // NOTE: from https://www.openssl.org/docs/man1.1.1/man3/RSA_public_encrypt.html
        // flen(data size) must be less than RSA_size(rsa) - 11 for the PKCS #1 v1.5 based padding modes,
        // less than RSA_size(rsa) - 41 for RSA_PKCS1_OAEP_PADDING and exactly RSA_size(rsa) for RSA_NO_PADDING .
        if ((Padding::kRsaPkcs1Padding == m_padding) &&
            static_cast<size_t>(data.size()) >= m_rsaSize - 11) {
            ThrowException("Input buffer is too big for provided key.");
        } else if (Padding::kRsaPkcs1OaepPadding == m_padding && static_cast<size_t>(data.size()) >= m_rsaSize - 41) {
            ThrowException("Input buffer is too big for provided key.");
        }

        size = RSA_public_encrypt(static_cast<int>(data.size()),
                                  toUchar(data.data()),
                                  toUchar(encryptedData.data()),
                                  m_publicKey.get(),
                                  static_cast<int>(m_padding));
        if (-1 == size) {
            ThrowOpensslException("Encryption failed using asymmetric encoding.");
        }
    }

    /// Returns minimal buffer size required for output.
    size_t AsymEncryption::getOutBufferSize() const noexcept {
        return m_rsaSize;
    }

    /// Return public key in PEM format.
    std::string AsymEncryption::pemFormat() const {
        BIO_free_ptr bio{BIO_new(BIO_s_mem())};

        if (1 != PEM_write_bio_RSA_PUBKEY(bio.get(), m_publicKey.get())) {
            ThrowOpensslException("Failed to retrieve public key data.");
        }

        // Read the public key from the buffer and put it in the string
        std::string publicKeyPem(BIO_pending(bio.get()), '\0');
        auto readResult = BIO_read(bio.get(), publicKeyPem.data(), static_cast<int>(publicKeyPem.size()));
        if (readResult <= 0) {
            ThrowOpensslException("Failed to read public key data.");
        }

        return publicKeyPem;
    }
}  // namespace virtru::crypto


/* Copyright © 2017 Apple Inc. All rights reserved.
 *
 * Use of this source code is governed by a BSD-3-clause license that can
 * be found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
 */
#include <set>
#include <iostream>
#include <fstream>
#include <chrono>

#include <boost/random.hpp>
#include <boost/integer_traits.hpp>

#include <core/parallel/pthread_tools.hpp>
#include <core/parallel/atomic.hpp>
#include <core/logging/assertions.hpp>
#include <timer/timer.hpp>

#include <core/random/random.hpp>

#ifdef _WIN32
#include <cross_platform/windows_wrapper.hpp>
#include <Wincrypt.h>
#endif

#include <core/util/syserr_reporting.hpp>
#include <core/util/cityhash_tc.hpp>

namespace turi {
  namespace random {

    /** Get as close to a true source of randomness as possible.
     *
     *  nanoseconds clock from program start.  This should be pretty good.
     *
     *  In case subsequent calls on a platform that does not support nanosecond resolution
     *  happen, also increment a base count to make sure that subsequent seeds
     *  are never the same.
     *
     *  hash all these together to get a final seed hash.
     */
    uint64_t pure_random_seed() {
      static auto base_start_time = std::chrono::high_resolution_clock::now();
      static uint64_t base_seed = hash64(time(NULL));
      static atomic<uint64_t> base_count = 0;

      ++base_count;

      auto now = std::chrono::high_resolution_clock::now();

      uint64_t cur_seed = std::chrono::duration_cast<std::chrono::nanoseconds>(now-base_start_time).count();

      return hash64(base_seed, hash64(cur_seed, base_count));
    }

    /**
     * A truely nondeterministic generator
     */
    class nondet_generator {
    public:
      static nondet_generator& global() {
        static nondet_generator global_gen;
        return global_gen;
      }

      typedef size_t result_type;
      BOOST_STATIC_CONSTANT(result_type, min_value =
                            boost::integer_traits<result_type>::const_min);
      BOOST_STATIC_CONSTANT(result_type, max_value =
                            boost::integer_traits<result_type>::const_max);
      result_type min BOOST_PREVENT_MACRO_SUBSTITUTION () const { return min_value; }
      result_type max BOOST_PREVENT_MACRO_SUBSTITUTION () const { return max_value; }

      nondet_generator() {
#ifndef _WIN32
        rnd_dev.open("/dev/urandom", std::ios::binary | std::ios::in);
        ASSERT_TRUE(rnd_dev.good());
#else
        auto rnd_dev_ret = CryptAcquireContext(&rnd_dev,
                                                NULL,
                                                NULL,
                                                PROV_RSA_FULL,
                                                0);
        if(!rnd_dev_ret) {
          auto err_code = GetLastError();
          if(err_code == NTE_BAD_KEYSET) {
            if(!CryptAcquireContext(&rnd_dev,
                                    NULL,
                                    NULL,
                                    PROV_RSA_FULL,
                                    CRYPT_NEWKEYSET)) {
            } else {
              err_code = GetLastError();
              log_and_throw(get_last_err_str(err_code));
            }
          } else {
            log_and_throw(get_last_err_str(err_code));
          }
        }
#endif
      }
      // Close the random number generator
#ifndef _WIN32
      ~nondet_generator() { rnd_dev.close(); }
#else
      ~nondet_generator() { CryptReleaseContext(rnd_dev, 0); }
#endif
      // read a size_t from the source
      result_type operator()() {
        // read a machine word into result
        result_type result(0);
#ifndef _WIN32
        mut.lock();
        ASSERT_TRUE(rnd_dev.good());
        rnd_dev.read(reinterpret_cast<char*>(&result), sizeof(result_type));
        ASSERT_TRUE(rnd_dev.good());
        mut.unlock();
        //        std::cerr << result << std::endl;
        return result;
#else
        mut.lock();
        ASSERT_TRUE(CryptGenRandom(rnd_dev,8,(BYTE *)&result));
        mut.unlock();
        return result;
#endif
      }
    private:
#ifndef _WIN32
      std::ifstream rnd_dev;
#else
      HCRYPTPROV rnd_dev;
#endif
      mutex mut;
    };
    //nondet_generator global_nondet_rng;

    /**
     * This class represents a master registery of all active random
     * number generators
     */
    struct source_registry {
      std::set<generator*> generators;
      generator master;
      mutex mut;

      static source_registry& global() {
        static source_registry registry;
        return registry;
      }
      /**
       * Seed all threads using the default seed
       */
      void seed() {
        get_source().seed();
      }

      /**
       * Seed all threads using the default seed
       */
      void nondet_seed() {
        get_source().nondet_seed();
      }

      /**
       * Seed all threads using the default seed
       */
      void time_seed() {
        get_source().time_seed();
      }

      /**
       *  Seed all threads with a fixed number
       */
      void seed(const size_t number) {
        get_source().seed(number);
      }

      /**
       * Register a source with the registry and seed it based on the
       * master.
       */
      void register_generator(generator* tls_ptr) {
        ASSERT_TRUE(tls_ptr != NULL);
        mut.lock();
        generators.insert(tls_ptr);
        tls_ptr->seed(master);
        // std::cerr << "Generator created" << std::endl;
        // __print_back_trace();
        mut.unlock();
      }

      /**
       * Unregister a source from the registry
       */
      void unregister_source(generator* tls_ptr) {
        mut.lock();
        generators.erase(tls_ptr);
        mut.unlock();
      }
    };
    // source_registry registry;

    //////////////////////////////////////////////////////////////
    /// Pthread TLS code

    /**
     * this function is responsible for destroying the random number
     * generators
     */
    void destroy_tls_data(void* ptr) {
      generator* tls_rnd_ptr =
        reinterpret_cast<generator*>(ptr);
      if(tls_rnd_ptr != NULL) {
        /// TOFIX: This has issues... The global may have been destroyed already
        //source_registry::global().unregister_source(tls_rnd_ptr);
        delete tls_rnd_ptr;
      }
    }

    /**
     * Simple struct used to construct the thread local storage at
     * startup.
     */
    struct tls_key_creator {
      pthread_key_t TLS_RANDOM_SOURCE_KEY;
      tls_key_creator() : TLS_RANDOM_SOURCE_KEY(0) {
        pthread_key_create(&TLS_RANDOM_SOURCE_KEY,
                           destroy_tls_data);
      }
    };
    // This function is to be called prior to any access to the random
    // source
    static pthread_key_t get_random_source_key() {
      static const tls_key_creator key;
      return key.TLS_RANDOM_SOURCE_KEY;
    }
    // This forces __init_keys__ to be called prior to main.
    static TURI_ATTRIBUTE_UNUSED pthread_key_t __unused_init_keys__(get_random_source_key());

  // the combination of the two mechanisms above will force the
  // thread local store to be initialized
  // 1: before main
  // 2: before any use of random by global variables.
  // KNOWN_ISSUE: if a global variable (initialized before main)
  //               spawns threads which then call random. Things explode.

    /////////////////////////////////////////////////////////////
    //// Implementation of header functions

    generator& get_source() {
      // get the thread local storage
      generator* tls_rnd_ptr =
        reinterpret_cast<generator*>
        (pthread_getspecific(get_random_source_key()));
      // Create a tls_random_source if none was provided
      if(tls_rnd_ptr == NULL) {
        tls_rnd_ptr = new generator();
        assert(tls_rnd_ptr != NULL);
        // This will seed it with the master rng
        source_registry::global().register_generator(tls_rnd_ptr);
        pthread_setspecific(get_random_source_key(),
                            tls_rnd_ptr);
      }
      // assert(tls_rnd_ptr != NULL);
      return *tls_rnd_ptr;
    } // end of get local random source

    void seed() { source_registry::global().seed();  }

    void nondet_seed() { source_registry::global().nondet_seed(); }

    void time_seed() { source_registry::global().time_seed(); }

    void seed(const size_t seed_value) {
      source_registry::global().seed(seed_value);
    }

    void generator::nondet_seed() {
      // Get the global nondeterministic random number generator.
      nondet_generator& nondet_rnd(nondet_generator::global());
      mut.lock();
      // std::cerr << "initializing real rng" << std::endl;
      real_rng.seed(static_cast<uint32_t>(nondet_rnd()));
      // std::cerr << "initializing discrete rng" << std::endl;
      discrete_rng.seed(static_cast<uint32_t>(nondet_rnd()));
      // std::cerr << "initializing fast discrete rng" << std::endl;
      fast_discrete_rng.seed(static_cast<uint32_t>(nondet_rnd()));
      mut.unlock();
    }

    void pdf2cdf(std::vector<double>& pdf) {
      double Z = 0;
      for(size_t i = 0; i < pdf.size(); ++i) Z += pdf[i];
      for(size_t i = 0; i < pdf.size(); ++i)
        pdf[i] = pdf[i]/Z + ((i>0)? pdf[i-1] : 0);
    } // end of pdf2cdf

  }; // end of namespace random

};// end of namespace turi


#include "dnn.hpp"
#include "fixedwavefunction.hpp"
#include "hardspherewavefunction.hpp"
#include "inputsorter.hpp"
#include "jastrowmcmillian.hpp"
#include "jastroworion.hpp"
#include "jastrowpade.hpp"
#include "rbmsymmetricwavefunction.hpp"
#include "rbmwavefunction.hpp"
#include "simplegaussian.hpp"
#include "wavefunction.hpp"
#include "wavefunctionpooling.hpp"
#include "wavefunctionproduct.hpp"

#include <Eigen/Dense>
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

namespace py = pybind11;

void init_nn(py::module&);

void init_wavefunction(py::module& main)
{
    auto m  = main.def_submodule("wavefunctions");
    m.doc() = R"doc(
Wavefunctions
-------------

The wavefunction module defines a set of pre-built components that can be used to
set up a variety of different trial wavefunctions.

)doc";

    init_nn(m);

    py::class_<Wavefunction>(m, "Wavefunction", R"doc(
The :class:`Wavefunction` class provides a unified abstraction for all wavefunctions.

All other code expecting a wavefunction instance will take a :class:`Wavefunction` reference.

In this context, a trial wavefunction is defied to be a class that can:

- Have zero or more parameters, :math:`\vec\alpha`
- Evaluate given some system configuraiton :math:`\mathbf{X}`:

.. math::
    \text{Evaluation}_\Psi(\mathbf{X} = \Psi(\mathbf{X})

- Compute the gradient w.r.t. each variational parameter, :math:`\vec{\alpha}`:

.. math::

    \text{Gradient}_\Psi(\mathbf{X}) = \frac{1}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial\vec\alpha}

- Compute the drift force w.r.t. particle :math:`k`'s dimensions ional coordinate :math:`l`:

.. math::

    \text{Drift}_{\Psi(\mathbf{X}), k,l} = \frac{2}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial X_{kl}}

- Compute the laplacian of a system of :math:`N` particles in :math:`D` dimensions:

.. math::

    \text{Laplacian}_\Psi(\mathbf{X}) = \sum_{k=1}^N\sum_{l=1}^D
        \frac{1}{\Psi(\mathbf{X})}\frac{\partial^2 \Psi(\mathbf{X})}{\partial X_{kl}^2}

.. note::
    Wavefunctions are limited to real-valued functions only - there is no
    support for complex valued wavefunctions at this time.

)doc")

        .def("__call__", &Wavefunction::operator(), py::arg("system"), R"doc(
Return the evaluation of the wavefunction for the given system

.. math::

    \text{Evaluation}_\Psi(\mathbf{X} = \Psi(\mathbf{X})
)doc")

        .def("gradient", &Wavefunction::gradient, py::arg("system"), R"doc(
Return the gradient w.r.t. each variational parameter, divided by the evaluation:

.. math::
    \text{Gradient}_\Psi(\mathbf{X}) = \frac{1}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial\vec\alpha}
)doc")

        .def("laplacian", &Wavefunction::laplacian, py::arg("system"), R"doc(

.. math::

    \text{Laplacian}_\Psi(\mathbf{X}) = \sum_{k=1}^N\sum_{l=1}^D
        \frac{1}{\Psi(\mathbf{X})}\frac{\partial^2 \Psi(\mathbf{X})}{\partial X_{kl}^2}
)doc")
        .def("drift_force",
             py::overload_cast<const System&, int, int>(&Wavefunction::drift_force),
             py::arg("system"),
             py::arg("k"),
             py::arg("l"),
             R"doc(

Return the drift force w.r.t. particle :math:`k`'s dimensions ional coordinate :math:`l`:

.. math::

    \text{Drift}_{\Psi(\mathbf{X}), k, l} = \frac{2}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial X_{kl}}
 )doc")

        .def("drift_force",
             py::overload_cast<const System&>(&Wavefunction::drift_force),
             py::arg("system"),
             R"doc(
Return a list of drift forces for all particles and all dimensions.

The list will be one-dimensional, such that
:math:`\text{Drift}_{\Psi(\mathbf{X}), k,l}` is at index ``k * D + l``.
 )doc")
        .def("symmetry_metric",
             &Wavefunction::symmetry_metric,
             py::arg("sampler"),
             py::arg("samples"),
             py::arg("max_permutations") = 100,
             R"doc(
Return an estimate of the symmetry of the wavefunction.

This is defined as follows:

.. math::
    S(\Psi) =  \frac{\int_{-\infty}^\infty \text{d}\mathbf X\left|\frac{1}{n!}
    \sum_{\mathbf\alpha\in \mathcal{P}_n} P_{\mathbf\alpha}\Psi\right|^2}{\int_{-\infty}^\infty\text{d}\mathbf X
    \max_{\mathbf\alpha\in \mathcal{P}} \left|P_{\mathbf\alpha}\Psi\right|^2}

Here :math:`P_\mathbf{\alpha}\Psi` denotes applying the permutation
:math:`\mathbf{\alpha}` to the input system before evaluating the wavefunction,
and :math:`\mathcal{P}_n` is the set of all permutations of :math:`n` particles.

Properties of this metric:

    - It takes values in :math:`[0, 1]`
    - For symmetric wavefunctions, it equals exactly 1
    - For anti-symmetric wavefunctions, it equals exactly 0
    - Any other function evaluates in :math:`(0, 1)`

This integral will be approximated with Monte Carlo integration using samples from the provided
sampling strategy, and using the specified number of samples. While the integral is defied over all
permutations, a maximum of ``max_permutations`` will be used to make it tractable for large ``N``.
 )doc")
        .def_property(
            "parameters",
            py::overload_cast<>(&Wavefunction::get_parameters, py::const_),
            py::overload_cast<const RowVector&>(&Wavefunction::set_parameters),
            R"doc(
List of all variational parameters. Supports both read and write.
)doc");

    py::class_<FixedWavefunction, Wavefunction>(m, "FixedWavefunction", R"doc(
Wrapper for any wavefunction that fixes the variational parameters.

This means that if this wavefunction is used in a optimization call, its
parameters will not change. This is useful if parts of a wavefunction should be
held constant, while another is allowed to be optimized.
)doc")
        .def(py::init<Wavefunction&>(), py::arg("wavefunction"));

    py::class_<WavefunctionProduct, Wavefunction>(m, "WavefunctionProduct", R"doc(
Defines a wavefunction that acts like the product of two others.

All derivatives will be suitably derived, so this is a simple way to produce
compound expressions.
)doc")
        .def(py::init<Wavefunction&, Wavefunction&>(),
             py::arg("Psi_1"),
             py::arg("Psi_2"));

    py::class_<SumPooling, Wavefunction>(m, "SumPooling", R"doc(
This wavefunction expects a wavefunction of two particles,
:math:`f(\mathbf{x}_1, \mathbf{x}_2)`, and represents the following compound
expression:

.. math::
    \Psi(\mathbf{X}) = \sum_{i \neq j}^N f(\mathbf{X}_i, \mathbf{X}_j)

This is guaranteed to produce a permutation symmetric wavefunction, given any suitable
inner wavefunction :math:`f`.
)doc")
        .def(py::init<Wavefunction&>());

    py::class_<InputSorter, Wavefunction>(m, "InputSorter")
        .def(py::init<Wavefunction&>());

    py::class_<SimpleGaussian, Wavefunction>(m, "SimpleGaussian", R"doc(
A product of gaussians:

.. math::

    \Psi(\mathbf{X}) = \prod_{i=1}^N e^{-\alpha ||\mathbf{X}_i||^2}

The only variational parameter is :math:`\alpha`.

For 3D systems, an optional *fixed* parameter :math:`\beta` can be specified, which changes the above definition to:

.. math::

    \Psi(\mathbf{X}) = \prod_{i=1}^N e^{-\alpha\left(X_{i,1}^2 + X_{i,2}^2 + \beta X_{i,3}^2\right)}
)doc")
        .def(py::init<Real, Real>(), py::arg("alpha") = 0.5, py::arg("beta") = 1);

    py::class_<HardSphereWavefunction, SimpleGaussian>(
        m, "HardSphereWavefunction", R"doc(
A product of gaussians and pairwise correlation factors:

.. math::

    \Psi(\mathbf{X}) = \prod_{i=1}^N e^{-\alpha\left(X_{i,1}^2 + X_{i,2}^2 + \beta X_{i,3}^2\right)}
        \prod_{j = i + 1}^{N} \begin{cases} 0 &\text{if}\ \  ||\mathbf{X}_i - \mathbf{X}_j|| \leq a\\
                                            1 - \frac{a}{||\mathbf{X}_i - \mathbf{X}_j||} &\text{otherwise}
                              \end{cases}

Similarly to :class:`SimpleGaussian`, the only variational parameter is
:math:`\alpha`, while the other two parameters :math:`\beta` and :math:`a` are
assumed constant.
)doc")
        .def(py::init<Real, Real, Real>(),
             py::arg("alpha") = 0.5,
             py::arg("beta")  = 1,
             py::arg("a")     = 0);

    py::class_<JastrowPade, Wavefunction>(m, "JastrowPade", R"doc(
A correlation term meant to be suitable for particles in a harmonic oscillator potential with
a repulsive Coulomb force causing interactions:

.. math::

    \Psi(\mathbf{X}) = \prod_{i < j}^N e^{\frac{\alpha r_ij}{1 + \beta r_ij}}

)doc")
        .def(py::init<Real, Real, bool>(),
             py::arg("alpha")             = 0.5,
             py::arg("beta")              = 1,
             py::arg("alpha_is_constant") = true);

    py::class_<JastrowOrion, Wavefunction>(m, "JastrowOrion", R"doc(
A correlation term meant to be suitable for particles in a harmonic oscillator potential with
a repulsive Coulomb force causing interactions:

.. math::

    \Psi(\mathbf{X}) = \prod_{i < j}^N
        e^{-\frac{\beta^2}{2}||\mathbf{X_i}-\mathbf{X_j}||^2 +
            |\beta\gamma|||\mathbf{X_i}-\mathbf{X_j}||}

)doc")
        .def(py::init<Real, Real>(), py::arg("beta") = 1.0, py::arg("gamma") = 0);

    py::class_<JastrowMcMillian, Wavefunction>(m, "JastrowMcMillian")
        .def(py::init<int, Real, Real>(), py::arg("n"), py::arg("beta"), py::arg("L"));

    py::class_<RBMWavefunction, Wavefunction>(m, "RBMWavefunction")
        .def(py::init<int, int, Real, Real>(),
             py::arg("M"),
             py::arg("N"),
             py::arg("sigma2")      = 1,
             py::arg("root_factor") = 1);

    py::class_<RBMSymmetricWavefunction, Wavefunction>(m, "RBMSymmetricWavefunction")
        .def(py::init<int, int, int, Real, Real>(),
             py::arg("M"),
             py::arg("N"),
             py::arg("f"),
             py::arg("sigma2")      = 1,
             py::arg("root_factor") = 1);

    py::class_<Dnn, Wavefunction>(m, "Dnn")
        .def(py::init<>())
        .def("add_layer", &Dnn::addLayer)
        .def_property_readonly("layers", &Dnn::getLayers);
}


#define BOOST_TEST_MODULE "C++ Unit Tests for metaLBM"
#include <boost/test/included/unit_test.hpp>
#include <boost/test/unit_test.hpp>

#include <math.h>

#include "structure.h"

using namespace lbm;

BOOST_AUTO_TEST_SUITE(Structure)

BOOST_AUTO_TEST_CASE(UnrollFor) {
  constexpr ::lbm::LatticeType latticeType = ::lbm::LatticeType::D1Q3;

  typedef Parameters<T, L> Param;

  std::vector<T> fEq(Param::dimQ * Param::lX_g * Param::lY_g * Param::lZ_g, (T)0);

  const int idx = Param::lX_g * Param::lY_g * Param::lZ_g - 1;

  const auto density_ = computeDensity<T, L>(fEq.data(), idx);
  BOOST_CHECK_EQUAL(density_, (T)0);

}

BOOST_AUTO_TEST_SUITE_END()


#ifndef STAN_MATH_PRIM_MAT_FUN_MDIVIDE_LEFT_TRI_LOW_HPP
#define STAN_MATH_PRIM_MAT_FUN_MDIVIDE_LEFT_TRI_LOW_HPP

#include <boost/math/tools/promotion.hpp>
#include <stan/math/prim/mat/fun/Eigen.hpp>
#include <stan/math/prim/mat/fun/mdivide_left_tri.hpp>
#include <stan/math/prim/mat/err/check_square.hpp>
#include <stan/math/prim/mat/err/check_multiplicable.hpp>

namespace stan {
namespace math {

/**
 * Return the result of left dividing the second argument by the
 * first argument.  Calling mdivide_left_tri_low(A,
 * b) with divisor A and dividend
 * b is more arithmetically stable than calling
 * inv(A) * b.
 *
 * @tparam T1 type of divisor matrix
 * @tparam T2 type of dividend matrix
 * @tparam R1 rows of divisor matrix
 * @tparam C1 columns of divisor matrix
 * @tparam R2 rows of dividen matrix
 * @tparam C2 rows of divisor matrix
 * @param A divisor, an invertible square matrix
 * @param b dividend, a matrix or vector with the same number of
 *   rows as the divisor has columns
 * @return left division of b by A
 * @throws std::invalid_argument if the divisor is not square or
 *   the dividend does not have the same number of rows as the
 *   divisor has columns.
 */
template <typename T1, typename T2, int R1, int C1, int R2, int C2>
inline Eigen::Matrix<return_type_t<T1, T2>, R1, C2> mdivide_left_tri_low(
    const Eigen::Matrix<T1, R1, C1> &A, const Eigen::Matrix<T2, R2, C2> &b) {
  check_square("mdivide_left_tri_low", "A", A);
  check_multiplicable("mdivide_left_tri_low", "A", A, "b", b);
  return mdivide_left_tri<Eigen::Lower>(A, b);
}

template <typename T, int R1, int C1>
inline Eigen::Matrix<T, R1, C1> mdivide_left_tri_low(
    const Eigen::Matrix<T, R1, C1> &A) {
  check_square("mdivide_left_tri_low", "A", A);
  return mdivide_left_tri<Eigen::Lower>(A);
}

}  // namespace math
}  // namespace stan
#endif