// Copyright (c) 2017 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE StdOperatorsTest
#include <boost/program_options.hpp>
#include <boost/random.hpp>
#include <cmath>
#include <iostream>
#include <limits>
#include <poplar/Engine.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/Util.hpp>
#include <popops/AllTrue.hpp>
#include <popops/ElementWise.hpp>
#include <popops/Expr.hpp>
#include <popops/ScaledAdd.hpp>
#include <popops/codelets.hpp>
#include <poputil/TileMapping.hpp>
#include <poputil/exceptions.hpp>
#include <random>
#include <stdexcept>
#include <string>

using namespace poplar;
using namespace poplar::program;
using namespace poputil;
using namespace popops;
using namespace poplibs_test::util;
using namespace poplibs_support;
namespace br = boost::random;
namespace pe = popops::expr;

static DeviceType deviceType;

poplar::OptionFlags options;

#define DIM_SIZE 3

#define STRINGIFY(x) #x
#define TOSTRING(x) STRINGIFY(x)

#define CHECK(pred)                                                            \
  do {                                                                         \
    if (!(pred)) {                                                             \
      throw std::runtime_error("failed check '" #pred                          \
                               "' on line " TOSTRING(__LINE__));               \
    }                                                                          \
  } while (false)

#define CHECK_CLOSE(a, b)                                                      \
  do {                                                                         \
    if (!checkIsClose<                                                         \
            typename std::common_type<decltype(a), decltype(b)>::type>(        \
            a, b, 0.01)) {                                                     \
      throw std::runtime_error(                                                \
          "failed close check on line " TOSTRING(__LINE__));                   \
    }                                                                          \
  } while (false)

static Tensor mapUnaryOpTensor(Graph &graph, const Type &type) {
  auto in = graph.addVariable(type, {DIM_SIZE, DIM_SIZE}, "in0");
  mapTensorLinearly(graph, in);

  return in.dimShuffle({1, 0});
}

static std::pair<Tensor, Tensor> mapBinaryOpTensors(Graph &graph,
                                                    const Type &type) {
  auto in1 = graph.addVariable(type, {DIM_SIZE, DIM_SIZE}, "in1");
  mapTensorLinearly(graph, in1);

  auto in2 = graph.addVariable(type, {DIM_SIZE, DIM_SIZE}, "in2");
  mapTensorLinearly(graph, in2);

  return std::make_pair(in1.dimShuffle({1, 0}), in2.dimShuffle({1, 0}));
}

/* Generates a 2D matrix of size DIM_SIZE x DIM_SIZE containing linearly
 * increasing absolute values with a fixed slope and sign alternating sign
 * for each element in a row. The first two elements are +/-Infinity
 */
static void setUnaryOpInput(float hIn[DIM_SIZE][DIM_SIZE]) {
  float val = -100;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      float sign = 1.0 - 2.0 * ((r + c) & 1);
      hIn[r][c] = (val + (r * DIM_SIZE + c) * .1) * sign;
    }
  }
}

static void setUnaryOpInput(int hIn[DIM_SIZE][DIM_SIZE]) {
  int val = -100;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      int sign = 1 - 2 * ((r + c) & 1);
      hIn[r][c] = (val + (r * DIM_SIZE + c)) * sign;
    }
  }
}

static void setUnaryOpInput(bool hIn[DIM_SIZE][DIM_SIZE]) {
  std::mt19937 randomEngine;
  br::uniform_int_distribution<> dist(0, 1);
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hIn[r][c] = dist(randomEngine);
    }
  }
}

static void setUnaryOpInputHalf(float hIn[DIM_SIZE][DIM_SIZE]) {
  float val = -100;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      float sign = 1.0 - 2.0 * ((r + c) & 1);
      hIn[r][c] = (val + (r * DIM_SIZE + c) * .1) * sign;
    }
  }
}

/* Generates two 2D matrix of size DIM_SIZE x DIM_SIZE containing linearly
 * increasing absolute values with a fixed slope and sign alternating sign
 * for each element in a row. The start value at position [0][0] for each
 * of the matrices is different.
 */
static void setBinaryOpInputs(float hIn1[DIM_SIZE][DIM_SIZE],
                              float hIn2[DIM_SIZE][DIM_SIZE],
                              bool isShift = false) {
  float val1 = 1000;
  float val2 = 50;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      float sign1 = (1.0 - 2.0 * ((c + 1) & 1));
      float sign2 = (1.0 - 2.0 * ((r + c) & 1));
      hIn1[r][c] = (val1 + (r * DIM_SIZE + c) * .1) * sign1;
      hIn2[r][c] = (val2 + (r * DIM_SIZE + c) * .1) * sign2;
    }
  }
}

static void setBinaryOpInputsHalf(float hIn1[DIM_SIZE][DIM_SIZE],
                                  float hIn2[DIM_SIZE][DIM_SIZE],
                                  bool isShift = false) {
  float val1 = -100;
  float val2 = 50;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      float sign1 = (1.0 - 2.0 * ((c + 1) & 1));
      float sign2 = (1.0 - 2.0 * ((r + c) & 1));
      hIn1[r][c] = (val1 + (r * DIM_SIZE + c) * .1) * sign1;
      hIn2[r][c] = (val2 + (r * DIM_SIZE + c) * .1) * sign2;
    }
  }
}

/* Generates two 2D matrix of size DIM_SIZE x DIM_SIZE containing
 * boolean values. All combinations of boolean values are produced
 */
static void setBinaryOpInputs(bool hIn1[DIM_SIZE][DIM_SIZE],
                              bool hIn2[DIM_SIZE][DIM_SIZE],
                              bool isShift = false) {
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hIn1[r][c] = r & 1;
      hIn2[r][c] = (r + c) & 1;
    }
  }
}

static void setBinaryOpInputs(int hIn1[DIM_SIZE][DIM_SIZE],
                              int hIn2[DIM_SIZE][DIM_SIZE],
                              bool isShift = false) {
  int val1 = -100;
  int val2 = 59;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hIn1[r][c] = (1 - 2 * (r & 1)) * (r + val1);
      hIn2[r][c] = (1 - 2 * ((r + c) & 1)) * (r + c + val2);

      // Shifting by more than 32 is undefined.
      if (isShift) {
        hIn2[r][c] = hIn2[r][c] % 32;
        if (hIn2[r][c] < 0)
          hIn2[r][c] = -hIn2[r][c];
      }
    }
  }
}

template <typename T> void convertToPositive(T array[DIM_SIZE][DIM_SIZE]) {
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      array[i][j] = std::abs(array[i][j]);
    }
  }
}

template <> void convertToPositive(bool array[DIM_SIZE][DIM_SIZE]) {}

template <> void convertToPositive(float array[DIM_SIZE][DIM_SIZE]) {
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      array[i][j] = std::fabs(array[i][j]);
    }
  }
}

using UnaryOpFn = std::function<Tensor(Graph &, const Tensor &, Sequence &,
                                       const poplar::DebugContext &,
                                       const poplar::OptionFlags &)>;

template <typename T, typename TestT>
void unaryOpTest(const UnaryOpFn &op, const std::function<TestT(T)> &testFn,
                 bool positiveInputs = false, bool disableFpException = false) {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto in = mapUnaryOpTensor(graph, equivalent_device_type<T>().value);
  auto prog = Sequence();

  auto out = op(graph, in, prog, "unaryOp", {});
  graph.createHostWrite("in", in);
  graph.createHostRead("out", out);

  // Some inbuild functions will trigger an float point exception hence need
  // to disable FP exception for specific tests
  if (disableFpException) {
    options.set("debug.floatPointOpException", "false");
  }
  Engine eng(graph, prog, options);
  T hIn[DIM_SIZE][DIM_SIZE];
  T hOut[DIM_SIZE][DIM_SIZE];
  device.bind([&](const Device &d) {
    eng.load(d);

    setUnaryOpInput(hIn);
    if (positiveInputs) {
      convertToPositive(hIn);
    }
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto res = testFn(hIn[i][j]);
      CHECK_CLOSE(res, hOut[i][j]);
    }
  }
}

using BinaryOpFn = std::function<Tensor(
    Graph &, const Tensor &, const Tensor &, Sequence &,
    const poplar::DebugContext &, const poplar::OptionFlags &)>;

template <typename BinaryOpFn> struct BinaryOpFnPtr;

template <typename R, typename... Args>
struct BinaryOpFnPtr<std::function<R(Args...)>> {
  using type = R (*)(Args...);
};

using BinaryOpFnPtr_t = BinaryOpFnPtr<BinaryOpFn>::type;

template <typename T, typename TestT, typename OutT = T>
void binaryOpTest(const BinaryOpFn &op,
                  const std::function<TestT(T, T)> &testFn,
                  bool isShift = false) {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  Tensor in1, in2;
  auto type = equivalent_device_type<T>().value;

  std::tie(in1, in2) = mapBinaryOpTensors(graph, type);

  auto prog = Sequence();

  auto out = op(graph, in1, in2, prog, "binaryOp", {});
  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.createHostRead("out", out);

  Engine eng(graph, prog, options);
  T hIn1[DIM_SIZE][DIM_SIZE], hIn2[DIM_SIZE][DIM_SIZE];
  OutT hOut[DIM_SIZE][DIM_SIZE];
  device.bind([&](const Device &d) {
    eng.load(d);
    setBinaryOpInputs(hIn1, hIn2, isShift);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cout, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      T res = testFn(hIn1[i][j], hIn2[i][j]);
      CHECK_CLOSE(static_cast<TestT>(hOut[i][j]), res);
    }
  }
}

void binaryOpTestHalf(const BinaryOpFn &op,
                      const std::function<float(float, float)> &testFn) {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  const auto &target = device.getTarget();
  popops::addCodelets(graph);

  Tensor in1, in2;
  std::tie(in1, in2) = mapBinaryOpTensors(graph, poplar::HALF);

  auto prog = Sequence();

  auto out = op(graph, in1, in2, prog, "binaryOp", {});
  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.createHostRead("out", out);

  Engine eng(graph, prog, options);
  float hIn1[DIM_SIZE][DIM_SIZE], hIn2[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];
  device.bind([&](const Device &d) {
    eng.load(d);
    setBinaryOpInputsHalf(hIn1, hIn2);
    auto rawBufSize = target.getTypeSize(HALF) * DIM_SIZE * DIM_SIZE;
    std::vector<char> rawIn1(rawBufSize), rawIn2(rawBufSize),
        rawOut(rawBufSize);
    poplar::copyFloatToDeviceHalf(target, &hIn1[0][0], rawIn1.data(),
                                  DIM_SIZE * DIM_SIZE);
    poplar::copyFloatToDeviceHalf(target, &hIn2[0][0], rawIn2.data(),
                                  DIM_SIZE * DIM_SIZE);
    eng.writeTensor("in1", rawIn1.data(), rawIn1.data() + rawIn1.size());
    eng.writeTensor("in2", rawIn2.data(), rawIn2.data() + rawIn2.size());
    eng.run();
    eng.readTensor("out", rawOut.data(), rawOut.data() + rawOut.size());
    poplar::copyDeviceHalfToFloat(target, rawOut.data(), &hOut[0][0],
                                  DIM_SIZE * DIM_SIZE);
  });
  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto res = testFn(hIn1[i][j], hIn2[i][j]);
      CHECK_CLOSE(static_cast<float>(hOut[i][j]), res);
    }
  }
}

void powTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hIn1[DIM_SIZE][DIM_SIZE], hIn2[DIM_SIZE][DIM_SIZE];

  float x = 10.23;
  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c, x += .1f) {
      hIn1[r][c] = x;
      hIn2[r][c] = 2;
    }
  }

  Tensor in1, in2;
  std::tie(in1, in2) = mapBinaryOpTensors(graph, FLOAT);

  auto prog = Sequence();
  auto out = popops::pow(graph, in1, in2, prog);
  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.createHostRead("out", out);

  Engine eng(graph, prog, options);
  float hOut[DIM_SIZE][DIM_SIZE];
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      double res = std::pow(static_cast<double>(hIn1[i][j]),
                            static_cast<double>(hIn2[i][j]));
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

template <typename InType> void selectTest() {
  auto type = equivalent_device_type<InType>().value;
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  InType hIn1[DIM_SIZE][DIM_SIZE], hIn2[DIM_SIZE][DIM_SIZE];
  setBinaryOpInputs(hIn1, hIn2);
  bool hIn3[DIM_SIZE][DIM_SIZE];

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hIn3[r][c] = (c)&0x1;
    }
  }

  Tensor in1, in2;
  std::tie(in1, in2) = mapBinaryOpTensors(graph, type);
  Tensor in3 = mapUnaryOpTensor(graph, BOOL);

  auto prog = Sequence();
  auto out = select(graph, in1, in2, in3, prog);
  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.createHostWrite("in3", in3);
  graph.createHostRead("out", out);

  InType hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.writeTensor("in3", hIn3, &hIn3[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      InType res = hIn3[i][j] ? hIn1[i][j] : hIn2[i][j];
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void selectTestFloatLHSConst() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hLhs = 1.f;
  float hRhs[DIM_SIZE][DIM_SIZE];
  setUnaryOpInput(hRhs);
  bool hPred[DIM_SIZE][DIM_SIZE];

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hPred[r][c] = (c)&0x1;
    }
  }

  Tensor rhs = mapUnaryOpTensor(graph, FLOAT);
  Tensor pred = mapUnaryOpTensor(graph, BOOL);

  auto prog = Sequence();
  auto out = map(
      graph,
      pe::TernaryOp(pe::TernaryOpType::SELECT, pe::Const(hLhs), pe::_1, pe::_2),
      {rhs, pred}, prog);
  graph.createHostWrite("rhs", rhs);
  graph.createHostWrite("pred", pred);
  graph.createHostRead("out", out);

  float hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("rhs", hRhs, &hRhs[DIM_SIZE]);
    eng.writeTensor("pred", hPred, &hPred[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      float res = hPred[i][j] ? hLhs : hRhs[i][j];
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void selectTestFloatRHSConst() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hRhs = 1.f;
  float hLhs[DIM_SIZE][DIM_SIZE];
  setUnaryOpInput(hLhs);
  bool hPred[DIM_SIZE][DIM_SIZE];

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hPred[r][c] = (c)&0x1;
    }
  }

  Tensor lhs = mapUnaryOpTensor(graph, FLOAT);
  Tensor pred = mapUnaryOpTensor(graph, BOOL);

  auto prog = Sequence();
  auto out = map(
      graph,
      pe::TernaryOp(pe::TernaryOpType::SELECT, pe::_1, pe::Const(hRhs), pe::_2),
      {lhs, pred}, prog);
  graph.createHostWrite("lhs", lhs);
  graph.createHostWrite("pred", pred);
  graph.createHostRead("out", out);

  float hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("lhs", hLhs, &hLhs[DIM_SIZE]);
    eng.writeTensor("pred", hPred, &hPred[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      float res = hPred[i][j] ? hLhs[i][j] : hRhs;
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void selectTestFloatLHSAndRHSConst() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hLhs = 2.f;
  float hRhs = 1.f;
  bool hPred[DIM_SIZE][DIM_SIZE];

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hPred[r][c] = (c)&0x1;
    }
  }

  Tensor pred = mapUnaryOpTensor(graph, BOOL);

  auto prog = Sequence();
  auto out = map(graph,
                 pe::TernaryOp(pe::TernaryOpType::SELECT, pe::Const(hLhs),
                               pe::Const(hRhs), pe::_1),
                 {pred}, prog);
  graph.createHostWrite("pred", pred);
  graph.createHostRead("out", out);

  float hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("pred", hPred, &hPred[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      float res = hPred[i][j] ? hLhs : hRhs;
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void selectTestHalfLHSAndRHSConst() {
  auto device = createTestDevice(deviceType);
  auto target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);

  float hLhs = 2.f;
  float hRhs = 1.f;
  bool hPred[DIM_SIZE][DIM_SIZE];

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hPred[r][c] = (c)&0x1;
    }
  }

  Tensor pred = mapUnaryOpTensor(graph, BOOL);

  auto prog = Sequence();
  auto out = map(graph,
                 pe::TernaryOp(pe::TernaryOpType::SELECT, pe::ConstHalf(hLhs),
                               pe::ConstHalf(hRhs), pe::_1),
                 {pred}, prog);
  CHECK(out.elementType() == HALF);
  graph.createHostWrite("pred", pred);
  graph.createHostRead("out", out);

  auto rawBufSize = target.getTypeSize(HALF) * DIM_SIZE * DIM_SIZE;
  std::vector<char> rawOut(rawBufSize);

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("pred", hPred, &hPred[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", rawOut.data(), rawOut.data() + rawOut.size());
  });

  float hOut[DIM_SIZE][DIM_SIZE];
  poplar::copyDeviceHalfToFloat(target, rawOut.data(), &hOut[0][0],
                                DIM_SIZE * DIM_SIZE);
  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      float res = hPred[i][j] ? hLhs : hRhs;
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

template <typename InType> void broadcastSelectorSelectTest(bool inPlace) {
  auto type = equivalent_device_type<InType>().value;
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  InType hIn1[DIM_SIZE][DIM_SIZE];
  InType hIn2[DIM_SIZE][DIM_SIZE];
  setBinaryOpInputs(hIn1, hIn2);
  bool hIn3 = true;

  Tensor in1, in2;
  std::tie(in1, in2) = mapBinaryOpTensors(graph, type);

  Tensor in3 = graph.addVariable(BOOL, {}, "pred"); // selector
  mapTensorLinearly(graph, in3);

  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.setInitialValue(in3, hIn3);

  auto prog = Sequence();

  if (inPlace) {
    selectInPlace(graph, in1, in2, in3, prog);
    graph.createHostRead("out", in1);
  } else {
    auto out = select(graph, in1, in2, in3, prog);
    graph.createHostRead("out", out);
  }

  InType hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      InType res = hIn3 ? hIn1[i][j] : hIn2[i][j];
      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void clampTestFloatMinConst() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hMin = -0.5f;
  float hMax[DIM_SIZE][DIM_SIZE];
  float hIn[DIM_SIZE][DIM_SIZE];
  setUnaryOpInput(hIn);

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hMax[r][c] = 0.5;
    }
  }

  Tensor max = mapUnaryOpTensor(graph, FLOAT);
  Tensor in = mapUnaryOpTensor(graph, FLOAT);

  auto prog = Sequence();
  auto out = map(
      graph,
      pe::TernaryOp(pe::TernaryOpType::CLAMP, pe::_1, pe::Const(hMin), pe::_2),
      {in, max}, prog);
  graph.createHostWrite("max", max);
  graph.createHostWrite("in", in);
  graph.createHostRead("out", out);

  float hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("max", hMax, &hMax[DIM_SIZE]);
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      float res = hIn[i][j];
      if (res < hMin)
        res = hMin;
      if (res > hMax[i][j])
        res = hMax[i][j];

      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void clampTestFloatMaxConst() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hMax = 0.5f;
  float hMin[DIM_SIZE][DIM_SIZE];
  float hIn[DIM_SIZE][DIM_SIZE];
  setUnaryOpInput(hIn);

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hMin[r][c] = 0.5;
    }
  }

  Tensor min = mapUnaryOpTensor(graph, FLOAT);
  Tensor in = mapUnaryOpTensor(graph, FLOAT);

  auto prog = Sequence();
  auto out = map(
      graph,
      pe::TernaryOp(pe::TernaryOpType::CLAMP, pe::_1, pe::_2, pe::Const(hMax)),
      {in, min}, prog);
  graph.createHostWrite("min", min);
  graph.createHostWrite("in", in);
  graph.createHostRead("out", out);

  float hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("min", hMin, &hMin[DIM_SIZE]);
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      float res = hIn[i][j];
      if (res < hMin[i][j])
        res = hMin[i][j];
      if (res > hMax)
        res = hMax;

      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

template <typename InType> void clampTest(bool inPlace) {
  auto type = equivalent_device_type<InType>().value;
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  InType hIn1[DIM_SIZE][DIM_SIZE];
  setUnaryOpInput(hIn1);
  InType hIn2[DIM_SIZE][DIM_SIZE], hIn3[DIM_SIZE][DIM_SIZE];

  for (auto r = 0U; r != DIM_SIZE; ++r) {
    for (auto c = 0U; c != DIM_SIZE; ++c) {
      hIn2[r][c] = static_cast<InType>(-5.5);
      hIn3[r][c] = static_cast<InType>(5.5);
    }
  }

  Tensor in1, in2;
  std::tie(in1, in2) = mapBinaryOpTensors(graph, type);
  Tensor in3 = mapUnaryOpTensor(graph, type);

  auto prog = Sequence();

  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.createHostWrite("in3", in3);

  if (inPlace) {
    clampInPlace(graph, in1, in2, in3, prog);
    graph.createHostRead("out", in1);
  } else {
    auto out = clamp(graph, in1, in2, in3, prog);
    graph.createHostRead("out", out);
  }

  InType hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.writeTensor("in3", hIn3, &hIn3[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      InType res = hIn1[i][j];
      if (res < hIn2[i][j])
        res = hIn2[i][j];
      if (res > hIn3[i][j])
        res = hIn3[i][j];

      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

template <typename InType>
void broadcastClampTest(bool inPlace, size_t dim = DIM_SIZE) {
  auto type = equivalent_device_type<InType>().value;
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  InType hIn1[DIM_SIZE][DIM_SIZE];
  InType hIn2 = InType(-5.5);
  InType hIn3 = InType(5.5);

  if (dim == DIM_SIZE) {
    setUnaryOpInput(hIn1);
  } else if (dim == 1) {
    hIn1[0][0] = 10;
  } else {
    throw std::runtime_error("broadcastClampTest: Unsupported dimension size");
  }

  Tensor in1 = graph.addVariable(type, {dim, dim}, "input"); // source
  Tensor in2 = graph.addVariable(type, {}, "lower");         // lower
  Tensor in3 = graph.addVariable(type, {}, "upper");         // upper

  mapTensorLinearly(graph, in1);
  mapTensorLinearly(graph, in2);
  mapTensorLinearly(graph, in3);

  graph.createHostWrite("in1", in1);
  graph.setInitialValue(in2, hIn2);
  graph.setInitialValue(in3, hIn3);

  auto prog = Sequence();

  if (inPlace) {
    clampInPlace(graph, in1, in2, in3, prog);
    graph.createHostRead("out", in1);
  } else {
    auto out = clamp(graph, in1, in2, in3, prog);
    graph.createHostRead("out", out);
  }

  InType hOut[DIM_SIZE][DIM_SIZE];

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in1", hIn1, &hIn1[dim]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[dim]);
  });

  /* Check result */
  for (auto i = 0U; i < dim; ++i) {
    for (auto j = 0U; j < dim; ++j) {
      InType res = hIn1[i][j];
      if (res < hIn2)
        res = hIn2;
      if (res > hIn3)
        res = hIn3;

      CHECK_CLOSE(hOut[i][j], res);
    }
  }
}

void binaryOutputMapChoiceTest() {
  auto device = createTestDevice(deviceType, 1, 4);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  Tensor in1, in2;
  in1 = graph.addVariable(FLOAT, {2, 2}, "t1");
  in2 = graph.addVariable(FLOAT, {2, 2}, "t2");

  // Tensor in1 all on tile 0
  graph.setTileMapping(in1, 0);

  // Tensor in2 spread out
  graph.setTileMapping(in2.index({0, 0}), 0);
  graph.setTileMapping(in2.index({0, 1}), 1);
  graph.setTileMapping(in2.index({1, 0}), 2);
  graph.setTileMapping(in2.index({1, 1}), 3);

  auto prog = Sequence();
  auto out1 = add(graph, in1, in2, prog);
  auto out2 = add(graph, in2, in1, prog);

  const auto &tile1 = graph.getTileMapping(out1);
  CHECK(tile1[0].size() > 0);
  CHECK(tile1[1].size() > 0);
  CHECK(tile1[2].size() > 0);
  CHECK(tile1[3].size() > 0);

  const auto &tile2 = graph.getTileMapping(out2);
  CHECK(tile2[0].size() > 0);
  CHECK(tile2[1].size() > 0);
  CHECK(tile2[2].size() > 0);
  CHECK(tile2[3].size() > 0);
}

void trinaryOutputMapChoiceTest() {
  auto device = createTestDevice(deviceType, 1, 4);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  Tensor in1, in2, in3, in4;
  in1 = graph.addVariable(FLOAT, {2, 2}, "t1");
  in2 = graph.addVariable(FLOAT, {2, 2}, "t2");
  in3 = graph.addVariable(BOOL, {2, 2}, "pred1");
  in4 = graph.addVariable(BOOL, {2, 2}, "pred2");

  // Tensor in1 all on tile 0
  graph.setTileMapping(in1, 0);

  // Tensor in2 spread out
  graph.setTileMapping(in2.index({0, 0}), 0);
  graph.setTileMapping(in2.index({0, 1}), 1);
  graph.setTileMapping(in2.index({1, 0}), 2);
  graph.setTileMapping(in2.index({1, 1}), 3);

  // Tensor pred1 all on tile 1
  graph.setTileMapping(in3, 1);

  // Tensor pred2 spread out
  graph.setTileMapping(in4.index({0, 0}), 0);
  graph.setTileMapping(in4.index({0, 1}), 1);
  graph.setTileMapping(in4.index({1, 0}), 2);
  graph.setTileMapping(in4.index({1, 1}), 3);

  auto prog = Sequence();
  auto out1 = select(graph, in1, in2, in3, prog);
  auto out2 = select(graph, in2, in1, in3, prog);
  auto out3 = select(graph, in1, in1, in4, prog);

  const auto &tile1 = graph.getTileMapping(out1);
  CHECK(tile1[0].size() > 0);
  CHECK(tile1[1].size() > 0);
  CHECK(tile1[2].size() > 0);
  CHECK(tile1[3].size() > 0);

  const auto &tile2 = graph.getTileMapping(out2);
  CHECK(tile2[0].size() > 0);
  CHECK(tile2[1].size() > 0);
  CHECK(tile2[2].size() > 0);
  CHECK(tile2[3].size() > 0);

  const auto &tile3 = graph.getTileMapping(out3);
  CHECK(tile3[0].size() > 0);
  CHECK(tile3[1].size() > 0);
  CHECK(tile3[2].size() > 0);
  CHECK(tile3[3].size() > 0);
}

void allTrueBadTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  Tensor in = graph.addVariable(FLOAT, {2, 2}, "t1");
  auto prog = Sequence();

  bool throws = false;
  try {
    allTrue(graph, in, prog, "all_true");
  } catch (const poplibs_error &) {
    throws = true;
  }
  CHECK(throws);
}

void allTrueTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  constexpr auto SIZE = 2;
  Tensor in = graph.addVariable(INT, {SIZE}, "t1");
  Tensor ones = graph.addConstant(INT, {SIZE}, 1);
  Tensor zeros = graph.addConstant(INT, {SIZE}, 0);
  graph.setTileMapping(ones, 0);
  graph.setTileMapping(zeros, 0);
  graph.setTileMapping(in, 0);

  auto bodyProg = Sequence();
  subInPlace(graph, in, ones, bodyProg);

  auto condProg = Sequence();
  Tensor neZero = neq(graph, in, zeros, condProg);
  auto predicate = allTrue(graph, neZero, condProg, "all_true");

  int init[SIZE] = {10, 8};
  int output[SIZE] = {0, 0};

  auto mainProg = Sequence();
  mainProg.add(RepeatWhileTrue(condProg, predicate, bodyProg));
  graph.createHostWrite("in", in);
  graph.createHostRead("out", in);

  Engine eng(graph, mainProg, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", init, &init[SIZE]);
    eng.run();
    eng.readTensor("out", output, &output[SIZE]);
  });

  CHECK(output[0] == 2);
  CHECK(output[1] == 0);
}

void isFiniteTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  float hIn[DIM_SIZE][DIM_SIZE];
  setUnaryOpInput(hIn);
  hIn[0][0] = INFINITY;
  hIn[1][0] = -INFINITY;
  hIn[0][1] = NAN;
  hIn[1][1] = NAN;
  hIn[0][2] = 0.0f;

  auto in = mapUnaryOpTensor(graph, FLOAT);
  auto prog = Sequence();
  auto out = isFinite(graph, in, prog);
  graph.createHostWrite("in", in);
  graph.createHostRead("out", out);

  bool hOut[DIM_SIZE][DIM_SIZE];

  // Disabling float point exception here cause it needs to check for INF and
  // that will automatically trigger an exception
  options.set("debug.floatPointOpException", "false");
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      bool expected = !(i <= 1 && j <= 1);
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

using namespace popops::expr;

void mapTestCast(bool inPlace, poplar::Type in1Type, poplar::Type in2Type) {
  auto device = createTestDevice(deviceType);
  auto target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto in1 = graph.addVariable(in1Type, {DIM_SIZE, DIM_SIZE}, "in1");
  auto in2 = graph.addVariable(in2Type, {DIM_SIZE, DIM_SIZE}, "in2");
  auto in3 = graph.addVariable(in1Type, {DIM_SIZE, DIM_SIZE}, "in3");
  mapTensorLinearly(graph, in1);
  mapTensorLinearly(graph, in2);
  mapTensorLinearly(graph, in3);
  Tensor out, out2;
  if (inPlace) {
    mapInPlace(graph, Add(_1, Cast(_2, in1Type)), {in1, in2}, prog);
    mapInPlace(graph, Add(Cast(_2, in1Type), _1), {in3, in2}, prog);
  } else {
    out = map(graph, Add(_1, Cast(_2, in1Type)), {in1, in2}, prog);
    // The type cast to determines the output type.
    out2 =
        map(graph, Add(Cast(_2, in1Type), Cast(_1, in1Type)), {in1, in2}, prog);
  }
  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in3", in3);
  graph.createHostWrite("in2", in2);
  if (inPlace) {
    graph.createHostRead("out", in1);
    graph.createHostRead("out2", in3);
  } else {
    graph.createHostRead("out", out);
    graph.createHostRead("out2", out2);
  }
  float hIn1[DIM_SIZE][DIM_SIZE];
  float hIn2[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];
  float hOut2[DIM_SIZE][DIM_SIZE];

  setBinaryOpInputsHalf(hIn1, hIn2);
  auto rawBufSize = target.getTypeSize(HALF) * DIM_SIZE * DIM_SIZE;
  std::vector<char> rawIn1(rawBufSize), rawIn2(rawBufSize), rawOut(rawBufSize),
      rawOut2(rawBufSize);
  if (in1Type == HALF) {
    poplar::copyFloatToDeviceHalf(target, &hIn1[0][0], rawIn1.data(),
                                  DIM_SIZE * DIM_SIZE);
  }
  if (in2Type == HALF) {
    poplar::copyFloatToDeviceHalf(target, &hIn2[0][0], rawIn2.data(),
                                  DIM_SIZE * DIM_SIZE);
  }
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    if (in1Type == HALF) {
      eng.writeTensor("in1", rawIn1.data(), rawIn1.data() + rawIn1.size());
      eng.writeTensor("in3", rawIn1.data(), rawIn1.data() + rawIn1.size());
    } else {
      eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
      eng.writeTensor("in3", hIn1, &hIn1[DIM_SIZE]);
    }
    if (in2Type == HALF) {
      eng.writeTensor("in2", rawIn2.data(), rawIn2.data() + rawIn2.size());
    } else {
      eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    }
    eng.run();
    if (in1Type == HALF) {
      eng.readTensor("out", rawOut.data(), rawOut.data() + rawOut.size());
      eng.readTensor("out2", rawOut2.data(), rawOut2.data() + rawOut2.size());
    } else {
      eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
      eng.readTensor("out2", hOut2, &hOut2[DIM_SIZE]);
    }
  });
  if (in1Type == HALF) {
    poplar::copyDeviceHalfToFloat(target, rawOut.data(), &hOut[0][0],
                                  DIM_SIZE * DIM_SIZE);
    poplar::copyDeviceHalfToFloat(target, rawOut2.data(), &hOut2[0][0],
                                  DIM_SIZE * DIM_SIZE);
  }
  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = hIn1[i][j] + hIn2[i][j];
      CHECK_CLOSE(hOut[i][j], expected);
      CHECK_CLOSE(hOut2[i][j], expected);
    }
  }
}

template <typename T, typename TestT>
void unaryMapExprTest(const expr::Expr &expr,
                      const std::function<TestT(T)> &testFn,
                      bool positiveInputs = false) {
  auto op = [&](Graph &graph, const Tensor &t, Sequence &prog,
                const poplar::DebugContext &,
                const poplar::OptionFlags &) -> Tensor {
    return map(graph, expr, {t}, prog);
  };
  return unaryOpTest(op, testFn, positiveInputs);
}

template <typename T, typename TestT>
void binaryMapExprTest(const expr::Expr &expr,
                       const std::function<TestT(T, T)> &testFn,
                       bool positiveInputs = false) {
  auto op = [&](Graph &graph, const Tensor &t0, const Tensor &t1,
                Sequence &prog, const poplar::DebugContext &,
                const poplar::OptionFlags &) -> Tensor {
    return map(graph, expr, {t0, t1}, prog);
  };
  return binaryOpTest(op, testFn, positiveInputs);
}

void mapTestCastIntToFloat() {
  auto device = createTestDevice(deviceType);
  auto target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto in1 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in1");
  auto in2 = graph.addVariable(INT, {DIM_SIZE, DIM_SIZE}, "in2");
  mapTensorLinearly(graph, in1);
  mapTensorLinearly(graph, in2);
  auto out = map(graph, Add(_1, Cast(_2, FLOAT)), {in1, in2}, prog);
  auto out2 = map(graph, Cast(_1, FLOAT), {in2}, prog);

  graph.createHostWrite("in1", in1);
  graph.createHostWrite("in2", in2);
  graph.createHostRead("out", out);
  graph.createHostRead("out2", out2);

  float hIn1[DIM_SIZE][DIM_SIZE];
  int hIn2[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];
  float hOut2[DIM_SIZE][DIM_SIZE];

  setUnaryOpInput(hIn1);
  setUnaryOpInput(hIn2);
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
    eng.readTensor("out2", hOut2, &hOut2[DIM_SIZE]);
  });
  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = hIn1[i][j] + hIn2[i][j];
      CHECK_CLOSE(hOut[i][j], expected);

      CHECK_CLOSE(hOut2[i][j], hIn2[i][j]);
    }
  }
}
void mapTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto in = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in");
  mapTensorLinearly(graph, in);
  auto out = map(graph, Add(Abs(_1), Const(3)), {in}, prog);

  graph.createHostWrite("in", in);
  graph.createHostRead("out", out);

  float hIn[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];

  setUnaryOpInput(hIn);

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = fabs(hIn[i][j]) + 3;
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

void mapTestMultiTensor() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto a = graph.addVariable(FLOAT, {DIM_SIZE}, "a");
  mapTensorLinearly(graph, a);
  auto b = graph.addVariable(FLOAT, {DIM_SIZE}, "b");
  mapTensorLinearly(graph, b);
  auto c = graph.addVariable(FLOAT, {DIM_SIZE}, "c");
  mapTensorLinearly(graph, c);
  auto out = map(graph, Select(Add(_3, _2), Add(_1, _2), Gte(_1, _2)),
                 {a, b, c}, prog);

  graph.createHostWrite("a", a);
  graph.createHostWrite("b", b);
  graph.createHostWrite("c", c);
  graph.createHostRead("out", out);

  float aIn[DIM_SIZE], bIn[DIM_SIZE], cIn[DIM_SIZE];
  float hOut[DIM_SIZE];

  std::mt19937 randomEngine;
  br::uniform_real_distribution<> dist(-10.0, +10.0);
  for (unsigned i = 0; i < DIM_SIZE; ++i) {
    aIn[i] = dist(randomEngine);
    bIn[i] = dist(randomEngine);
    cIn[i] = dist(randomEngine);
  }

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("a", aIn, &aIn[DIM_SIZE]);
    eng.writeTensor("b", bIn, &bIn[DIM_SIZE]);
    eng.writeTensor("c", cIn, &cIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    auto expected = (aIn[i] >= bIn[i]) ? cIn[i] + bIn[i] : aIn[i] + bIn[i];
    CHECK_CLOSE(hOut[i], expected);
  }
}

void mapInPlaceTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto t = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in");
  mapTensorLinearly(graph, t);
  mapInPlace(graph, Add(Abs(_1), _1), {t}, prog);

  graph.createHostWrite("in", t);
  graph.createHostRead("out", t);

  float hIn[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];

  setUnaryOpInput(hIn);

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = fabs(hIn[i][j]) + hIn[i][j];
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

void mapInPlaceBroadcastTest() {
  auto device = createTestDevice(deviceType, 1, 4);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto t = graph.addVariable(FLOAT, {2, 2, 2, 2}, "in");

  // force there to be more than one contiguous region
  graph.setTileMapping(t, 0);
  graph.setTileMapping(t[0][1][1], 3);

  const float scale = 0.5f;
  mapInPlace(graph, Mul(_1, Const(scale)), {t}, prog);

  Engine eng(graph, prog, options);
}

void mapFusedOpsWithScalarsTest(unsigned dimSize) {
  auto device = createTestDevice(deviceType, 1, 1);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);
  const float scalarVal = 3.1;
  auto prog = Sequence();
  auto a = graph.addVariable(FLOAT, {dimSize}, "a");
  auto b = graph.addConstant(FLOAT, {1}, scalarVal, "b");
  graph.setTileMapping(a, 0);
  graph.setTileMapping(b, 0);

  // These expressions should use generateAndExecuteMappedOperations, and
  // therefore test that the mapped operation code generator can deal with
  // expressiosn containing some scalars and some non-scalars.
  auto c = map(graph, Add(_1, Square(_2)), {a, b}, prog);
  mapInPlace(graph, Add(_1, Square(_2)), {a, b}, prog);

  graph.createHostWrite("in a", a);
  graph.createHostRead("out a", a);
  graph.createHostRead("out c", c);
  float resultA[DIM_SIZE], resultC[DIM_SIZE], inA[DIM_SIZE];
  for (unsigned i = 0; i < dimSize; i++) {
    inA[i] = i + 1;
  }
  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in a", inA, &inA[dimSize]);
    eng.run();
    eng.readTensor("out a", resultA, &resultA[dimSize]);
    eng.readTensor("out c", resultC, &resultC[dimSize]);
  });
  for (unsigned i = 0; i < dimSize; i++) {
    const float result = inA[i] + scalarVal * scalarVal;
    CHECK_CLOSE(resultA[i], result);
    CHECK_CLOSE(resultC[i], result);
  }
}
void mapInferTypeTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto a = graph.addVariable(FLOAT, {DIM_SIZE}, "a");
  mapTensorLinearly(graph, a);
  auto b = graph.addVariable(FLOAT, {DIM_SIZE}, "b");
  mapTensorLinearly(graph, b);
  auto c = graph.addVariable(BOOL, {DIM_SIZE}, "c");
  mapTensorLinearly(graph, c);
  auto out = map(graph, And(Equal(_1, _2), _3), {a, b, c}, prog);

  graph.createHostWrite("a", a);
  graph.createHostWrite("b", b);
  graph.createHostWrite("c", c);
  graph.createHostRead("out", out);

  float aIn[DIM_SIZE], bIn[DIM_SIZE];
  bool cIn[DIM_SIZE];
  bool hOut[DIM_SIZE];

  std::mt19937 randomEngine;
  br::uniform_real_distribution<> dist(-10.0, +10.0);
  for (unsigned i = 0; i < DIM_SIZE; ++i) {
    aIn[i] = dist(randomEngine);
    bIn[i] = dist(randomEngine);
    cIn[i] = dist(randomEngine) > 0.0;
  }

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("a", aIn, &aIn[DIM_SIZE]);
    eng.writeTensor("b", bIn, &bIn[DIM_SIZE]);
    eng.writeTensor("c", cIn, &cIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    auto expected = (aIn[i] == bIn[i]) && cIn[i];
    CHECK_CLOSE(hOut[i], expected);
  }
}

void addInPlaceTest() {
  auto tdevice = createTestDevice(deviceType);
  Graph graph(tdevice.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto t1 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in1");
  auto t2 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in2");

  mapTensorLinearly(graph, t1);
  mapTensorLinearly(graph, t2);
  addInPlace(graph, t1, t2, prog);

  graph.createHostWrite("in1", t1);
  graph.createHostWrite("in2", t2);
  graph.createHostRead("out", t1);

  float hIn1[DIM_SIZE][DIM_SIZE];
  float hIn2[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];

  setUnaryOpInput(hIn1);
  setUnaryOpInput(hIn2);

  Engine eng(graph, prog, options);
  tdevice.bind([&](const Device &device) {
    eng.load(device);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);

    if (isIpuModel(deviceType)) {
      eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
    }
  });
  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = hIn1[i][j] + hIn2[i][j];
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

void multiplyFloatInPlaceConstScalarTest() {
  auto device = createTestDevice(deviceType);
  Graph graph(device.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto t = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in");
  mapTensorLinearly(graph, t);
  const float b = 1.2;
  mulInPlace(graph, t, b, prog);

  graph.createHostWrite("in", t);
  graph.createHostRead("out", t);

  float hIn[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];

  setUnaryOpInput(hIn);

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", hIn, &hIn[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }
  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = hIn[i][j] * b;
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

void addHalfConstScalarTest() {
  auto device = createTestDevice(deviceType);
  auto target = device.getTarget();
  Graph graph(target);
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto tIn = graph.addVariable(HALF, {DIM_SIZE, DIM_SIZE}, "in");
  mapTensorLinearly(graph, tIn);
  const float b = 1.0;
  auto tOut = add(graph, tIn, b, prog);

  graph.createHostWrite("in", tIn);
  graph.createHostRead("out", tOut);

  float hIn[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][DIM_SIZE];

  setUnaryOpInputHalf(hIn);

  auto rawBufSize = target.getTypeSize(HALF) * DIM_SIZE * DIM_SIZE;
  std::vector<char> rawIn(rawBufSize), rawOut(rawBufSize);
  poplar::copyFloatToDeviceHalf(target, &hIn[0][0], rawIn.data(),
                                DIM_SIZE * DIM_SIZE);

  Engine eng(graph, prog, options);
  device.bind([&](const Device &d) {
    eng.load(d);
    eng.writeTensor("in", rawIn.data(), rawIn.data() + rawIn.size());
    eng.run();
    eng.readTensor("out", rawOut.data(), rawOut.data() + rawOut.size());
  });

  if (isIpuModel(deviceType)) {
    eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
  }
  poplar::copyDeviceHalfToFloat(target, rawOut.data(), &hOut[0][0],
                                DIM_SIZE * DIM_SIZE);
  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < DIM_SIZE; ++j) {
      auto expected = hIn[i][j] + b;
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

void binaryConcatTest() {
  auto tdevice = createTestDevice(deviceType);
  Graph graph(tdevice.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto t1 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in1");
  auto t2 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in2");
  auto t3 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in3");
  auto t4 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in4");

  mapTensorLinearly(graph, t1);
  mapTensorLinearly(graph, t2);
  mapTensorLinearly(graph, t3);
  mapTensorLinearly(graph, t4);
  auto t5 = add(graph, concat(t1, t2, 1), concat(t3, t4, 1), prog);

  graph.createHostWrite("in1", t1);
  graph.createHostWrite("in2", t2);
  graph.createHostWrite("in3", t3);
  graph.createHostWrite("in4", t4);
  graph.createHostRead("out", t5);

  float hIn1[DIM_SIZE][DIM_SIZE];
  float hIn2[DIM_SIZE][DIM_SIZE];
  float hIn3[DIM_SIZE][DIM_SIZE];
  float hIn4[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][2 * DIM_SIZE];

  setUnaryOpInput(hIn1);
  setUnaryOpInput(hIn2);
  setUnaryOpInput(hIn3);
  setUnaryOpInput(hIn4);

  Engine eng(graph, prog, options);
  tdevice.bind([&](const Device &device) {
    eng.load(device);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.writeTensor("in3", hIn3, &hIn3[DIM_SIZE]);
    eng.writeTensor("in4", hIn4, &hIn4[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);

    if (isIpuModel(deviceType)) {
      eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
    }
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < 2 * DIM_SIZE; ++j) {
      auto expected = j < DIM_SIZE
                          ? hIn1[i][j] + hIn3[i][j]
                          : hIn2[i][j - DIM_SIZE] + hIn4[i][j - DIM_SIZE];
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

void unaryConcatTest() {
  auto tdevice = createTestDevice(deviceType);
  Graph graph(tdevice.getTarget());
  popops::addCodelets(graph);

  auto prog = Sequence();
  auto t1 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in1");
  auto t2 = graph.addVariable(FLOAT, {DIM_SIZE, DIM_SIZE}, "in2");

  mapTensorLinearly(graph, t1);
  mapTensorLinearly(graph, t2);
  auto t3 = neg(graph, concat(t1, t2, 1), prog);

  graph.createHostWrite("in1", t1);
  graph.createHostWrite("in2", t2);
  graph.createHostRead("out", t3);

  float hIn1[DIM_SIZE][DIM_SIZE];
  float hIn2[DIM_SIZE][DIM_SIZE];
  float hOut[DIM_SIZE][2 * DIM_SIZE];

  setUnaryOpInput(hIn1);
  setUnaryOpInput(hIn2);

  Engine eng(graph, prog, options);
  tdevice.bind([&](const Device &device) {
    eng.load(device);
    eng.writeTensor("in1", hIn1, &hIn1[DIM_SIZE]);
    eng.writeTensor("in2", hIn2, &hIn2[DIM_SIZE]);
    eng.run();
    eng.readTensor("out", hOut, &hOut[DIM_SIZE]);

    if (isIpuModel(deviceType)) {
      eng.printProfileSummary(std::cerr, {{"showExecutionSteps", "true"}});
    }
  });

  /* Check result */
  for (auto i = 0U; i < DIM_SIZE; ++i) {
    for (auto j = 0U; j < 2 * DIM_SIZE; ++j) {
      auto expected = j < DIM_SIZE ? -hIn1[i][j] : -hIn2[i][j - DIM_SIZE];
      CHECK_CLOSE(hOut[i][j], expected);
    }
  }
}

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

  std::string test;

  po::options_description desc("Options");
  // clang-format off
  desc.add_options() ("help", "Print help")
    ("device-type",
     po::value<DeviceType>(&deviceType)->required(),
     "Device Type")
    ("test",
     po::value<std::string>(&test)->required(),
     "The test to run");
  // 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";
      return 1;
    }
    po::notify(vm);
  } catch (std::exception &e) {
    std::cerr << "error: " << e.what() << "\n";
    return 1;
  }

  if (test == "AbsFloat") {
    unaryOpTest<float, double>(popops::abs, [](float x) -> double {
      double res = fabs(static_cast<double>(x));
      return res;
    });
  } else if (test == "AbsInt") {
    unaryOpTest<int, int>(popops::abs,
                          [](int x) -> int { return std::abs(x); });
  } else if (test == "AddFloat") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::add),
                                [](float x, float y) -> double {
                                  double res = x + y;
                                  return res;
                                });
  } else if (test == "Atan2Float") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::atan2),
                                [](float x, float y) -> double {
                                  double res = std::atan2(x, y);
                                  return res;
                                });
  } else if (test == "AddInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::add),
                           [](int x, int y) -> int { return x + y; });
  } else if (test == "BitwiseAndInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::bitwiseAnd),
                           [](int x, int y) -> int { return x & y; });
  } else if (test == "BitwiseOrInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::bitwiseOr),
                           [](int x, int y) -> int { return x | y; });
  } else if (test == "BitwiseXorInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::bitwiseXor),
                           [](int x, int y) -> int { return x ^ y; });
  } else if (test == "BitwiseXnorInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::bitwiseXnor),
                           [](int x, int y) -> int { return ~(x ^ y); });
  } else if (test == "BitwiseNotInt") {
    unaryOpTest<int, int>(popops::bitwiseNot, [](int x) -> int { return ~x; });
  } else if (test == "Ceil") {
    unaryOpTest<float, double>(popops::ceil, [](float x) -> double {
      double res = std::ceil(static_cast<double>(x));
      return res;
    });

  } else if (test == "Cos") {
    unaryOpTest<float, double>(popops::cos, [](float x) -> double {
      double res = std::cos(static_cast<double>(x));
      return res;
    });
  } else if (test == "CountLeadingZeros") {
    unaryOpTest<int, int>(popops::countLeadingZeros, [](int x) -> int {
      return x ? __builtin_clz(x) : 0;
    });
  } else if (test == "DivideFloat") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::div),
                                [](float x, float y) -> double {
                                  double res = x / y;
                                  return res;
                                });
  } else if (test == "DivideHalf") {
    binaryOpTestHalf(static_cast<BinaryOpFnPtr_t>(popops::div),
                     [](float x, float y) -> double {
                       double res = x / y;
                       return res;
                     });
  } else if (test == "DivideInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::div),
                           [](int x, int y) -> int {
                             int res = x / y;
                             return res;
                           });
  } else if (test == "EqualFloat") {
    binaryOpTest<float, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::eq),
        [](float x, float y) -> bool { return x == y; });
  } else if (test == "GreaterThanBool") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::gt),
        [](bool x, bool y) -> bool { return x > y; });
  } else if (test == "GreaterThanEqualBool") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::gteq),
        [](bool x, bool y) -> bool { return x >= y; });
  } else if (test == "LessThanBool") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::lt),
        [](bool x, bool y) -> bool { return x < y; });
  } else if (test == "LessThanEqualBool") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::lteq),
        [](bool x, bool y) -> bool { return x <= y; });
  } else if (test == "Exponent") {
    unaryOpTest<float, float>(popops::exp, [](float x) -> float {
      double res = std::exp(static_cast<double>(x));
      return res;
    });
  } else if (test == "ExponentMinus1") {
    unaryOpTest<float, float>(popops::expm1, [](float x) -> float {
      double res = std::expm1(static_cast<double>(x));
      return res;
    });
  } else if (test == "Floor") {
    unaryOpTest<float, double>(popops::floor, [](float x) -> double {
      double res = std::floor(static_cast<double>(x));
      return res;
    });
  } else if (test == "GreaterThanFloat") {
    binaryOpTest<float, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::gt),
        [](float x, float y) -> bool { return x > y; });
  } else if (test == "GreaterThanInt") {
    binaryOpTest<int, bool, bool>(static_cast<BinaryOpFnPtr_t>(popops::gt),
                                  [](int x, int y) -> bool { return x > y; });
  } else if (test == "GreaterThanEqualFloat") {
    binaryOpTest<float, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::gteq),
        [](float x, float y) -> bool { return x >= y; });
  } else if (test == "LessThanFloat") {
    binaryOpTest<float, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::lt),
        [](float x, float y) -> bool { return x < y; });
  } else if (test == "LessThanEqualFloat") {
    binaryOpTest<float, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::lteq),
        [](float x, float y) -> bool { return x <= y; });
  } else if (test == "Logarithm") {
    unaryOpTest<float, double>(
        popops::log,
        [](float x) -> double {
          double res = std::log(static_cast<double>(x));
          return res;
        },
        true /* positive inputs */);
  } else if (test == "Logarithm1Plus") {
    unaryOpTest<float, double>(
        popops::log1p,
        [](float x) -> double {
          double res = std::log1p(static_cast<double>(x));
          return res;
        },
        true /* positive inputs */);
  } else if (test == "LogicalAnd") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::logicalAnd),
        [](bool x, bool y) -> bool { return x && y; });
  } else if (test == "LogicalNot") {
    unaryOpTest<bool, bool>(popops::logicalNot,
                            [](bool x) -> bool { return !x; });
  } else if (test == "LogicalOr") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::logicalOr),
        [](bool x, bool y) -> bool { return x || y; });
  } else if (test == "MaxFloat") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::max),
                                [](float x, float y) -> double {
                                  double res = std::max(x, y);
                                  return res;
                                });
  } else if (test == "MaxInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::max),
                           [](int x, int y) -> int {
                             auto res = std::max(x, y);
                             return res;
                           });
  } else if (test == "MinFloat") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::min),
                                [](float x, float y) -> double {
                                  double res = std::min(x, y);
                                  return res;
                                });
  } else if (test == "MinInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::min),
                           [](int x, int y) -> int {
                             auto res = std::min(x, y);
                             return res;
                           });
  } else if (test == "Multiply") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::mul),
                                [](float x, float y) -> double {
                                  double res = x * y;
                                  return res;
                                });
  } else if (test == "NotEqualFloat") {
    binaryOpTest<float, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::neq),
        [](float x, float y) -> bool { return x != y; });
  } else if (test == "NotEqualBool") {
    binaryOpTest<bool, bool, bool>(
        static_cast<BinaryOpFnPtr_t>(popops::neq),
        [](bool x, bool y) -> bool { return x != y; });
  } else if (test == "NegateFloat") {
    unaryOpTest<float, double>(popops::neg,
                               [](float x) -> double { return -x; });
  } else if (test == "NegateInt") {
    unaryOpTest<int, int>(popops::neg, [](int x) -> int { return -x; });
  } else if (test == "Popcount") {
    unaryOpTest<int, int>(popops::popcount,
                          [](int x) -> int { return __builtin_popcount(x); });
  } else if (test == "Power") {
    powTest();
  } else if (test == "RemainderFloat") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::rem),
                                [](float x, float y) -> double {
                                  double res =
                                      std::fmod(static_cast<double>(x),
                                                static_cast<double>(y));
                                  return res;
                                });
  } else if (test == "RemainderInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::rem),
                           [](int x, int y) -> double { return x % y; });
  } else if (test == "ShiftLeftInt") {
    binaryOpTest<int, int, int>(
        static_cast<BinaryOpFnPtr_t>(popops::shiftLeft),
        [](int x, int y) -> int { return x << y; }, true);
  } else if (test == "ShiftRightInt") {
    binaryOpTest<int, int, int>(
        static_cast<BinaryOpFnPtr_t>(popops::shiftRight),
        [](int x, int y) -> int { return (unsigned)x >> y; }, true);
  } else if (test == "ShiftRightSignExtendInt") {
    binaryOpTest<int, int, int>(
        static_cast<BinaryOpFnPtr_t>(popops::shiftRightSignExtend),
        [](int x, int y) -> int { return x >> y; }, true);
  } else if (test == "SignumFloat") {
    unaryOpTest<float, double>(popops::signum, [](float x) -> double {
      double res = (0 < x) - (x < 0);
      return res;
    });
  } else if (test == "SignumInt") {
    unaryOpTest<int, int>(popops::signum, [](int x) -> int {
      int res = (0 < x) - (x < 0);
      return res;
    });
  } else if (test == "Sin") {
    unaryOpTest<float, double>(popops::sin, [](float x) -> double {
      double res = std::sin(static_cast<double>(x));
      return res;
    });
  } else if (test == "Asin") {
    unaryOpTest<float, double>(
        popops::asin,
        [](float x) -> double {
          double res = std::asin(static_cast<double>(x));
          return res;
        },
        false, true);
  } else if (test == "Tan") {
    unaryOpTest<float, double>(popops::tan, [](float x) -> double {
      double res = std::tan(static_cast<double>(x));
      return res;
    });
  } else if (test == "Tanh") {
    unaryOpTest<float, double>(popops::tanh, [](float x) -> double {
      double res = std::tanh(static_cast<double>(x));
      return res;
    });
  } else if (test == "Square") {
    unaryOpTest<float, double>(popops::square, [](float x) -> double {
      double xd = static_cast<double>(x);
      double res = xd * xd;
      return res;
    });
  } else if (test == "SquareRoot") {
    unaryOpTest<float, double>(
        popops::sqrt,
        [](float x) -> double {
          double xd = static_cast<double>(x);
          double res = std::sqrt(xd);
          return res;
        },
        true /* positive inputs */);
  } else if (test == "Sigmoid") {
    unaryOpTest<float, double>(popops::sigmoid, [](float x) -> double {
      double xd = static_cast<double>(x);
      double res = (1.0 / (1.0 + std::exp(-xd)));
      return res;
    });
  } else if (test == "Rsqrt") {
    unaryOpTest<float, double>(popops::rsqrt, [](float x) -> double {
      double xd = static_cast<double>(x);
      double res = 1.0 / std::sqrt(xd);
      return res;
    });
  } else if (test == "SubtractFloat") {
    binaryOpTest<float, double>(static_cast<BinaryOpFnPtr_t>(popops::sub),
                                [](float x, float y) -> double {
                                  return static_cast<double>(x) -
                                         static_cast<double>(y);
                                });
  } else if (test == "SubtractHalf") {
    binaryOpTestHalf(static_cast<BinaryOpFnPtr_t>(popops::sub),
                     [](float x, float y) -> float { return x - y; });
  } else if (test == "SubtractInt") {
    binaryOpTest<int, int>(static_cast<BinaryOpFnPtr_t>(popops::sub),
                           [](int x, int y) -> int { return (x - y); });
  } else if (test == "RoundFloat") {
    unaryOpTest<float, double>(popops::round,
                               [](float x) -> double { return std::round(x); });
  } else if (test == "SelectFloat") {
    selectTest<float>();
  } else if (test == "SelectFloatLHSConst") {
    selectTestFloatLHSConst();
  } else if (test == "SelectFloatRHSConst") {
    selectTestFloatRHSConst();
  } else if (test == "SelectFloatLHSAndRHSConst") {
    selectTestFloatLHSAndRHSConst();
  } else if (test == "SelectHalfLHSAndRHSConst") {
    selectTestHalfLHSAndRHSConst();
  } else if (test == "SelectInt") {
    selectTest<int>();
  } else if (test == "BroadcastSelectorSelectInt") {
    broadcastSelectorSelectTest<int>(false);
  } else if (test == "BroadcastSelectorSelectFloat") {
    broadcastSelectorSelectTest<float>(false);
  } else if (test == "BroadcastSelectorSelectInPlaceInt") {
    broadcastSelectorSelectTest<int>(true);
  } else if (test == "BroadcastSelectorSelectInPlaceFloat") {
    broadcastSelectorSelectTest<float>(true);
  } else if (test == "ClampFloat") {
    clampTest<float>(false);
  } else if (test == "ClampFloatMinConst") {
    clampTestFloatMinConst();
  } else if (test == "ClampFloatMaxConst") {
    clampTestFloatMaxConst();
  } else if (test == "ClampInt") {
    clampTest<int>(false);
  } else if (test == "ClampInPlaceFloat") {
    clampTest<float>(true);
  } else if (test == "BroadcastClampInt") {
    broadcastClampTest<int>(false);
  } else if (test == "BroadcastClampInPlaceInt") {
    broadcastClampTest<int>(true);
  } else if (test == "BroadcastClampFloat") {
    broadcastClampTest<float>(false);
  } else if (test == "BroadcastClampInPlaceFloat") {
    broadcastClampTest<float>(true);
  } else if (test == "BroadcastClampSingleElementSrcFloat") {
    broadcastClampTest<float>(false, 1);
  } else if (test == "BinaryOutputMapChoice") {
    binaryOutputMapChoiceTest();
  } else if (test == "TrinaryOutputMapChoice") {
    trinaryOutputMapChoiceTest();
  } else if (test == "AllTrueBad") {
    allTrueBadTest();
  } else if (test == "AllTrue") {
    allTrueTest();
  } else if (test == "IsFinite") {
    isFiniteTest();
  } else if (test == "Map") {
    mapTest();
  } else if (test == "MapCast") {
    mapTestCast(false, FLOAT, HALF);
  } else if (test == "MapCastIntToFloat") {
    mapTestCastIntToFloat();
  } else if (test == "MapCastInPlace") {
    mapTestCast(true, HALF, FLOAT);
  } else if (test == "MapMultiTensor") {
    mapTestMultiTensor();
  } else if (test == "MapInPlace") {
    mapInPlaceTest();
  } else if (test == "MapInPlaceBroadcast") {
    mapInPlaceBroadcastTest();
  } else if (test == "MapAllScalar") {
    mapFusedOpsWithScalarsTest(1);
  } else if (test == "MapSomeScalar") {
    mapFusedOpsWithScalarsTest(DIM_SIZE);

  } else if (test == "MapInferType") {
    mapInferTypeTest();
  } else if (test == "MapInferTypeNot") {
    unaryMapExprTest<bool, bool>(Equal(Const(0), Not(_1)),
                                 [](bool x) -> bool { return false == !x; });
  } else if (test == "MapInferTypeEqual") {
    binaryMapExprTest<bool, bool>(
        Equal(Const(0), Equal(_1, _2)),
        [](bool x, bool y) -> bool { return false == (x == y); });
  } else if (test == "MapInferTypeCast") {
    unaryMapExprTest<float, float>(Cast(Add(_1, Const(1)), FLOAT),
                                   [](float x) -> float { return x + 1; });
  } else if (test == "AddInPlace") {
    addInPlaceTest();
  } else if (test == "BinaryConcat") {
    binaryConcatTest();
  } else if (test == "UnaryConcat") {
    unaryConcatTest();
  } else if (test == "MultiplyFloatInPlaceConstScalarTest") {
    multiplyFloatInPlaceConstScalarTest();
  } else if (test == "AddHalfConstScalarTest") {
    addHalfConstScalarTest();
  } else {
    throw std::runtime_error("Unknown test '" + test + "'");
  }

  return 0;
}


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2017 Cheng Li

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.  See the license for more details.
*/

#include <boost/make_shared.hpp>
#include <ql/cashflows/iborcoupon.hpp>
#include <ql/indexes/ibor/shibor.hpp>
#include <ql/pricingengines/swap/discountingswapengine.hpp>
#include <ql/time/calendars/china.hpp>
#include <ql/utilities/null_deleter.hpp>
#include <qlext/instruments/shiborswap.hpp>
#include <qlext/instruments/subperiodsswap.hpp>
#include <qlext/termstructures/yield/ratehelpers.hpp>

namespace QuantLib {

    ShiborSwapRateHelper::ShiborSwapRateHelper(const Handle<Quote>& rate, const Period& swapTenor, Frequency fixedFreq,
                                               const boost::shared_ptr<Shibor>& shiborIndex, const Period& fwdStart,
                                               const Handle<YieldTermStructure>& discountingCurve)
        : RelativeDateRateHelper(rate), tenor_(swapTenor), fixedFrequency_(fixedFreq), fwdStart_(fwdStart),
          discountHandle_(discountingCurve) {

        settlementDays_ = shiborIndex->fixingDays();
        shiborIndex_ = boost::dynamic_pointer_cast<Shibor>(shiborIndex->clone(termStructureHandle_));
        shiborIndex_->unregisterWith(termStructureHandle_);

        registerWith(shiborIndex_);
        registerWith(discountHandle_);

        initializeDates();
    }

    void ShiborSwapRateHelper::initializeDates() {

        Date refDate = Settings::instance().evaluationDate();
        China floatCalendar(China::IB);
        refDate = floatCalendar.adjust(refDate);

        Date spotDate = floatCalendar.advance(refDate, settlementDays_ * Days);
        Date startDate = spotDate + fwdStart_;

        if (fwdStart_.length() < 0)
            startDate = floatCalendar.adjust(startDate, Preceding);
        else
            startDate = floatCalendar.adjust(startDate, Following);

        swap_ = boost::make_shared<ShiborSwap>(VanillaSwap::Payer, 1., startDate, tenor_, Period(fixedFrequency_), 0.,
                                               shiborIndex_);

        bool includeSettlementDateFlows = false;
        boost::shared_ptr<PricingEngine> engine(
            new DiscountingSwapEngine(discountRelinkableHandle_, includeSettlementDateFlows));

        swap_->setPricingEngine(engine);

        earliestDate_ = swap_->startDate();
        maturityDate_ = swap_->maturityDate();

        boost::shared_ptr<IborCoupon> lastCoupon = boost::dynamic_pointer_cast<IborCoupon>(swap_->floatingLeg().back());
        latestRelevantDate_ = std::max(maturityDate_, lastCoupon->fixingEndDate());
        latestDate_ = maturityDate_;
    }

    void ShiborSwapRateHelper::accept(AcyclicVisitor& v) {
        Visitor<ShiborSwapRateHelper>* v1 = dynamic_cast<Visitor<ShiborSwapRateHelper>*>(&v);
        if (v1 != 0)
            v1->visit(*this);
        else
            RateHelper::accept(v);
    }

    Real ShiborSwapRateHelper::impliedQuote() const {
        QL_REQUIRE(termStructure_ != 0, "term structure not set");
        // we didn't register as observers - force calculation
        swap_->recalculate();
        // weak implementation... to be improved
        static const Spread basisPoint = 1.0e-4;
        Real floatingLegNPV = swap_->floatingLegNPV();
        Real totNPV = -floatingLegNPV;
        Real result = totNPV / (swap_->fixedLegBPS() / basisPoint);
        return result;
    }

    void ShiborSwapRateHelper::setTermStructure(YieldTermStructure* t) {
        // do not set the relinkable handle as an observer -
        // force recalculation when needed---the index is not lazy
        bool observer = false;

        boost::shared_ptr<YieldTermStructure> temp(t, null_deleter());
        termStructureHandle_.linkTo(temp, observer);

        if (discountHandle_.empty())
            discountRelinkableHandle_.linkTo(temp, observer);
        else
            discountRelinkableHandle_.linkTo(*discountHandle_, observer);

        RelativeDateRateHelper::setTermStructure(t);
    }

    SubPeriodsSwapRateHelper::SubPeriodsSwapRateHelper(
        const Handle<Quote>& rate, const Period& swapTenor, Frequency fixedFreq, const Calendar& fixedCalendar,
        const DayCounter& fixedDayCount, BusinessDayConvention fixedConvention, const Period& floatPayTenor,
        const boost::shared_ptr<IborIndex>& iborIndex, const DayCounter& floatingDayCount, DateGeneration::Rule rule,
        Ext::SubPeriodsCoupon::Type type, const Period& fwdStart, const Handle<YieldTermStructure>& discountingCurve)
        : RelativeDateRateHelper(rate), tenor_(swapTenor), fixedFrequency_(fixedFreq), fixedCalendar_(fixedCalendar),
          fixedDayCount_(fixedDayCount), fixedConvention_(fixedConvention), floatPayTenor_(floatPayTenor),
          floatingDayCount_(floatingDayCount), rule_(rule), type_(type), fwdStart_(fwdStart),
          discountHandle_(discountingCurve) {

        settlementDays_ = iborIndex->fixingDays();
        iborIndex_ = boost::dynamic_pointer_cast<IborIndex>(iborIndex->clone(termStructureHandle_));
        iborIndex_->unregisterWith(termStructureHandle_);

        registerWith(iborIndex_);
        registerWith(discountHandle_);

        initializeDates();
    }

    void SubPeriodsSwapRateHelper::initializeDates() {

        Date refDate = Settings::instance().evaluationDate();
        Calendar floatCalendar = iborIndex_->fixingCalendar();
        refDate = floatCalendar.adjust(refDate);

        Date spotDate = floatCalendar.advance(refDate, settlementDays_ * Days);
        Date startDate = spotDate + fwdStart_;

        if (fwdStart_.length() < 0)
            startDate = floatCalendar.adjust(startDate, Preceding);
        else
            startDate = floatCalendar.adjust(startDate, Following);

        swap_ = boost::shared_ptr<SubPeriodsSwap>(new SubPeriodsSwap(startDate, 1., tenor_, true, Period(fixedFrequency_), 0.,
                                                   fixedCalendar_, fixedDayCount_, fixedConvention_, floatPayTenor_,
                                                   iborIndex_, floatingDayCount_, rule_, type_));

        bool includeSettlementDateFlows = false;
        boost::shared_ptr<PricingEngine> engine(
            new DiscountingSwapEngine(discountRelinkableHandle_, includeSettlementDateFlows));

        swap_->setPricingEngine(engine);

        earliestDate_ = swap_->startDate();
        maturityDate_ = swap_->maturityDate();
        Leg floatLeg = swap_->floatLeg();

        boost::shared_ptr<Ext::SubPeriodsCoupon> lastCoupon =
            boost::dynamic_pointer_cast<Ext::SubPeriodsCoupon>(floatLeg[floatLeg.size() - 1]);
        latestRelevantDate_ = std::max(latestDate_, lastCoupon->latestRelevantDate());
        latestDate_ = maturityDate_;
    }

    void SubPeriodsSwapRateHelper::accept(AcyclicVisitor& v) {
        Visitor<SubPeriodsSwapRateHelper>* v1 = dynamic_cast<Visitor<SubPeriodsSwapRateHelper>*>(&v);
        if (v1 != 0)
            v1->visit(*this);
        else
            RateHelper::accept(v);
    }

    Real SubPeriodsSwapRateHelper::impliedQuote() const {
        QL_REQUIRE(termStructure_ != 0, "term structure not set");
        // we didn't register as observers - force calculation
        swap_->recalculate();
        // weak implementation... to be improved
        static const Spread basisPoint = 1.0e-4;
        Real floatingLegNPV = swap_->floatLegNPV();
        Real totNPV = -floatingLegNPV;
        Real result = totNPV / (swap_->fixedLegBPS() / basisPoint);
        return result;
    }

    void SubPeriodsSwapRateHelper::setTermStructure(YieldTermStructure* t) {
        // do not set the relinkable handle as an observer -
        // force recalculation when needed---the index is not lazy
        bool observer = false;

        boost::shared_ptr<YieldTermStructure> temp(t, null_deleter());
        termStructureHandle_.linkTo(temp, observer);

        if (discountHandle_.empty())
            discountRelinkableHandle_.linkTo(temp, observer);
        else
            discountRelinkableHandle_.linkTo(*discountHandle_, observer);

        RelativeDateRateHelper::setTermStructure(t);
    }

} // namespace QuantLib


#ifndef __solverEhqp_primal_h__
#define __solverEhqp_primal_h__

#include <Eigen/Dense>
#include <vector>
#include <iostream>
#include "weighted_hqp/HCod.hpp"

namespace hcod{
    class Ehqp_primal{
        public:
            Ehqp_primal(const std::vector<h_structure> &h, const Eigen::MatrixXd & Y);
            Ehqp_primal(const std::vector<h_structure> &h);
            ~Ehqp_primal(){};

        private:

        public:
            void setProblem(const std::vector<h_structure> &h, const Eigen::MatrixXd & Y){
                h_ = h;
                Y_ = Y;

                nh_ = Y.cols();
                p_ = h.size();
                int num=0;

                for (int i =0; i<p_; i++)
                    num += h_[i].m;

                y_.setZero(num);
            }
            void compute();
            Eigen::VectorXd getx(){
                return x_;
            }
            Eigen::VectorXd gety(){
                return y_;
            }
            void setWProblem(const std::vector<h_structure> &h){
                h_ = h;

                p_ = h.size();
                int num=0;

                for (int i =0; i<p_; i++)
                    num += h_[i].m;

                y_.setZero(num);
            }

        private:
           std::vector<h_structure> h_;
           Eigen::MatrixXd Y_, L_, M1_, W1_;
           int p_, nh_;
           Eigen::VectorXd y_, b_, e_, x_;
           bool _isweighted;

    };
}

#endif


// Copyright (c) 2017 Pieter Wuille
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <blech32.h>
#include <random.h>
#include <util/strencodings.h>
#include <test/util/setup_common.h>

#include <boost/test/unit_test.hpp>

BOOST_FIXTURE_TEST_SUITE(blech32_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(blech32_polymod_sanity)
{
    std::vector<unsigned char> data(40);
    // GetRandBytes only allows 32 bytes at a time
    GetRandBytes(data.data(), 32);
    GetRandBytes(data.data() + 32, data.size() - 32);

    std::vector<unsigned char> base32;
    ConvertBits<8, 5, true>([&](unsigned char c) { base32.push_back(c); }, data.begin(), data.end());
    uint64_t plm1 = blech32::PolyMod(base32);

    // Now add 1023 zeros.
    for (auto i = 0; i < 1023; i++) {
        base32.push_back(0);
    }
    uint64_t plm2 = blech32::PolyMod(base32);

    BOOST_CHECK_EQUAL(plm1, plm2);
}

BOOST_AUTO_TEST_CASE(blech32_checksum)
{
    std::vector<unsigned char> vector{7,2,3,4,5,6,7,8,9,234,123,213,16};
    std::vector<unsigned char> b32;
    ConvertBits<8, 5, true>([&](unsigned char c) { b32.push_back(c); }, vector.begin(), vector.end());
    std::vector<unsigned char> cs = blech32::CreateChecksum("lq", b32);

    std::vector<unsigned char> expected_cs{22,13,13,5,4,4,23,7,28,21,30,12};
    for (size_t i = 0; i < expected_cs.size(); i++) {
        BOOST_CHECK_EQUAL(expected_cs[i], cs[i]);
    }
}

BOOST_AUTO_TEST_SUITE_END()


#ifndef MMTBX_GEOMETRY_SPHERE_SURFACE_SAMPLING_H
#define MMTBX_GEOMETRY_SPHERE_SURFACE_SAMPLING_H

#include <boost/math/constants/constants.hpp>

#include <vector>
#include <cmath>

namespace mmtbx
{

namespace geometry
{

namespace sphere_surface_sampling
{

template< typename Vector >
class GoldenSpiral
{
public:
  typedef Vector vector_type;
  typedef typename Vector::value_type value_type;
  typedef std::vector< vector_type > storage_type;
  typedef typename storage_type::const_iterator const_storage_iterator;

private:
  size_t count_;
  value_type unit_area_;
  storage_type points_;

  static const value_type pitch;

public:
  GoldenSpiral(size_t count);
  ~GoldenSpiral();

  inline const value_type& unit_area() const;
  inline const size_t& count() const;
  inline const storage_type& points() const;
};

#include "sphere_surface_sampling.hxx"

} // namespace sphere_surface_sampling
} // namespace geometry
} // namespace mmtbx

#endif // MMTBX_GEOMETRY_SPHERE_SURFACE_SAMPLING_H


#include <boost/numeric/ublas/matrix_sparse.hpp>


// Copyright 2015-2019 Hans Dembinski
//
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt
// or copy at http://www.boost.org/LICENSE_1_0.txt)

#include <benchmark/benchmark.h>
#include <boost/histogram/axis/regular.hpp>
#include <boost/histogram/storage_adaptor.hpp>
#include <boost/histogram/unlimited_storage.hpp>
#include <random>
#include "../test/utility_histogram.hpp"

using namespace boost::histogram;
using reg = axis::regular<>;
using uniform = std::uniform_real_distribution<>;
using normal = std::normal_distribution<>;

template <class Distribution>
Distribution init();

template <>
uniform init<uniform>() {
  return uniform{0.0, 1.0};
}

template <>
normal init<normal>() {
  return normal{0.5, 0.3};
}

template <class Tag, class Storage, class Distribution>
static void fill_1d(benchmark::State& state) {
  auto h = make_s(Tag(), Storage(), reg(100, 0, 1));
  std::default_random_engine gen(1);
  Distribution dis = init<Distribution>();
  for (auto _ : state) h(dis(gen));
}

template <class Tag, class Storage, class Distribution>
static void fill_2d(benchmark::State& state) {
  auto h = make_s(Tag(), Storage(), reg(100, 0, 1), reg(100, 0, 1));
  std::default_random_engine gen(1);
  Distribution dis = init<Distribution>();
  for (auto _ : state) { h(dis(gen), dis(gen)); }
}

template <class Tag, class Storage, class Distribution>
static void fill_3d(benchmark::State& state) {
  auto h = make_s(Tag(), Storage(), reg(100, 0, 1), reg(100, 0, 1), reg(100, 0, 1));
  std::default_random_engine gen(1);
  Distribution dis = init<Distribution>();
  for (auto _ : state) { h(dis(gen), dis(gen), dis(gen)); }
}

template <class Tag, class Storage, class Distribution>
static void fill_6d(benchmark::State& state) {
  auto h = make_s(Tag(), Storage(), reg(10, 0, 1), reg(10, 0, 1), reg(10, 0, 1),
                  reg(10, 0, 1), reg(10, 0, 1), reg(10, 0, 1));
  std::default_random_engine gen(1);
  Distribution dis = init<Distribution>();
  for (auto _ : state) { h(dis(gen), dis(gen), dis(gen), dis(gen), dis(gen), dis(gen)); }
}

using SStore = std::vector<int>;
using DStore = unlimited_storage<>;

BENCHMARK_TEMPLATE(fill_1d, static_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_1d, static_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_1d, dynamic_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_1d, dynamic_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_2d, static_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_2d, static_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_2d, dynamic_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_2d, dynamic_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_3d, static_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_3d, static_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_3d, dynamic_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_3d, dynamic_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_6d, static_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_6d, static_tag, DStore, uniform);
BENCHMARK_TEMPLATE(fill_6d, dynamic_tag, SStore, uniform);
BENCHMARK_TEMPLATE(fill_6d, dynamic_tag, DStore, uniform);

BENCHMARK_TEMPLATE(fill_1d, static_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_1d, static_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_1d, dynamic_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_1d, dynamic_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_2d, static_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_2d, static_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_2d, dynamic_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_2d, dynamic_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_3d, static_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_3d, static_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_3d, dynamic_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_3d, dynamic_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_6d, static_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_6d, static_tag, DStore, normal);
BENCHMARK_TEMPLATE(fill_6d, dynamic_tag, SStore, normal);
BENCHMARK_TEMPLATE(fill_6d, dynamic_tag, DStore, normal);


/**
* @file lldateutil.cpp
*
* $LicenseInfo:firstyear=2009&license=viewerlgpl$
* Second Life Viewer Source Code
* Copyright (C) 2010, Linden Research, Inc.
*
* This library is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation;
* version 2.1 of the License only.
*
* This library is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with this library; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*
* Linden Research, Inc., 945 Battery Street, San Francisco, CA  94111  USA
* $/LicenseInfo$
*/

#include "llviewerprecompiledheaders.h"

#include "lldateutil.h"

#include <boost/date_time/gregorian/gregorian.hpp>
#include <boost/date_time/posix_time/ptime.hpp>

// Linden libraries
#include "lltrans.h"
#include "llui.h"

using namespace boost::gregorian;
using namespace boost::posix_time;

static S32 DAYS_PER_MONTH_NOLEAP[] =
	{ 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 };
static S32 DAYS_PER_MONTH_LEAP[] =
	{ 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 };

static S32 days_from_month(S32 year, S32 month)
{
	llassert_always(1 <= month);
	llassert_always(month <= 12);

	if (year % 4 == 0
		&& year % 100 != 0)
	{
		// leap year
		return DAYS_PER_MONTH_LEAP[month - 1];
	}
	else
	{
		return DAYS_PER_MONTH_NOLEAP[month - 1];
	}
}

bool LLDateUtil::dateFromPDTString(LLDate& date, const std::string& str)
{
	S32 month, day, year;
	S32 matched = sscanf(str.c_str(), "%d/%d/%d", &month, &day, &year);
	if (matched != 3) return false;
	date.fromYMDHMS(year, month, day);
	F64 secs_since_epoch = date.secondsSinceEpoch();
	// Correct for the fact that specified date is in Pacific time, == UTC - 8
	secs_since_epoch += 8.0 * 60.0 * 60.0;
	date.secondsSinceEpoch(secs_since_epoch);
	return true;
}

std::string LLDateUtil::ageFromDate(const LLDate& born_date, const LLDate& now)
{
	S32 born_month, born_day, born_year;
	// explode out to month/day/year again
	born_date.split(&born_year, &born_month, &born_day);

	S32 now_year, now_month, now_day;
	now.split(&now_year, &now_month, &now_day);

	// Do grade-school subtraction, from right-to-left, borrowing from the left
	// when things go negative
	S32 age_days = (now_day - born_day);
	if (age_days < 0)
	{
		now_month -= 1;
		if (now_month == 0)
		{
			now_year -= 1;
			now_month = 12;
		}
		age_days += days_from_month(now_year, now_month);
	}
	S32 age_months = (now_month - born_month);
	if (age_months < 0)
	{
		now_year -= 1;
		age_months += 12;
	}
	S32 age_years = (now_year - born_year);

	// Ansariel: We want total days too!
	S32 age_days_total = (S32)((now.secondsSinceEpoch() - born_date.secondsSinceEpoch()) / 86400);

	// Noun pluralization depends on language
	std::string lang = LLUI::getLanguage();

	// Try for age in round number of years
	LLStringUtil::format_map_t args;

	if (age_months > 0 || age_years > 0)
	{
		args["[AGEYEARS]"] =
			LLTrans::getCountString(lang, "AgeYears", age_years);
		args["[AGEMONTHS]"] =
			LLTrans::getCountString(lang, "AgeMonths", age_months);
		args["[AGEDAYSTOTAL]"] =
			LLTrans::getCountString(lang, "AgeDays", age_days_total);

		// We want to display times like:
		// 2 year 2 months
		// 2 years (implicitly 0 months)
		// 11 months
		if (age_years > 0)
		{
			if (age_months > 0)
			{
				return LLTrans::getString("YearsMonthsOld", args) + LLTrans::getString("TotalDaysOld", args);
			}
			else
			{
				return LLTrans::getString("YearsOld", args) + LLTrans::getString("TotalDaysOld", args);
			}
		}
		else // age_years == 0
		{
			return LLTrans::getString("MonthsOld", args) + LLTrans::getString("TotalDaysOld", args);
		}
	}
	// you're 0 months old, display in weeks or days

	// Now for age in weeks
	S32 age_weeks = age_days / 7;
	age_days = age_days % 7;
	if (age_weeks > 0)
	{
		args["[AGEWEEKS]"] =
			LLTrans::getCountString(lang, "AgeWeeks", age_weeks);

		args["[AGEDAYSTOTAL]"] =
			LLTrans::getCountString(lang, "AgeDays", age_days_total);

		return LLTrans::getString("WeeksOld", args) + LLTrans::getString("TotalDaysOld", args);
	}

	// Down to days now
	if (age_days > 0)
	{
		args["[AGEDAYS]"] =
			LLTrans::getCountString(lang, "AgeDays", age_days);
		return LLTrans::getString("DaysOld", args);
	}

	return LLTrans::getString("TodayOld");
}

std::string LLDateUtil::ageFromDate(const std::string& date_string, const LLDate& now)
{
	LLDate born_date;

	if (!dateFromPDTString(born_date, date_string))
		return "???";

	return ageFromDate(born_date, now);
}

std::string LLDateUtil::ageFromDate(const std::string& date_string)
{
	return ageFromDate(date_string, LLDate::now());
}

//std::string LLDateUtil::ageFromDateISO(const std::string& date_string,
//									   const LLDate& now)
//{
//	S32 born_month, born_day, born_year;
//	S32 matched = sscanf(date_string.c_str(), "%d-%d-%d",
//			&born_year, &born_month, &born_day);
//	if (matched != 3) return "???";
//	date.fromYMDHMS(year, month, day);
//	F64 secs_since_epoch = date.secondsSinceEpoch();
//	// Correct for the fact that specified date is in Pacific time, == UTC - 8
//	secs_since_epoch += 8.0 * 60.0 * 60.0;
//	date.secondsSinceEpoch(secs_since_epoch);
//	return ageFromDate(born_year, born_month, born_day, now);
//}
//
//std::string LLDateUtil::ageFromDateISO(const std::string& date_string)
//{
//	return ageFromDateISO(date_string, LLDate::now());
//}

S32 LLDateUtil::secondsSinceEpochFromString(const std::string& format, const std::string& str)
{
	date_input_facet *facet = new date_input_facet(format);

	std::stringstream ss;
	ss << str;
	ss.imbue(std::locale(ss.getloc(), facet));

	date d;
	ss >> d;

	ptime time_t_date(d);
	ptime time_t_epoch(date(1970,1,1));

	// We assume that the date defined by str is in UTC, so the difference
	// is calculated with no time zone corrections.
	time_duration diff = time_t_date - time_t_epoch;

	return diff.total_seconds();
}


#include <iostream>
#include <algorithm>
#include <numeric>
#include <cassert>

#include <boost/lambda/lambda.hpp>

#include "RandomSource.h"
#include "ContrastiveDivergenceTraining.h"

using namespace boost::lambda;

namespace Grante {

ContrastiveDivergenceTraining::ContrastiveDivergenceTraining(
	FactorGraphModel* fg_model, unsigned int cd_k,
	unsigned int mini_batch_size, double stepsize)
	: ParameterEstimationMethod(fg_model), mini_batch_size(mini_batch_size),
		cd(fg_model, cd_k), stepsize(stepsize) {
	assert(cd_k > 0);
}

ContrastiveDivergenceTraining::~ContrastiveDivergenceTraining() {
	// nothing to do
}

void ContrastiveDivergenceTraining::SetupTrainingData(
	const std::vector<labeled_instance_type>& training_data,
	const std::vector<InferenceMethod*> inference_methods) {
	assert(training_data.size() > 0);
	assert(inference_methods.empty());
	// Call parent
	ParameterEstimationMethod::SetupTrainingData(training_data,
		inference_methods);
}

void ContrastiveDivergenceTraining::SetupPartiallyObservedTrainingData(
	const std::vector<partially_labeled_instance_type>& pobs_training_data) {
	assert(training_data.empty());
	assert(pobs_training_data.size() > 0);

	this->pobs_training_data = pobs_training_data;
}

double ContrastiveDivergenceTraining::Train(double conv_tol,
	unsigned int max_iter) {
	assert(pobs_training_data.empty() ^ training_data.empty());

	size_t instance_count = 0;
	if (training_data.empty() == false) {
		// fully observed training data
		instance_count = training_data.size();
	} else {
		// partially observed training data
		instance_count = pobs_training_data.size();
	}

	// Global scaling constants
	double scale = 1.0 / static_cast<double>(instance_count);

	// Mini batch indices
	std::vector<unsigned int> instance_idx(instance_count);
	for (unsigned int ni = 0; ni < instance_count; ++ni)
		instance_idx[ni] = ni;

	// Parameter gradient
	std::unordered_map<std::string, std::vector<double> >
		parameter_gradient;
	double mean_nabla_w_norm = 0.0;
	for (unsigned int iter = 1; max_iter == 0 || (iter <= max_iter); ++iter) {
		// 1. Setup mini-batches
		RandomSource::ShuffleRandom(instance_idx);

		size_t mb_count = 1;	// number of mini batches
		if (mini_batch_size > 0) {
			mb_count = instance_idx.size() / mini_batch_size;
			if (instance_idx.size() % mini_batch_size > 0)
				mb_count += 1;
		}

		// Process all mini batches
		mean_nabla_w_norm = 0.0;
		for (unsigned int mbi = 0; mbi < mb_count; ++mbi) {
			GradientSetup(parameter_gradient);	// reset gradient

			size_t mb_start = 0;
			size_t mb_end = instance_idx.size();
			if (mb_count > 1) {
				mb_start = mbi*mini_batch_size;
				mb_end = (mbi+1)*mini_batch_size;
				if (mb_end > instance_idx.size())
					mb_end = instance_idx.size();
			}
			assert(mb_end > mb_start);

			// composite scaling term for energy gradients
			double scale_mb = scale;
			if (mb_count > 1)
				scale_mb = 1.0 / static_cast<double>(mb_end-mb_start);

			// 2a. Compute CD parameter gradient, scaled, in order to obtain
			//     an unbiased estimator of the full 1/N \sum_{n=1}^N
			//     CD-gradient.
			for (size_t mb_si = mb_start; mb_si < mb_end; ++mb_si) {
				unsigned int ni = instance_idx[mb_si];
				if (training_data.empty() == false) {
					// fully observed
					cd.ComputeGradientFullyObserved(parameter_gradient,
						training_data[ni].first, training_data[ni].second);
				} else {
					// partially observed
					cd.ComputeGradientPartiallyObserved(parameter_gradient,
						pobs_training_data[ni].first,	// fg
						pobs_training_data[ni].second);	// pobs
				}
			}
			GradientScale(parameter_gradient, scale_mb);

			// 2b. Compute parameter prior gradient: (1/N) \nabla_w -log p(w)
			for (std::multimap<std::string, Prior*>::const_iterator
				prior = priors.begin(); prior != priors.end(); ++prior) {
				FactorType* ft = fg_model->FindFactorType(prior->first);
				prior->second->EvaluateNegLogP(ft->Weights(),
					parameter_gradient[prior->first], scale);
			}

			// 3. Select step size
			double alpha = stepsize;
			//double alpha = 1.0e-1 / std::sqrt(static_cast<double>(iter));

			// 4. Update parameters: w <-- w - alpha * \nabla_w E
			double nabla_w_norm = 0.0;	// norm of approximate gradient
			for (std::unordered_map<std::string,
				std::vector<double> >::iterator pgi =
				parameter_gradient.begin(); pgi != parameter_gradient.end();
				++pgi) {
				nabla_w_norm += std::inner_product(pgi->second.begin(),
					pgi->second.end(), pgi->second.begin(), 0.0);

				// Update model weights
				FactorType* ft = fg_model->FindFactorType(pgi->first);
				std::transform(ft->Weights().begin(), ft->Weights().end(),
					pgi->second.begin(), ft->Weights().begin(),
					_1 - alpha * _2);
			}
			nabla_w_norm = std::sqrt(nabla_w_norm);
			mean_nabla_w_norm += nabla_w_norm;

#if 0
			// TODO
			std::cout << "cd iter " << iter << ", mb " << mbi
				<< " of " << mb_count << ", |nabla w| " << nabla_w_norm
				<< std::endl;
#endif
		}
		mean_nabla_w_norm /= static_cast<double>(mb_count);
		std::cout << "cd iter " << iter << " completed, mean |nabla w| "
			<< mean_nabla_w_norm << std::endl;
	}
	return (mean_nabla_w_norm);
}

void ContrastiveDivergenceTraining::GradientSetup(
	std::unordered_map<std::string, std::vector<double> >&
		parameter_gradient) const {
	const std::vector<FactorType*>& factor_types = fg_model->FactorTypes();
	for (std::vector<FactorType*>::const_iterator fti = factor_types.begin();
		fti != factor_types.end(); ++fti) {
		const std::string& ftname = (*fti)->Name();
		std::unordered_map<std::string, std::vector<double> >::iterator
			pgi = parameter_gradient.find(ftname);
		if (pgi != parameter_gradient.end()) {
			// Reset to zero
			std::fill(pgi->second.begin(), pgi->second.end(), 0.0);
		} else {
			// Create
			parameter_gradient[ftname] =
				std::vector<double>((*fti)->WeightDimension(), 0.0);
		}
	}
}

void ContrastiveDivergenceTraining::GradientScale(
	std::unordered_map<std::string, std::vector<double> >&
		parameter_gradient, double scale) const {
	for (std::unordered_map<std::string, std::vector<double> >::iterator
		pgi = parameter_gradient.begin(); pgi != parameter_gradient.end();
		++pgi) {
		std::transform(pgi->second.begin(), pgi->second.end(),
			pgi->second.begin(), scale * _1);
	}
}

}