#ifdef USE_EIGEN_BACKEND

/** Eigen3 backend.
 *
 * Only supports float32 computation with NHWC memory layout (at runtime and as input).
 */

// CR lpuchallafiore: Add multi-threading support (see "Evaluating with a Thread Pool" in the Eigen Tensor docs).

#include "../neuralnet/nninterface.h"

#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>
#include <zstr/src/zstr.hpp>

#include "../neuralnet/desc.h"
#include "../neuralnet/modelversion.h"
#include "../neuralnet/nninputs.h"
#include "../neuralnet/nneval.h"

using namespace std;
using Eigen::Tensor;
using Eigen::TensorMap;

//Eigen doesn't seem to have a way to make a const tensor map out of a const float* ??
//So we have to cast away qualifiers to build it.
#pragma GCC diagnostic ignored "-Wcast-qual"

// Eigen tensors are stored in column-major order, so an NHWC memory layout is given by Tensor<4>(C,W,H,N).

#define SCALAR float
#define TENSOR2 Tensor<SCALAR, 2>
#define TENSOR3 Tensor<SCALAR, 3>
#define TENSOR4 Tensor<SCALAR, 4>
#define TENSORMAP2 TensorMap<Tensor<SCALAR, 2>>
#define TENSORMAP3 TensorMap<Tensor<SCALAR, 3>>
#define TENSORMAP4 TensorMap<Tensor<SCALAR, 4>>

#define CONSTTENSOR2 const Tensor<SCALAR, 2>
#define CONSTTENSOR3 const Tensor<SCALAR, 3>
#define CONSTTENSOR4 const Tensor<SCALAR, 4>
#define CONSTTENSORMAP2 const TensorMap<Tensor<SCALAR, 2>>
#define CONSTTENSORMAP3 const TensorMap<Tensor<SCALAR, 3>>
#define CONSTTENSORMAP4 const TensorMap<Tensor<SCALAR, 4>>

// Debugging -----------------------------------------------------------------------------------------------------------
// #define DEBUG true

template <typename T>
void printTensorShape(const string& name, const T* t) {
  auto d = t->dimensions();
  cout << name << " rank=" << d.size() << " - (";
  for (int i = 0; i < d.size(); i++) {
    cout << d[i] << ",";
  }
  cout << ")" << endl;
}

#if DEBUG
#define DSHAPE(n, x) printTensorShape(n,x)
#define DTENSOR(n, x) cout << n << *x << endl
#else
#define DSHAPE(n, x)
#define DTENSOR(n, x)
#endif

// LoadedModel / ModelDesc ---------------------------------------------------------------------------------------------

struct LoadedModel {
  ModelDesc modelDesc;

  LoadedModel(const string& fileName) {
    ModelDesc::loadFromFileMaybeGZipped(fileName,modelDesc);
  }

  LoadedModel() = delete;
  LoadedModel(const LoadedModel&) = delete;
  LoadedModel& operator=(const LoadedModel&) = delete;
};

LoadedModel* NeuralNet::loadModelFile(const string& file) {
  LoadedModel* loadedModel = new LoadedModel(file);
  return loadedModel;
}

void NeuralNet::freeLoadedModel(LoadedModel* loadedModel) {
  delete loadedModel;
}

string NeuralNet::getModelName(const LoadedModel* loadedModel) {
  return loadedModel->modelDesc.name;
}

int NeuralNet::getModelVersion(const LoadedModel* loadedModel) {
  return loadedModel->modelDesc.version;
}

Rules NeuralNet::getSupportedRules(const LoadedModel* loadedModel, const Rules& desiredRules, bool& supported) {
  return loadedModel->modelDesc.getSupportedRules(desiredRules, supported);
}

//------------------------------------------------------------------------------

struct ComputeContext {
  int nnXLen;
  int nnYLen;
};

ComputeContext* NeuralNet::createComputeContext(
  const std::vector<int>& gpuIdxs,
  Logger* logger,
  int nnXLen,
  int nnYLen,
  const string& openCLTunerFile,
  const string& homeDataDirOverride,
  bool openCLReTunePerBoardSize,
  enabled_t useFP16Mode,
  enabled_t useNHWCMode,
  const LoadedModel* loadedModel
) {
  (void)gpuIdxs;
  (void)logger;
  (void)openCLTunerFile;
  (void)homeDataDirOverride;
  (void)openCLReTunePerBoardSize;
  (void)loadedModel;

  bool useFP16 = useFP16Mode == enabled_t::True ? true : false;
  bool useNHWC = useNHWCMode == enabled_t::False ? false : true;

  if(useFP16)
    throw StringError("Eigen backend: useFP16 = true not supported");
  if(!useNHWC)
    throw StringError("Eigen backend: useNHWC = false not supported");

  ComputeContext* context = new ComputeContext();
  context->nnXLen = nnXLen;
  context->nnYLen = nnYLen;
  return context;
}

void NeuralNet::freeComputeContext(ComputeContext* computeContext) {
  delete computeContext;
}

// Helpers --------------------------------------------------------------------------------------------------------------

static void computeMaskSum(CONSTTENSORMAP3* mask, float* maskSum) {
  for (int n = 0; n < mask->dimension(2); n++) {
    float s = 0.f;
    for (int h = 0; h < mask->dimension(1); h++) {
      for (int w = 0; w < mask->dimension(0); w++) {
        s += (*mask)(w, h, n);
      }
    }
    maskSum[n] = s;
  }
}

// in NxHxWxC, bias NxC
static void addNCBiasInplace(TENSORMAP4* in, CONSTTENSORMAP2* bias) {
  assert(in->dimension(0) == bias->dimension(0) && in->dimension(3) == bias->dimension(1));
  for (int n = 0; n < in->dimension(3); n++) {
    for (int h = 0; h < in->dimension(2); h++) {
      for (int w = 0; w < in->dimension(1); w++) {
        for (int c = 0; c < in->dimension(0); c++) {
          (*in)(c,w,h,n) += (*bias)(c,n);
        }
      }
    }
  }
}

static void poolRowsGPool(CONSTTENSORMAP4* in, TENSORMAP2* out, const float* maskSum) {
  for (int n = 0; n < in->dimension(3); n++) {
    for (int c = 0; c < in->dimension(0); c++) {
      float s = 0.f;
      float m = 0.f;
      for (int h = 0; h < in->dimension(2); h++) {
        for (int w = 0; w < in->dimension(1); w++) {
          float x = (*in)(c, w, h, n);
          s += x;
          m = max(m, x);
        }
      }
      float div = maskSum[n];
      float sqrtdiv = sqrt(div);
      float mean = s / div;
      (*out)(c, n) = mean;
      (*out)(c + in->dimension(0), n) = mean * (sqrtdiv - 14.f) * 0.1f;
      (*out)(c + 2*in->dimension(0), n) = m;
    }
  }
}

// // Given input [n,w,h,c] fills output of shape [n,c] with sum over c.
// static void poolRowsSum(CONSTTENSORMAP4* in, TENSORMAP2* out, float scaleSum) {
//   for (int n = 0; n < in->dimension(3); n++) {
//     for (int c = 0; c < in->dimension(0); c++) {
//       float s = 0.f;
//       for (int h = 0; h < in->dimension(2); h++) {
//         for (int w = 0; w < in->dimension(1); w++) {
//           float x = in(c, w, h, n);
//           s += x;
//         }
//       }
//       out(c, n) = s * scaleSum;
//     }
//   }
// }

static void poolRowsValueHead(CONSTTENSORMAP4* in, TENSORMAP2* out, const float* maskSum) {
  for (int n = 0; n < in->dimension(3); n++) {
    for (int c = 0; c < in->dimension(0); c++) {
      float s = 0.f;
      for (int h = 0; h < in->dimension(2); h++) {
        for (int w = 0; w < in->dimension(1); w++) {
          float x = (*in)(c, w, h, n);
          s += x;
        }
      }
      float div = maskSum[n];
      float sqrtdiv = sqrt(div);
      float mean = s / div;
      (*out)(c, n) = mean;
      (*out)(c + in->dimension(0), n) = mean * (sqrtdiv - 14.f) * 0.1f;
      (*out)(c + 2*in->dimension(0), n) = mean * ((sqrtdiv - 14.0f) * (sqrtdiv - 14.0f) * 0.01f - 0.1f);
    }
  }
}

// Layers --------------------------------------------------------------------------------------------------------------

// Convolution layer with zero-padding.
struct ConvLayer {
  string name;

  Eigen::array<pair<int, int>, 4> paddings;
  TENSOR4 kernel;
  int inChannels, outChannels;

  ConvLayer() = delete;
  ConvLayer(const ConvLayer&) = delete;
  ConvLayer& operator=(const ConvLayer&) = delete;

  ConvLayer(const ConvLayerDesc& desc) {
    name = desc.name;
    int convYSize = desc.convYSize;
    int convXSize = desc.convXSize;
    inChannels = desc.inChannels;
    outChannels = desc.outChannels;
    //Currently eigen impl doesn't support dilated convs
    int dilationY = desc.dilationY;
    int dilationX = desc.dilationX;
    int paddingX = (convXSize / 2) * dilationX;
    int paddingY = (convYSize / 2) * dilationY;

    if(dilationX != 1 || dilationY != 1)
      throw StringError("Eigen backend: Encountered convolution dilation factors other than 1, not supported");

    assert(convXSize % 2 == 1);
    assert(convYSize % 2 == 1);

    paddings[0] = make_pair(0, 0);                // C
    paddings[1] = make_pair(paddingX, paddingX);  // W
    paddings[2] = make_pair(paddingY, paddingY);  // H
    paddings[3] = make_pair(0, 0);                // N

    // CR-someday lpuchallafiore: optimize NHWC vs NCHW, etc.
    kernel = TensorMap<const Tensor<const SCALAR, 4>>(
      desc.weights.data(), convXSize, convYSize, inChannels, outChannels);
  }

  void apply(CONSTTENSORMAP4* input, TENSORMAP4* output, bool accumulate) const {
    auto padded = input->pad(paddings);
    assert(output->dimension(0) == outChannels);
    for(int n = 0; n < input->dimension(3); n++) {
      auto inN = padded.chip(n, 3);
      for(int oc = 0; oc < outChannels; oc++) {
        TENSOR2 sum(input->dimension(1), input->dimension(2));
        sum.setZero();

        for(int ic = 0; ic < inChannels; ic++) {
          Eigen::array<ptrdiff_t, 2> dims({0, 1});
          auto kChip = kernel.chip(oc, 3).chip(ic, 2);
          auto inNC = inN.chip(ic, 0);
          sum += inNC.convolve(kChip, dims);
        }

        if(accumulate)
          output->chip(n, 3).chip(oc, 0) += sum;
        else
          output->chip(n, 3).chip(oc, 0) = sum;
      }
    }
  }
};

//--------------------------------------------------------------

struct BatchNormLayer {
  string name;

  vector<float> mergedScale;
  vector<float> mergedBias;

  BatchNormLayer() = delete;
  BatchNormLayer(const BatchNormLayer&) = delete;
  BatchNormLayer& operator=(const BatchNormLayer&) = delete;

  BatchNormLayer(const BatchNormLayerDesc& desc) {
    name = desc.name;
    int numChannels = desc.numChannels;
    float epsilon = desc.epsilon;

    mergedScale.resize(numChannels);
    mergedBias.resize(numChannels);
    for(int c = 0; c < numChannels; c++) {
      mergedScale[c] = desc.scale[c] / sqrt(desc.variance[c] + epsilon);
      mergedBias[c] = desc.bias[c] - mergedScale[c] * desc.mean[c];
    }
  }

  // Mask should be in 'NHW' format (no "C" channel).
  void apply(
    bool applyRelu,
    CONSTTENSORMAP4* input,
    TENSORMAP4* output,
    CONSTTENSORMAP3* mask
  ) const {
    for(int c = 0; c < input->dimension(0); c++) {
      auto inC = input->chip(c, 0);
      auto x = inC * mergedScale[c] + mergedBias[c];
      auto z = TENSOR3(mask->dimension(0), mask->dimension(1), mask->dimension(2)).setZero();
      if(applyRelu)
        output->chip(c, 0) = (*mask == 1.f).select(x.cwiseMax(0.f), z);
      else
        output->chip(c, 0) = (*mask == 1.f).select(x, z);
    }
  }
};

//--------------------------------------------------------------

struct ActivationLayer {
  string name;

  ActivationLayer() = delete;
  ActivationLayer(const ActivationLayer&) = delete;
  ActivationLayer& operator=(const ActivationLayer&) = delete;

  ActivationLayer(const ActivationLayerDesc& desc) { name = desc.name; }

  template <int N>
  void apply(const Tensor<SCALAR, N>* input, Tensor<SCALAR, N>* output) const { *output = input->cwiseMax(0.f); }
  template <int N>
  void apply(const TensorMap<Tensor<SCALAR, N>>* input, TensorMap<Tensor<SCALAR, N>>* output) const { *output = input->cwiseMax(0.f); }
};

//--------------------------------------------------------------

struct MatMulLayer {
  string name;
  TENSOR2 weights;

  MatMulLayer() = delete;
  MatMulLayer(const MatMulLayer&) = delete;
  MatMulLayer& operator=(const MatMulLayer&) = delete;

  MatMulLayer(const MatMulLayerDesc& desc)
    : name(desc.name)
  {
    weights = TENSOR2(desc.outChannels, desc.inChannels);
    memcpy(weights.data(), desc.weights.data(), sizeof(SCALAR) * weights.size());
  }

  void apply(CONSTTENSORMAP2* in, TENSORMAP2* out) const {
    Eigen::array<Eigen::IndexPair<int>, 1> product_dims = { Eigen::IndexPair<int>(1, 0) };
    *out = weights.contract(*in, product_dims);
  }
};

struct MatBiasLayer {
  string name;
  std::vector<float> weights;

  MatBiasLayer() = delete;
  MatBiasLayer(const MatBiasLayer&) = delete;
  MatBiasLayer& operator=(const MatBiasLayer&) = delete;

  MatBiasLayer(const MatBiasLayerDesc& desc)
    : name(desc.name),
      weights(desc.weights) {}

  void apply(TENSORMAP2* mat) const {
    for(int n = 0; n < mat->dimension(1); n++) {
      for(int c = 0; c < mat->dimension(0); c++) {
        (*mat)(c, n) += weights[c];
      }
    }
  }
};

// Blocks
// --------------------------------------------------------------------------------------------------------------

struct ResidualBlockIntf {
  virtual ~ResidualBlockIntf(){}

  virtual void apply(
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const = 0;
};

struct ResidualBlock final : public ResidualBlockIntf {
  string name;
  BatchNormLayer preBN;
  ConvLayer regularConv;
  BatchNormLayer midBN;
  ConvLayer finalConv;

  ResidualBlock() = delete;
  ResidualBlock(const ResidualBlock&) = delete;
  ResidualBlock& operator=(const ResidualBlock&) = delete;

  ~ResidualBlock(){}

  ResidualBlock(const ResidualBlockDesc& desc)
    : name(desc.name),
      preBN(desc.preBN),
      regularConv(desc.regularConv),
      midBN(desc.midBN),
      finalConv(desc.finalConv) {}

  void apply(
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const override {
    (void)regularOut;
    (void)regularScratch;
    (void)gpoolOut;
    (void)gpoolOut2;
    (void)gpoolConcat;
    (void)gpoolBias;
    (void)maskSum;
    const bool applyBNRelu = true;
    preBN.apply(applyBNRelu, trunk, trunkScratch, mask);
    regularConv.apply(trunkScratch, midIn, false);
    midBN.apply(applyBNRelu, midIn, midScratch, mask);
    finalConv.apply(midScratch, trunk, true);
  }
};

// // Given two tensors with shapes inA: [n, h, w, cA] and inB: [n, h, w, cB]
// // Copy them into a single tensor out: [n, h, w, cA + cB]
// TENSOR4 concatTensors(CONSTTENSOR4& a, CONSTTENSOR4& b) {
//   assert(a->dimension(1) == b->dimension(1) && a->dimension(2) == b->dimension(2) && a->dimension(3) == b->dimension(3));
//   TENSOR4 x = TENSOR4(/* C */ a->dimension(0) + b->dimension(0),
//                                           /* W */ a->dimension(1),
//                                           /* H */ a->dimension(2),
//                                           /* N */ a->dimension(3));
//   for (int n = 0; n < a->dimension(3); n++) {
//     for (int h = 0; h < a->dimension(2); h++) {
//       for (int w = 0; w < a->dimension(1); w++) {
//         int c = 0;
//         for (int ca = 0; a->dimension(0); ca++, c++) {
//           x(c,w,h,n) = a(ca,w,h,n);
//         }
//         for (int cb = 0; b->dimension(0); cb++, c++) {
//           x(c,w,h,n) = b(cb,w,h,n);
//         }
//       }
//     }
//   }
//   return x;
// }

struct GlobalPoolingResidualBlock final : public ResidualBlockIntf {
  string name;
  BatchNormLayer preBN;
  ActivationLayer preActivation;
  ConvLayer regularConv;
  ConvLayer gpoolConv;
  BatchNormLayer gpoolBN;
  ActivationLayer gpoolActivation;
  MatMulLayer gpoolToBiasMul;
  BatchNormLayer midBN;
  ActivationLayer midActivation;
  ConvLayer finalConv;

  GlobalPoolingResidualBlock() = delete;
  GlobalPoolingResidualBlock(const GlobalPoolingResidualBlock&) = delete;
  GlobalPoolingResidualBlock& operator=(const GlobalPoolingResidualBlock&) = delete;

  ~GlobalPoolingResidualBlock(){}

  GlobalPoolingResidualBlock(const GlobalPoolingResidualBlockDesc& desc)
    : name(desc.name),
      preBN(desc.preBN),
      preActivation(desc.preActivation),
      regularConv(desc.regularConv),
      gpoolConv(desc.gpoolConv),
      gpoolBN(desc.gpoolBN),
      gpoolActivation(desc.gpoolActivation),
      gpoolToBiasMul(desc.gpoolToBiasMul),
      midBN(desc.midBN),
      midActivation(desc.midActivation),
      finalConv(desc.finalConv) {}

  void apply(
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const override {
    (void)midIn;
    (void)midScratch;
    const bool applyBNRelu = true;
    DTENSOR("trunk", trunk);
    DTENSOR("mask", mask);
    preBN.apply(applyBNRelu, trunk, trunkScratch, mask);
    DTENSOR("trunkScratch", trunkScratch);
    regularConv.apply(trunkScratch, regularOut, false);
    DTENSOR("regularOut", regularOut);
    gpoolConv.apply(trunkScratch, gpoolOut, false);
    DTENSOR("gpoolOut", gpoolOut);
    gpoolBN.apply(applyBNRelu, gpoolOut, gpoolOut2, mask);
    DTENSOR("gpoolOut2", gpoolOut2);
    poolRowsGPool(gpoolOut2, gpoolConcat, maskSum);
    gpoolToBiasMul.apply(gpoolConcat, gpoolBias);
    addNCBiasInplace(regularOut, gpoolBias);
    midBN.apply(applyBNRelu, regularOut, regularScratch, mask);
    finalConv.apply(regularScratch, trunk, true);
    DSHAPE("trunk", trunk);
    DSHAPE("trunkScratch", trunkScratch);
    DSHAPE("regularOut", regularOut);
    DSHAPE("gpoolOut", gpoolOut);
    DSHAPE("gpoolOut2", gpoolOut2);
    DSHAPE("gpoolConcat", gpoolConcat);
    DSHAPE("gpoolBias", gpoolBias);
    DSHAPE("mask", mask);
  }
};

struct Trunk {
  string name;
  int version;
  int numBlocks;

  ConvLayer initialConv;
  MatMulLayer initialMatMul;
  vector<pair<int, ResidualBlockIntf*>> blocks;
  BatchNormLayer trunkTipBN;
  ActivationLayer trunkTipActivation;

  Trunk() = delete;
  Trunk(const Trunk&) = delete;
  Trunk& operator=(const Trunk&) = delete;

  Trunk(const TrunkDesc& desc)
    : name(desc.name),
      version(desc.version),
      numBlocks(desc.numBlocks),
      initialConv(desc.initialConv),
      initialMatMul(desc.initialMatMul),
      trunkTipBN(desc.trunkTipBN),
      trunkTipActivation(desc.trunkTipActivation)
  {
    for (int i = 0; i < numBlocks; ++i) {
      if (desc.blocks[i].first == ORDINARY_BLOCK_KIND) {
        ResidualBlockDesc* blockDesc = (ResidualBlockDesc*)desc.blocks[i].second;
        ResidualBlockIntf* block = new ResidualBlock(*blockDesc);
        blocks.push_back(make_pair(ORDINARY_BLOCK_KIND, block));
      }
      else if (desc.blocks[i].first == DILATED_BLOCK_KIND) {
        throw StringError("Eigen backend: Dilated residual blocks are not supported right now");
      }
      else if (desc.blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) {
        GlobalPoolingResidualBlockDesc* blockDesc = (GlobalPoolingResidualBlockDesc*)desc.blocks[i].second;
        GlobalPoolingResidualBlock* block = new GlobalPoolingResidualBlock(*blockDesc);
        blocks.push_back(make_pair(GLOBAL_POOLING_BLOCK_KIND, block));
      }
      else {
        ASSERT_UNREACHABLE;
      }
    }
  }

  virtual ~Trunk() {
    for (auto p : blocks) {
      delete p.second;
    }
  }

  void apply(
    CONSTTENSORMAP4* input,
    CONSTTENSORMAP2* inputGlobal,
    TENSORMAP2* inputMatMulOut,
    TENSORMAP4* trunk,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const {

    initialConv.apply(input, trunkScratch, false);
    initialMatMul.apply(inputGlobal, inputMatMulOut);
    addNCBiasInplace(trunkScratch, inputMatMulOut);

    // apply blocks
    // Flip trunkBuf and trunkScratchBuf so that the result gets accumulated in trunkScratchBuf
    for (auto block : blocks) {
      block.second->apply(
        trunkScratch,
        trunk,
        regularOut,
        regularScratch,
        midIn,
        midScratch,
        gpoolOut,
        gpoolOut2,
        gpoolConcat,
        gpoolBias,
        mask,
        maskSum
      );
    }

    // And now with the final BN port it from trunkScratchBuf to trunkBuf.
    const bool applyBNRelu = true;
    trunkTipBN.apply(applyBNRelu, trunkScratch, trunk, mask);
  }
};

struct PolicyHead {
  string name;
  int version;

  ConvLayer p1Conv;
  ConvLayer g1Conv;
  BatchNormLayer g1BN;
  ActivationLayer g1Activation;
  MatMulLayer gpoolToBiasMul;
  BatchNormLayer p1BN;
  ActivationLayer p1Activation;
  ConvLayer p2Conv;
  MatMulLayer gpoolToPassMul;

  PolicyHead() = delete;
  PolicyHead(const PolicyHead&) = delete;
  PolicyHead& operator=(const PolicyHead&) = delete;

  PolicyHead(const PolicyHeadDesc& desc)
    : name(desc.name),
      version(desc.version),
      p1Conv(desc.p1Conv),
      g1Conv(desc.g1Conv),
      g1BN(desc.g1BN),
      g1Activation(desc.g1Activation),
      gpoolToBiasMul(desc.gpoolToBiasMul),
      p1BN(desc.p1BN),
      p1Activation(desc.p1Activation),
      p2Conv(desc.p2Conv),
      gpoolToPassMul(desc.gpoolToPassMul) {}

  void apply(
    CONSTTENSORMAP4* trunk,
    TENSORMAP4* p1Out,
    TENSORMAP4* p1Out2,
    TENSORMAP4* g1Out,
    TENSORMAP4* g1Out2,
    TENSORMAP2* g1Concat,
    TENSORMAP2* g1Bias,
    TENSORMAP2* policyPass,
    TENSORMAP4* policy,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const {
    const bool applyBNRelu = true;
    p1Conv.apply(trunk, p1Out, false);
    g1Conv.apply(trunk, g1Out, false);
    g1BN.apply(applyBNRelu, g1Out, g1Out2, mask);
    poolRowsGPool(g1Out2, g1Concat, maskSum);
    gpoolToBiasMul.apply(g1Concat, g1Bias);
    addNCBiasInplace(p1Out, g1Bias);
    p1BN.apply(true, p1Out, p1Out2, mask);
    p2Conv.apply(p1Out2, policy, false);
    gpoolToPassMul.apply(g1Concat, policyPass);
  }
};

struct ValueHead {
  string name;
  int version;

  ConvLayer v1Conv;
  BatchNormLayer v1BN;
  ActivationLayer v1Activation;
  MatMulLayer v2Mul;
  MatBiasLayer v2Bias;
  ActivationLayer v2Activation;
  MatMulLayer v3Mul;
  MatBiasLayer v3Bias;
  MatMulLayer sv3Mul;
  MatBiasLayer sv3Bias;
  ConvLayer vOwnershipConv;

  ValueHead() = delete;
  ValueHead(const ValueHead&) = delete;
  ValueHead& operator=(const ValueHead&) = delete;

  ValueHead(const ValueHeadDesc& desc)
    : name(desc.name),
      version(desc.version),
      v1Conv(desc.v1Conv),
      v1BN(desc.v1BN),
      v1Activation(desc.v1Activation),
      v2Mul(desc.v2Mul),
      v2Bias(desc.v2Bias),
      v2Activation(desc.v2Activation),
      v3Mul(desc.v3Mul),
      v3Bias(desc.v3Bias),
      sv3Mul(desc.sv3Mul),
      sv3Bias(desc.sv3Bias),
      vOwnershipConv(desc.vOwnershipConv) {}

  void apply(
    CONSTTENSORMAP4* trunk,
    TENSORMAP4* v1Out,
    TENSORMAP4* v1Out2,
    TENSORMAP2* v1Mean,
    TENSORMAP2* v2Out,
    TENSORMAP2* value,
    TENSORMAP2* scoreValue,
    TENSORMAP4* ownership,
    CONSTTENSORMAP3* mask,
    const float* maskSum
  ) const {
    bool applyBNRelu = true;
    v1Conv.apply(trunk, v1Out, false);
    v1BN.apply(applyBNRelu, v1Out, v1Out2, mask);
    poolRowsValueHead(v1Out2, v1Mean, maskSum);
    v2Mul.apply(v1Mean, v2Out);
    v2Bias.apply(v2Out);
    v2Activation.apply(v2Out, v2Out);
    v3Mul.apply(v2Out, value);
    v3Bias.apply(value);

    sv3Mul.apply(v2Out, scoreValue);
    sv3Bias.apply(scoreValue);

    vOwnershipConv.apply(v1Out2, ownership, false);
  }
};

// Model and Buffer I/O ------------------------------------------------------------------------------------------------

struct Model {
  string name;
  int version;
  int numInputChannels;
  int numInputGlobalChannels;
  int numValueChannels;
  int numScoreValueChannels;
  int numOwnershipChannels;

  Trunk trunk;
  PolicyHead policyHead;
  ValueHead valueHead;

  Model() = delete;
  Model(const Model&) = delete;
  Model& operator=(const Model&) = delete;

  Model(const ModelDesc& desc)
    : name(desc.name), version(desc.version), numInputChannels(desc.numInputChannels),
      numInputGlobalChannels(desc.numInputGlobalChannels),
      numValueChannels(desc.numValueChannels),
      numScoreValueChannels(desc.numScoreValueChannels),
      numOwnershipChannels(desc.numOwnershipChannels),
      trunk(desc.trunk),
      policyHead(desc.policyHead),
      valueHead(desc.valueHead) {}

  void apply(
    CONSTTENSORMAP4* input,
    CONSTTENSORMAP2* inputGlobal,
    TENSORMAP2* inputMatMulOut,
    TENSORMAP4* trunkBuf,
    TENSORMAP4* trunkScratch,
    TENSORMAP4* regularOut,
    TENSORMAP4* regularScratch,
    TENSORMAP4* midIn,
    TENSORMAP4* midScratch,
    TENSORMAP4* gpoolOut,
    TENSORMAP4* gpoolOut2,
    TENSORMAP2* gpoolConcat,
    TENSORMAP2* gpoolBias,

    TENSORMAP4* p1Out,
    TENSORMAP4* p1Out2,
    TENSORMAP4* g1Out,
    TENSORMAP4* g1Out2,
    TENSORMAP2* g1Concat,
    TENSORMAP2* g1Bias,
    TENSORMAP2* policyPass,
    TENSORMAP4* policy,

    TENSORMAP4* v1Out,
    TENSORMAP4* v1Out2,
    TENSORMAP2* v1Mean,
    TENSORMAP2* v2Out,
    TENSORMAP2* value,
    TENSORMAP2* scoreValue,
    TENSORMAP4* ownership,

    TENSORMAP3* mask,
    float* maskSum
  ) const {
    *mask = input->chip(0,0);
    computeMaskSum(mask,maskSum);

    trunk.apply(
      input,
      inputGlobal,
      inputMatMulOut,
      trunkBuf,
      trunkScratch,
      regularOut,
      regularScratch,
      midIn,
      midScratch,
      gpoolOut,
      gpoolOut2,
      gpoolConcat,
      gpoolBias,
      mask,
      maskSum
    );
    policyHead.apply(
      trunkBuf,
      p1Out,
      p1Out2,
      g1Out,
      g1Out2,
      g1Concat,
      g1Bias,
      policyPass,
      policy,
      mask,
      maskSum
    );
    valueHead.apply(
      trunkBuf,
      v1Out,
      v1Out2,
      v1Mean,
      v2Out,
      value,
      scoreValue,
      ownership,
      mask,
      maskSum
    );
  }
};

//--------------------------------------------------------------

struct Buffers {
  TENSOR2 inputMatMulOut;
  TENSOR4 trunk;
  TENSOR4 trunkScratch;
  TENSOR4 regularOut;
  TENSOR4 regularScratch;
  TENSOR4 midIn;
  TENSOR4 midScratch;
  TENSOR4 gpoolOut;
  TENSOR4 gpoolOut2;
  TENSOR2 gpoolConcat;
  TENSOR2 gpoolBias;

  TENSOR4 p1Out;
  TENSOR4 p1Out2;
  TENSOR4 g1Out;
  TENSOR4 g1Out2;
  TENSOR2 g1Concat;
  TENSOR2 g1Bias;
  TENSOR2 policyPass;
  TENSOR4 policy;

  TENSOR4 v1Out;
  TENSOR4 v1Out2;
  TENSOR2 v1Mean;
  TENSOR2 v2Out;
  TENSOR2 value;
  TENSOR2 scoreValue;
  TENSOR4 ownership;

  TENSOR3 mask;
  vector<float> maskSum;

  Buffers(
    const ModelDesc& desc,
    int maxBatchSize,
    int nnXLen,
    int nnYLen
  ) :
    inputMatMulOut(desc.trunk.trunkNumChannels, maxBatchSize),
    trunk(desc.trunk.trunkNumChannels, nnXLen, nnYLen, maxBatchSize),
    trunkScratch(desc.trunk.trunkNumChannels, nnXLen, nnYLen, maxBatchSize),
    regularOut(desc.trunk.regularNumChannels, nnXLen, nnYLen, maxBatchSize),
    regularScratch(desc.trunk.regularNumChannels, nnXLen, nnYLen, maxBatchSize),
    midIn(desc.trunk.midNumChannels, nnXLen, nnYLen, maxBatchSize),
    midScratch(desc.trunk.midNumChannels, nnXLen, nnYLen, maxBatchSize),
    gpoolOut(desc.trunk.gpoolNumChannels, nnXLen, nnYLen, maxBatchSize),
    gpoolOut2(desc.trunk.gpoolNumChannels, nnXLen, nnYLen, maxBatchSize),
    gpoolConcat(desc.trunk.gpoolNumChannels*3, maxBatchSize),
    gpoolBias(desc.trunk.regularNumChannels, maxBatchSize),

    p1Out(desc.policyHead.p1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    p1Out2(desc.policyHead.p1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    g1Out(desc.policyHead.g1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    g1Out2(desc.policyHead.g1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    g1Concat(desc.policyHead.g1Conv.outChannels*3, maxBatchSize),
    g1Bias(desc.policyHead.gpoolToBiasMul.outChannels, maxBatchSize),
    policyPass(desc.policyHead.gpoolToPassMul.outChannels, maxBatchSize),
    policy(desc.policyHead.p2Conv.outChannels, nnXLen, nnYLen, maxBatchSize),

    v1Out(desc.valueHead.v1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    v1Out2(desc.valueHead.v1Conv.outChannels, nnXLen, nnYLen, maxBatchSize),
    v1Mean(desc.valueHead.v1Conv.outChannels*3, maxBatchSize),
    v2Out(desc.valueHead.v2Mul.outChannels, maxBatchSize),
    value(desc.valueHead.v3Mul.outChannels, maxBatchSize),
    scoreValue(desc.valueHead.sv3Mul.outChannels, maxBatchSize),
    ownership(desc.valueHead.vOwnershipConv.outChannels, nnXLen, nnYLen, maxBatchSize),

    mask(nnXLen, nnYLen, maxBatchSize),
    maskSum(maxBatchSize)
  {}
};

//--------------------------------------------------------------

struct InputBuffers {
  int maxBatchSize;

  size_t singleInputElts;
  size_t singleInputGlobalElts;

  size_t singlePolicyPassResultElts;
  size_t singlePolicyResultElts;
  size_t singleValueResultElts;
  size_t singleScoreValueResultElts;
  size_t singleOwnershipResultElts;

  std::vector<float> spatialInput;
  std::vector<float> globalInput;

  InputBuffers(const LoadedModel* loadedModel, int maxBatchSz, int nnXLen, int nnYLen) {
    const ModelDesc& m = loadedModel->modelDesc;

    int xSize = nnXLen;
    int ySize = nnYLen;

    maxBatchSize = maxBatchSz;
    singleInputElts = m.numInputChannels * xSize * ySize;
    singleInputGlobalElts = m.numInputGlobalChannels;

    singlePolicyPassResultElts = (size_t)(1);
    singlePolicyResultElts = (size_t)(xSize * ySize);
    singleValueResultElts = (size_t)m.numValueChannels;
    singleScoreValueResultElts = (size_t)m.numScoreValueChannels;
    singleOwnershipResultElts = (size_t)m.numOwnershipChannels * xSize * ySize;

    assert(NNModelVersion::getNumSpatialFeatures(m.version) == m.numInputChannels);
    assert(NNModelVersion::getNumGlobalFeatures(m.version) == m.numInputGlobalChannels);

    spatialInput = vector<float>(m.numInputChannels * xSize * ySize * maxBatchSize);
    globalInput = vector<float>(m.numInputGlobalChannels * maxBatchSize);
  }

  ~InputBuffers() { }

  InputBuffers() = delete;
  InputBuffers(const InputBuffers&) = delete;
  InputBuffers& operator=(const InputBuffers&) = delete;
};

InputBuffers* NeuralNet::createInputBuffers(const LoadedModel* loadedModel, int maxBatchSize, int nnXLen, int nnYLen) {
  return new InputBuffers(loadedModel, maxBatchSize, nnXLen, nnYLen);
}
void NeuralNet::freeInputBuffers(InputBuffers* inputBuffers) {
  delete inputBuffers;
}

// float* NeuralNet::getBatchEltSpatialInplace(InputBuffers* inputBuffers, int nIdx) {
//   assert(nIdx < inputBuffers->maxBatchSize);
//   return inputBuffers->spatialInput.data() + (inputBuffers->singleInputElts * nIdx);
// }

// float* NeuralNet::getBatchEltGlobalInplace(InputBuffers* inputBuffers, int rowIdx) {
//   assert(rowIdx < inputBuffers->maxBatchSize);
//   return inputBuffers->globalInput.data() + (inputBuffers->singleInputGlobalElts * rowIdx);
// }

// int NeuralNet::getBatchEltSpatialLen(const InputBuffers* inputBuffers) {
//   return inputBuffers->singleInputElts;
// }
// int NeuralNet::getBatchEltGlobalLen(const InputBuffers* inputBuffers) {
//   return inputBuffers->singleInputGlobalElts;
// }

// bool* NeuralNet::getSymmetriesInplace(InputBuffers* inputBuffers) {
//   return inputBuffers->symmetriesBuffer;
// }

// NeuralNet -----------------------------------------------------------------------------------------------------------

void NeuralNet::globalInitialize() {
  // no-op for cpu
}

void NeuralNet::globalCleanup() {
  // no-op for cpu
}

struct ComputeHandle {
  const ComputeContext* context;
  int maxBatchSize;
  bool inputsUseNHWC;
  Model model;
  Buffers buffers;

  ComputeHandle() = delete;
  ComputeHandle(const ComputeHandle&) = delete;
  ComputeHandle& operator=(const ComputeHandle&) = delete;

  ComputeHandle(const ComputeContext* ctx, const LoadedModel& loadedModel, int maxBSize, bool iNHWC)
    : context(ctx),
      maxBatchSize(maxBSize),
      inputsUseNHWC(iNHWC),
      model(loadedModel.modelDesc),
      buffers(loadedModel.modelDesc,maxBSize,ctx->nnXLen,ctx->nnYLen)
  {}
};

ComputeHandle* NeuralNet::createComputeHandle(
  ComputeContext* context,
  const LoadedModel* loadedModel,
  Logger* logger,
  int maxBatchSize,
  bool requireExactNNLen,
  bool inputsUseNHWC,
  int gpuIdxForThisThread
) {
  if(logger != NULL) {
    logger->write("Eigen (CPU) backend: Model version " + Global::intToString(loadedModel->modelDesc.version));
    logger->write("Eigen (CPU) backend: Model name: " + loadedModel->modelDesc.name);
  }

  (void)requireExactNNLen; //We don't bother with mask optimizations if we know exact sizes right now.
  (void)gpuIdxForThisThread; //Doesn't matter

  if(!inputsUseNHWC)
    throw StringError("Eigen backend: inputsUseNHWC = false unsupported");
  return new ComputeHandle(context, *loadedModel, maxBatchSize, inputsUseNHWC);
}

void NeuralNet::freeComputeHandle(ComputeHandle* gpuHandle) {
  delete gpuHandle;
}

void NeuralNet::getOutput(
  ComputeHandle* computeHandle,
  InputBuffers* inputBuffers,
  int numBatchEltsFilled,
  NNResultBuf** inputBufs,
  int symmetry,
  vector<NNOutput*>& outputs
) {
  assert(numBatchEltsFilled <= inputBuffers->maxBatchSize);
  assert(numBatchEltsFilled > 0);
  int batchSize = numBatchEltsFilled;
  int nnXLen = computeHandle->context->nnXLen;
  int nnYLen = computeHandle->context->nnYLen;
  int version = computeHandle->model.version;

  int numSpatialFeatures = NNModelVersion::getNumSpatialFeatures(version);
  int numGlobalFeatures = NNModelVersion::getNumGlobalFeatures(version);
  assert(numSpatialFeatures == computeHandle->model.numInputChannels);
  assert(numSpatialFeatures * nnXLen * nnYLen == inputBuffers->singleInputElts);
  assert(numGlobalFeatures == inputBuffers->singleInputGlobalElts);

  for(int nIdx = 0; nIdx<batchSize; nIdx++) {
    float* rowSpatialInput = inputBuffers->spatialInput.data() + (inputBuffers->singleInputElts * nIdx);
    float* rowGlobalInput = inputBuffers->globalInput.data() + (inputBuffers->singleInputGlobalElts * nIdx);

    const float* rowGlobal = inputBufs[nIdx]->rowGlobal;
    const float* rowSpatial = inputBufs[nIdx]->rowSpatial;
    std::copy(rowGlobal,rowGlobal+numGlobalFeatures,rowGlobalInput);
    SymmetryHelpers::copyInputsWithSymmetry(rowSpatial, rowSpatialInput, 1, nnYLen, nnXLen, numSpatialFeatures, computeHandle->inputsUseNHWC, symmetry);
  }

  Buffers& buffers = computeHandle->buffers;

  CONSTTENSORMAP4 input(inputBuffers->spatialInput.data(), numSpatialFeatures, nnXLen, nnYLen, batchSize);
  CONSTTENSORMAP2 inputGlobal(inputBuffers->globalInput.data(), numGlobalFeatures, batchSize);

#define MAP4(NAME) TENSORMAP4 NAME(buffers.NAME.data(), buffers.NAME.dimension(0), buffers.NAME.dimension(1), buffers.NAME.dimension(2), batchSize)
#define MAP3(NAME) TENSORMAP3 NAME(buffers.NAME.data(), buffers.NAME.dimension(0), buffers.NAME.dimension(1), batchSize)
#define MAP2(NAME) TENSORMAP2 NAME(buffers.NAME.data(), buffers.NAME.dimension(0), batchSize)

  MAP2(inputMatMulOut);
  MAP4(trunk);
  MAP4(trunkScratch);
  MAP4(regularOut);
  MAP4(regularScratch);
  MAP4(midIn);
  MAP4(midScratch);
  MAP4(gpoolOut);
  MAP4(gpoolOut2);
  MAP2(gpoolConcat);
  MAP2(gpoolBias);
  MAP4(p1Out);
  MAP4(p1Out2);
  MAP4(g1Out);
  MAP4(g1Out2);
  MAP2(g1Concat);
  MAP2(g1Bias);
  MAP2(policyPass);
  MAP4(policy);
  MAP4(v1Out);
  MAP4(v1Out2);
  MAP2(v1Mean);
  MAP2(v2Out);
  MAP2(value);
  MAP2(scoreValue);
  MAP4(ownership);
  MAP3(mask);
  vector<float>& maskSum = buffers.maskSum;

  computeMaskSum(&mask,maskSum.data());

  computeHandle->model.apply(
    &input,
    &inputGlobal,
    &inputMatMulOut,
    &trunk,
    &trunkScratch,
    &regularOut,
    &regularScratch,
    &midIn,
    &midScratch,
    &gpoolOut,
    &gpoolOut2,
    &gpoolConcat,
    &gpoolBias,
    &p1Out,
    &p1Out2,
    &g1Out,
    &g1Out2,
    &g1Concat,
    &g1Bias,
    &policyPass,
    &policy,
    &v1Out,
    &v1Out2,
    &v1Mean,
    &v2Out,
    &value,
    &scoreValue,
    &ownership,
    &mask,
    maskSum.data()
  );

  assert(outputs.size() == batchSize);

  float* policyData = policy.data();
  float* policyPassData = policyPass.data();
  float* valueData = value.data();
  float* scoreValueData = scoreValue.data();
  float* ownershipData = ownership.data();

  for(int row = 0; row < batchSize; row++) {
    NNOutput* output = outputs[row];
    assert(output->nnXLen == nnXLen);
    assert(output->nnYLen == nnYLen);

    const float* policySrcBuf = policyData + row * inputBuffers->singlePolicyResultElts;
    float* policyProbs = output->policyProbs;

    //These are not actually correct, the client does the postprocessing to turn them into
    //policy probabilities and white game outcome probabilities
    //Also we don't fill in the nnHash here either
    SymmetryHelpers::copyOutputsWithSymmetry(policySrcBuf, policyProbs, 1, nnYLen, nnXLen, symmetry);
    policyProbs[inputBuffers->singlePolicyResultElts] = policyPassData[row];

    int numValueChannels = computeHandle->model.numValueChannels;
    assert(numValueChannels == 3);
    output->whiteWinProb = valueData[row * numValueChannels];
    output->whiteLossProb = valueData[row * numValueChannels + 1];
    output->whiteNoResultProb = valueData[row * numValueChannels + 2];

    //As above, these are NOT actually from white's perspective, but rather the player to move.
    //As usual the client does the postprocessing.
    if(output->whiteOwnerMap != NULL) {
      const float* ownershipSrcBuf = ownershipData + row * nnXLen * nnYLen;
      assert(computeHandle->model.numOwnershipChannels == 1);
      SymmetryHelpers::copyOutputsWithSymmetry(ownershipSrcBuf, output->whiteOwnerMap, 1, nnYLen, nnXLen, symmetry);
    }

    if(version >= 8) {
      int numScoreValueChannels = computeHandle->model.numScoreValueChannels;
      assert(numScoreValueChannels == 4);
      output->whiteScoreMean = scoreValueData[row * numScoreValueChannels];
      output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1];
      output->whiteLead = scoreValueData[row * numScoreValueChannels + 2];
      output->varTimeLeft = scoreValueData[row * numScoreValueChannels + 3];
    }
    else if(version >= 4) {
      int numScoreValueChannels = computeHandle->model.numScoreValueChannels;
      assert(numScoreValueChannels == 2);
      output->whiteScoreMean = scoreValueData[row * numScoreValueChannels];
      output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1];
      output->whiteLead = output->whiteScoreMean;
      output->varTimeLeft = 0;
    }
    else if(version >= 3) {
      int numScoreValueChannels = computeHandle->model.numScoreValueChannels;
      assert(numScoreValueChannels == 1);
      output->whiteScoreMean = scoreValueData[row * numScoreValueChannels];
      //Version 3 neural nets don't have any second moment output, implicitly already folding it in, so we just use the mean squared
      output->whiteScoreMeanSq = output->whiteScoreMean * output->whiteScoreMean;
      output->whiteLead = output->whiteScoreMean;
      output->varTimeLeft = 0;
    }
    else {
      ASSERT_UNREACHABLE;
    }
  }
}

void NeuralNet::printDevices() {
}

// FOR TESTING ---------------------------------------------------------------------------------------------------------
bool NeuralNet::testEvaluateConv(
  const ConvLayerDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  std::vector<float>& outputBuffer
) {
  if(!useNHWC || useFP16)
    return false;
  ConvLayer layer(*desc);
  TENSORMAP4 inTensor(
    (float*)inputBuffer.data(), desc->inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 outTensorBuf(desc->outChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP4 outTensor(outTensorBuf);

  layer.apply(&inTensor, &outTensor, false);

  outputBuffer.resize(outTensorBuf.size());
  memcpy(outputBuffer.data(), outTensorBuf.data(), sizeof(SCALAR) * outTensorBuf.size());
  return true;
}

// Mask should be in 'NHW' format (no "C" channel).
bool NeuralNet::testEvaluateBatchNorm(
  const BatchNormLayerDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  const std::vector<float>& maskBuffer,
  std::vector<float>& outputBuffer
) {
  if(!useNHWC || useFP16)
    return false;
  BatchNormLayer layer(*desc);
  TENSORMAP4 inTensor((float*)inputBuffer.data(), desc->numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP3 mask((float*)maskBuffer.data(), nnXLen, nnYLen, batchSize);
  TENSOR4 outTensorBuf(desc->numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP4 outTensor(outTensorBuf);

  layer.apply(false, &inTensor, &outTensor, &mask);

  outputBuffer.resize(outTensorBuf.size());
  memcpy(outputBuffer.data(), outTensorBuf.data(), sizeof(SCALAR) * outTensorBuf.size());
  return true;
}

// CR lpuchallafiore: test evaluate activation layer.
// CR lpuchallafiore: test evaluate matmul layer.
// CR lpuchallafiore: test evaluate matbias layer.

bool NeuralNet::testEvaluateResidualBlock(
  const ResidualBlockDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  const std::vector<float>& maskBuffer,
  std::vector<float>& outputBuffer
) {
  if(!useNHWC || useFP16)
    return false;
  ResidualBlock block(*desc);
  TENSORMAP4 inTensor((float*)inputBuffer.data(), desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP3 mask((float*)maskBuffer.data(), nnXLen, nnYLen, batchSize);

  TENSOR4 trunkBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 trunkScratchBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 midBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 midScratchBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);

  TENSORMAP4 trunk(trunkBuf);
  TENSORMAP4 trunkScratch(trunkScratchBuf);
  TENSORMAP4 mid(midBuf);
  TENSORMAP4 midScratch(midScratchBuf);

  trunk = inTensor;

  block.apply(
    &trunk,
    &trunkScratch,
    NULL,
    NULL,
    &mid,
    &midScratch,
    NULL,
    NULL,
    NULL,
    NULL,
    &mask,
    NULL
  );

  outputBuffer.resize(trunk.size());
  memcpy(outputBuffer.data(), trunk.data(), sizeof(SCALAR) * trunk.size());
  return true;
}

bool NeuralNet::testEvaluateGlobalPoolingResidualBlock(
  const GlobalPoolingResidualBlockDesc* desc,
  int batchSize,
  int nnXLen,
  int nnYLen,
  bool useFP16,
  bool useNHWC,
  const std::vector<float>& inputBuffer,
  const std::vector<float>& maskBuffer,
  std::vector<float>& outputBuffer) {
  if(!useNHWC || useFP16)
    return false;

  GlobalPoolingResidualBlock block(*desc);

  TENSORMAP4 inTensor((float*)inputBuffer.data(), desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSORMAP3 mask((float*)maskBuffer.data(), nnXLen, nnYLen, batchSize);

  TENSOR4 trunkBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 trunkScratchBuf(desc->preBN.numChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 regularOutBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 regularScratchBuf(desc->finalConv.inChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 gpoolOutBuf(desc->gpoolConv.outChannels, nnXLen, nnYLen, batchSize);
  TENSOR4 gpoolOut2Buf(desc->gpoolConv.outChannels, nnXLen, nnYLen, batchSize);
  TENSOR2 gpoolConcatBuf(desc->gpoolConv.outChannels*3, batchSize);
  TENSOR2 gpoolBiasBuf(desc->gpoolToBiasMul.outChannels, batchSize);

  TENSORMAP4 trunk(trunkBuf);
  TENSORMAP4 trunkScratch(trunkScratchBuf);
  TENSORMAP4 regularOut(regularOutBuf);
  TENSORMAP4 regularScratch(regularScratchBuf);
  TENSORMAP4 gpoolOut(gpoolOutBuf);
  TENSORMAP4 gpoolOut2(gpoolOut2Buf);
  TENSORMAP2 gpoolConcat(gpoolConcatBuf);
  TENSORMAP2 gpoolBias(gpoolBiasBuf);

  std::vector<float> maskSum(batchSize);
  computeMaskSum(&mask,maskSum.data());

  trunk = inTensor;

  block.apply(
    &trunk,
    &trunkScratch,
    &regularOut,
    &regularScratch,
    NULL,
    NULL,
    &gpoolOut,
    &gpoolOut2,
    &gpoolConcat,
    &gpoolBias,
    &mask,
    maskSum.data()
  );

  outputBuffer.resize(trunk.size());
  memcpy(outputBuffer.data(), trunk.data(), sizeof(SCALAR) * trunk.size());

  return true;
}

#endif  // USE_EIGEN_BACKEND

// CR lpuchallafiore: test evaluate Trunk
// CR lpuchallafiore: test evaluate Policy Head
// CR lpuchallafiore: test evaluate Value Head


#include <cradle/io/standard_image_io.hpp>
#include <cradle/imaging/iterator.hpp>
#include <cradle/imaging/channel.hpp>
#include <boost/filesystem/operations.hpp>
#include <cmath>
#include <cradle/io/file.hpp>

#define BOOST_TEST_MODULE file_io
#include <cradle/imaging/test.hpp>

using namespace cradle;

template<class SP>
void make_color_square(image<2,alia::rgb8,SP> const& img)
{
    alia::rgb<double> p(0, 0, 0),
        x_inc(0, 0, 255.0 / (img.size[0] - 1)),
        y_inc(255.0 / (img.size[1] - 1), 0, 0);
    for (unsigned y = 0; y < img.size[1]; ++y)
    {
        auto row_start = p;
        for (auto i = get_row_begin(img, y); i != get_row_end(img, y); ++i)
        {
            p.g = std::pow(p.r * p.b, 0.75);
            i->r = cradle::uint8_t(p.r);
            p.r += x_inc.r;
            i->g = cradle::uint8_t(p.g);
            p.g += x_inc.g;
            i->b = cradle::uint8_t(p.b);
            p.b += x_inc.b;
        }
        p.r = row_start.r + y_inc.r;
        p.g = row_start.g + y_inc.g;
        p.b = row_start.b + y_inc.b;
    }
}

file_path const image_path = "img.png";

BOOST_AUTO_TEST_CASE(color_image_test)
{
    image<2,alia::rgb8,unique> img;
    create_image(img, make_vector<unsigned>(24, 24));
    make_color_square(img);
    write_image_file(image_path, img);

    image<2,alia::rgb8,shared> img2;
    read_image_file(img2, image_path);
    CRADLE_CHECK_IMAGE(img, img2.pixels.view,
        img2.pixels.view + product(img2.size))

    remove(image_path);
}

BOOST_AUTO_TEST_CASE(try_gray_load_test)
{
    image<2,alia::rgb8,unique> img;
    create_image(img, make_vector<unsigned>(24, 24));
    make_color_square(img);
    write_image_file(image_path, img);

    image<2,alia::rgb<cradle::uint16_t>,shared> img2;
    BOOST_CHECK_THROW(read_image_file(img2, image_path), image_type_mismatch);

    remove(image_path);
}


// Make sure we include Python.h before any system header
// to avoid _POSIX_C_SOURCE redefinition
#ifdef WITH_PYTHON_LAYER
#include <boost/python.hpp>
#endif
#include <string>

#include "caffe/layer.hpp"
#include "caffe/layer_factory.hpp"
//#include "caffe/layers/clip_layer.hpp"
#include "caffe/layers/conv_layer.hpp"
#include "caffe/layers/deconv_layer.hpp"
#include "caffe/layers/lrn_layer.hpp"
#include "caffe/layers/pooling_layer.hpp"
#include "caffe/layers/relu_layer.hpp"
#include "caffe/layers/sigmoid_layer.hpp"
#include "caffe/layers/softmax_layer.hpp"
#include "caffe/layers/tanh_layer.hpp"
#include "caffe/proto/caffe.pb.h"

#ifdef USE_CUDNN
#include "caffe/layers/cudnn_conv_layer.hpp"
#include "caffe/layers/cudnn_deconv_layer.hpp"
#include "caffe/layers/cudnn_lcn_layer.hpp"
#include "caffe/layers/cudnn_lrn_layer.hpp"
#include "caffe/layers/cudnn_pooling_layer.hpp"
#include "caffe/layers/cudnn_relu_layer.hpp"
#include "caffe/layers/cudnn_sigmoid_layer.hpp"
#include "caffe/layers/cudnn_softmax_layer.hpp"
#include "caffe/layers/cudnn_tanh_layer.hpp"
#endif

#ifdef WITH_PYTHON_LAYER
#include "caffe/layers/python_layer.hpp"
#endif

namespace caffe {

// Get convolution layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetConvolutionLayer(
    const LayerParameter& param) {
  ConvolutionParameter conv_param = param.convolution_param();
  ConvolutionParameter_Engine engine = conv_param.engine();
#ifdef USE_CUDNN
  bool use_dilation = false;
  for (int i = 0; i < conv_param.dilation_size(); ++i) {
    if (conv_param.dilation(i) > 1) {
      use_dilation = true;
    }
  }
#endif
  if (engine == ConvolutionParameter_Engine_DEFAULT) {
    engine = ConvolutionParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    if (!use_dilation) {
      engine = ConvolutionParameter_Engine_CUDNN;
    }
#endif
  }
  if (engine == ConvolutionParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new ConvolutionLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == ConvolutionParameter_Engine_CUDNN) {
    if (use_dilation) {
      LOG(FATAL) << "CuDNN doesn't support the dilated convolution at Layer "
                 << param.name();
    }
    return shared_ptr<Layer<Dtype> >(new CuDNNConvolutionLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(Convolution, GetConvolutionLayer);

// Get deconvolution layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetDeconvolutionLayer(const LayerParameter& param) {
  ConvolutionParameter conv_param = param.convolution_param();
  ConvolutionParameter_Engine engine = conv_param.engine();
#ifdef USE_CUDNN
  bool use_dilation = false;
  for (int i = 0; i < conv_param.dilation_size(); ++i) {
    if (conv_param.dilation(i) > 1) {
      use_dilation = true;
    }
  }
#endif
  if (engine == ConvolutionParameter_Engine_DEFAULT) {
    engine = ConvolutionParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    if (!use_dilation) {
      engine = ConvolutionParameter_Engine_CUDNN;
    }
#endif
  }
  if (engine == ConvolutionParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new DeconvolutionLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == ConvolutionParameter_Engine_CUDNN) {
    if (use_dilation) {
      LOG(FATAL) << "CuDNN doesn't support the dilated deconvolution at Layer "
                 << param.name();
    }
    return shared_ptr<Layer<Dtype> >(new CuDNNDeconvolutionLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(Deconvolution, GetDeconvolutionLayer);

// Get pooling layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
  PoolingParameter_Engine engine = param.pooling_param().engine();
  if (engine == PoolingParameter_Engine_DEFAULT) {
    engine = PoolingParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = PoolingParameter_Engine_CUDNN;
#endif
  }
  if (engine == PoolingParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == PoolingParameter_Engine_CUDNN) {
    if (param.top_size() > 1) {
      LOG(INFO) << "cuDNN does not support multiple tops. "
                << "Using Caffe's own pooling layer.";
      return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
    }
    // CuDNN assumes layers are not being modified in place, thus
    // breaking our index tracking for updates in some cases in Caffe.
    // Until there is a workaround in Caffe (index management) or
    // cuDNN, use Caffe layer to max pooling, or don't use in place
    // layers after max pooling layers
    if (param.pooling_param().pool() == PoolingParameter_PoolMethod_MAX) {
        return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
    } else {
        return shared_ptr<Layer<Dtype> >(new CuDNNPoolingLayer<Dtype>(param));
    }
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(Pooling, GetPoolingLayer);

// Get LRN layer according to engine
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetLRNLayer(const LayerParameter& param) {
  LRNParameter_Engine engine = param.lrn_param().engine();

  if (engine == LRNParameter_Engine_DEFAULT) {
#ifdef USE_CUDNN
    engine = LRNParameter_Engine_CUDNN;
#else
    engine = LRNParameter_Engine_CAFFE;
#endif
  }

  if (engine == LRNParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new LRNLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == LRNParameter_Engine_CUDNN) {
    LRNParameter lrn_param = param.lrn_param();

    if (lrn_param.norm_region() ==LRNParameter_NormRegion_WITHIN_CHANNEL) {
      return shared_ptr<Layer<Dtype> >(new CuDNNLCNLayer<Dtype>(param));
    } else {
      // local size is too big to be handled through cuDNN
      if (param.lrn_param().local_size() > CUDNN_LRN_MAX_N) {
        return shared_ptr<Layer<Dtype> >(new LRNLayer<Dtype>(param));
      } else {
        return shared_ptr<Layer<Dtype> >(new CuDNNLRNLayer<Dtype>(param));
      }
    }
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(LRN, GetLRNLayer);

// Get relu layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetReLULayer(const LayerParameter& param) {
  ReLUParameter_Engine engine = param.relu_param().engine();
  if (engine == ReLUParameter_Engine_DEFAULT) {
    engine = ReLUParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = ReLUParameter_Engine_CUDNN;
#endif
  }
  if (engine == ReLUParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new ReLULayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == ReLUParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNReLULayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(ReLU, GetReLULayer);

// Get sigmoid layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetSigmoidLayer(const LayerParameter& param) {
  SigmoidParameter_Engine engine = param.sigmoid_param().engine();
  if (engine == SigmoidParameter_Engine_DEFAULT) {
    engine = SigmoidParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = SigmoidParameter_Engine_CUDNN;
#endif
  }
  if (engine == SigmoidParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new SigmoidLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == SigmoidParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNSigmoidLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(Sigmoid, GetSigmoidLayer);

// Get softmax layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetSoftmaxLayer(const LayerParameter& param) {
  SoftmaxParameter_Engine engine = param.softmax_param().engine();
  if (engine == SoftmaxParameter_Engine_DEFAULT) {
    engine = SoftmaxParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = SoftmaxParameter_Engine_CUDNN;
#endif
  }
  if (engine == SoftmaxParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new SoftmaxLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == SoftmaxParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNSoftmaxLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(Softmax, GetSoftmaxLayer);

// Get tanh layer according to engine.
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetTanHLayer(const LayerParameter& param) {
  TanHParameter_Engine engine = param.tanh_param().engine();
  if (engine == TanHParameter_Engine_DEFAULT) {
    engine = TanHParameter_Engine_CAFFE;
#ifdef USE_CUDNN
    engine = TanHParameter_Engine_CUDNN;
#endif
  }
  if (engine == TanHParameter_Engine_CAFFE) {
    return shared_ptr<Layer<Dtype> >(new TanHLayer<Dtype>(param));
#ifdef USE_CUDNN
  } else if (engine == TanHParameter_Engine_CUDNN) {
    return shared_ptr<Layer<Dtype> >(new CuDNNTanHLayer<Dtype>(param));
#endif
  } else {
    LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
    throw;  // Avoids missing return warning
  }
}

REGISTER_LAYER_CREATOR(TanH, GetTanHLayer);

#ifdef WITH_PYTHON_LAYER
template <typename Dtype>
shared_ptr<Layer<Dtype> > GetPythonLayer(const LayerParameter& param) {
  Py_Initialize();
  try {
    bp::object module = bp::import(param.python_param().module().c_str());
    bp::object layer = module.attr(param.python_param().layer().c_str())(param);
    return bp::extract<shared_ptr<PythonLayer<Dtype> > >(layer)();
  } catch (bp::error_already_set) {
    PyErr_Print();
    throw;
  }
}

REGISTER_LAYER_CREATOR(Python, GetPythonLayer);
#endif

// Layers that use their constructor as their default creator should be
// registered in their corresponding cpp files. Do not register them here.
}  // namespace caffe


#define BOOST_TEST_MODULE
#include <boost/test/unit_test.hpp>
#include <util/test_macros.hpp>
#include <vector>
#include <string>
#include <functional>

#include <random/random.hpp>

#include <unity/toolkits/recsys/models.hpp>
#include <unity/toolkits/util/data_generators.hpp>
#include <unity/toolkits/ml_data_2/ml_data.hpp>
#include <unity/toolkits/ml_data_2/sframe_index_mapping.hpp>
#include <sframe/testing_utils.hpp>
#include <util/testing_utils.hpp>

#include <cfenv>

using namespace turi;
using namespace turi::recsys;

template <typename Model>
void run_test_weird_segments(const std::map<std::string, flexible_type>& opts
                             = std::map<std::string, flexible_type>()) {

  size_t n_items = 200;
  size_t n_users = 200;
  size_t n_obs   = 500;

  sframe train_data, test_data, test_data_2;

  size_t num_segments = 11;

  random::seed(0);

  for(sframe* sf_ptr : {&train_data, &test_data, &test_data_2} ) {
    std::vector<std::vector<flexible_type> > data(n_obs);

    for(size_t i = 0; i < n_obs; ++i) {
      size_t user = random::fast_uniform<size_t>(0, n_users-1);
      size_t item = random::fast_uniform<size_t>(0, n_items-1);

      data[i] = {user, item, 1.0 / (1.0 + user + item)};
    }

    sf_ptr->open_for_write({"user", "item", "target"},
                           {flex_type_enum::INTEGER, flex_type_enum::INTEGER, flex_type_enum::FLOAT},
                           "",
                           num_segments);

    for(size_t sidx = 0; sidx < num_segments; ++sidx) {

      auto it_out = sf_ptr->get_output_iterator(sidx);

      size_t start_idx = (sidx * data.size()) / num_segments;
      size_t end_idx   = ((sidx+1) * data.size()) / num_segments;

      for(size_t i = start_idx; i < end_idx; ++i, ++it_out)
        *it_out = data[i];
    }

    sf_ptr->close();

    // Increment the number of segments so they are all different
    ++num_segments;
  }

  std::shared_ptr<recsys::recsys_model_base> model(new Model);

  std::map<std::string, flexible_type> opts2;
  opts2["item_id"] = "item";
  opts2["user_id"] = "user";
  opts2["target"] = "target";
  model->init_options(opts2);
  model->setup_and_train(train_data);

  v2::ml_data test_data_ml = model->create_ml_data(test_data);
  model->predict(test_data_ml);

  v2::ml_data test_data_ml_2 = model->create_ml_data(test_data_2);
  model->predict(test_data_ml_2);

  auto user_indexer = model->metadata->indexer(model->USER_COLUMN_INDEX);

  std::string user_column_name = model->metadata->column_name(model->USER_COLUMN_INDEX);

  sframe users_1({v2::map_to_indexed_sarray(
      user_indexer, test_data.select_column(user_column_name))},
    {"user"});

  sframe users_2({v2::map_to_indexed_sarray(
      user_indexer, test_data_2.select_column(user_column_name))},
    {"user"});

  model->recommend(users_1, 5);

  model->recommend(users_2, 5);
}

struct recsys_new_user_tests  {
 public:
  // ////////////////////////////////////////////////////////////////////////////////

  void test_new_users_factorization_model() {
    run_test_weird_segments<recsys::recsys_ranking_factorization_model>({{"max_iterations", 5}});
  }

  void test_new_users_ranking_factorization_model() {
    run_test_weird_segments<recsys::recsys_factorization_model>({{"max_iterations", 5}});
  }

  void test_new_users_itemcf_jaccard() {
    run_test_weird_segments<recsys::recsys_itemcf>(
      { {"similarity_type", "jaccard"} });
  }

  void test_new_users_itemcf_jaccard_topk() {
    run_test_weird_segments<recsys::recsys_itemcf>(
      { {"similarity_type", "jaccard"},
        {"only_top_k", 100} });
  }

  void test_new_users_itemcf_cosine() {
    run_test_weird_segments<recsys::recsys_itemcf>(
      { {"similarity_type", "cosine"} });
  }

  void test_new_users_itemcf_cosine_topk() {
    run_test_weird_segments<recsys::recsys_itemcf>(
      { {"similarity_type", "cosine"},
        {"only_top_k", 100} });
  }

  void test_new_users_pearson_cosine() {
    run_test_weird_segments<recsys::recsys_itemcf>(
      { {"similarity_type", "pearson"} });
  }

  void test_new_users_itemcf_pearson_topk() {
    run_test_weird_segments<recsys::recsys_itemcf>(
      { {"similarity_type", "pearson"},
        {"only_top_k", 100} });
  }

  void test_new_users_popularity() {
    run_test_weird_segments<recsys::recsys_popularity>();
  }

};

BOOST_FIXTURE_TEST_SUITE(_recsys_new_user_tests, recsys_new_user_tests)
BOOST_AUTO_TEST_CASE(test_new_users_factorization_model) {
  recsys_new_user_tests::test_new_users_factorization_model();
}
BOOST_AUTO_TEST_CASE(test_new_users_ranking_factorization_model) {
  recsys_new_user_tests::test_new_users_ranking_factorization_model();
}
BOOST_AUTO_TEST_CASE(test_new_users_itemcf_jaccard) {
  recsys_new_user_tests::test_new_users_itemcf_jaccard();
}
BOOST_AUTO_TEST_CASE(test_new_users_itemcf_jaccard_topk) {
  recsys_new_user_tests::test_new_users_itemcf_jaccard_topk();
}
BOOST_AUTO_TEST_CASE(test_new_users_itemcf_cosine) {
  recsys_new_user_tests::test_new_users_itemcf_cosine();
}
BOOST_AUTO_TEST_CASE(test_new_users_itemcf_cosine_topk) {
  recsys_new_user_tests::test_new_users_itemcf_cosine_topk();
}
BOOST_AUTO_TEST_CASE(test_new_users_pearson_cosine) {
  recsys_new_user_tests::test_new_users_pearson_cosine();
}
BOOST_AUTO_TEST_CASE(test_new_users_itemcf_pearson_topk) {
  recsys_new_user_tests::test_new_users_itemcf_pearson_topk();
}
BOOST_AUTO_TEST_CASE(test_new_users_popularity) {
  recsys_new_user_tests::test_new_users_popularity();
}
BOOST_AUTO_TEST_SUITE_END()


#include "3x3_solver.h"
#include <boost/range/algorithm_ext/push_back.hpp>

#include "engine.h"
#include "log.h"
#include "nmms.h"
#include "state.h"
#include "system.h"

static const Vec3 kPlaneNeighbor8[] = {
    Vec3(-1, -1, 0),
    Vec3(-1, 0, 0),
    Vec3(-1, 1, 0),
    Vec3(0, -1, 0),
    Vec3(0, 1, 0),
    Vec3(1, -1, 0),
    Vec3(1, 0, 0),
    Vec3(1, 1, 0),
};

void DoFill(const Vec3& current_position, const Vec3& full_direction, Matrix& current_matrix, Trace& trace) {
    ASSERT(current_matrix.is_in_matrix(current_position + full_direction));
    ASSERT(current_matrix(current_position + full_direction) == Void);
    current_matrix(current_position + full_direction) = Full;
    trace.push_back(CommandFill{ full_direction });
}

void DoVoid(const Vec3& current_position, const Vec3& void_direction, Matrix& current_matrix, Trace& trace) {
    ASSERT(current_matrix.is_in_matrix(current_position + void_direction));
    ASSERT(current_matrix(current_position + void_direction) == Full);
    current_matrix(current_position + void_direction) = Void;
    trace.push_back(CommandVoid{ void_direction });
}

void DoSMove(const Matrix& current_matrix, const Vec3& move_direction, Vec3& current_position, Trace& trace) {
    ASSERT(current_matrix.is_in_matrix(current_position + move_direction));
    current_position += move_direction;
    trace.push_back(CommandSMove{ move_direction });
}

bool is_inside_bound(const Vec3 position, const Vec3 upper, const Vec3 lower){
  return ( position.x >= lower.x && position.x < upper.x && position.y >= lower.y && position.y < upper.y && position.z >= lower.z && position.z < upper.z);
}

void dig_to_x(const Matrix& dst_matrix, Matrix& current_matrix, Trace& trace, Vec3& current_position, int x_direction, const int xstep) {

    Vec3 front_direction(x_direction, 0, 0);
    Vec3 back_direction(-x_direction, 0, 0);
    for (int i = 0; i < xstep; ++i) {
        // Void if the front voxel if full.
        if (current_matrix(current_position + front_direction) == Full) {
            DoVoid(current_position, front_direction, current_matrix, trace);
        }

        // Move to the front voxel.
        DoSMove(current_matrix, front_direction, current_position, trace);

        // Full if the back voxel as needed.
        ASSERT(current_matrix(current_position + back_direction) == Void);
        if (dst_matrix(current_position + back_direction) == Full) {
            DoFill(current_position, back_direction, current_matrix, trace);
        }
    }
}

void dig_to_y(const Matrix& dst_matrix, Matrix& current_matrix, Trace& trace, Vec3& current_position, const int ystep) {

    Vec3 front_direction(0, 1, 0);
    Vec3 back_direction(0, -1, 0);
    for (int i = 0; i < ystep; ++i) {
      // Void if the front voxel if full.
        if (current_matrix(current_position + front_direction) == Full) {
            DoVoid(current_position, front_direction, current_matrix, trace);
        }

        // Move to the front voxel.
        DoSMove(current_matrix, front_direction, current_position, trace);

        // Full if the back voxel as needed.
        ASSERT(current_matrix(current_position + back_direction) == Void);
        if (dst_matrix(current_position + back_direction) == Full) {
            DoFill(current_position, back_direction, current_matrix, trace);
        }
    }
}

void dig_to_z(const Matrix& dst_matrix, Matrix& current_matrix, Trace& trace, Vec3& current_position, int z_direction, const Vec3 upper, const Vec3 lower) {
  Vec3 front_direction(0, 0, z_direction);
  Vec3 back_direction(0, 0, -z_direction);
  for (;;) {
    // Full or void x-y 8-neighbor voxels.
    for (const auto& direction : kPlaneNeighbor8) {
      if (!dst_matrix.is_in_matrix(current_position + direction)) {
	continue;
      }

      if (current_matrix(current_position + direction) == Full && dst_matrix(current_position + direction) == Void && is_inside_bound(current_position + direction, upper, lower)) {
	DoVoid(current_position, direction, current_matrix, trace);
      }
      else if (current_matrix(current_position + direction) == Void && dst_matrix(current_position + direction) == Full && is_inside_bound(current_position + direction, upper, lower)) {
	DoFill(current_position, direction, current_matrix, trace);
      }
    }

    if(is_inside_bound(current_position + back_direction, upper, lower)){
      // Full the back voxel as needed.
      ASSERT(current_matrix(current_position + back_direction) == Void);
      if (dst_matrix(current_position + back_direction) == Full) {
      DoFill(current_position, back_direction, current_matrix, trace);
      }
    }

    if (current_position.z + z_direction < lower.z || upper.z - 1 < current_position.z + z_direction) {
      break;
    }

    // Void the front voxel as needed.
    if (current_matrix(current_position + front_direction) == Full) {
      DoVoid(current_position, front_direction, current_matrix, trace);
    }

    // Move to the front voxel.
    ASSERT(current_matrix(current_position + front_direction) == Void);
    DoSMove(current_matrix, front_direction, current_position, trace);
  }
}

const Vec3 unitX(1, 0, 0);
const Vec3 unitY(0, 1, 0);
const Vec3 unitZ(0, 0, 1);

void naive_move(Matrix &current_matrix, const Vec3& destination, Vec3& position, Trace& trace) {
  //cout<<position.x<<","<<position.y<<","<<position.z<<","<<destination.x<<","<<destination.y<<","<<destination.z<<endl;
  const Vec3 dz = (destination.z > position.z) ? unitZ : -unitZ;
  while (position.z != destination.z) {
    if(current_matrix(position+dz) == Full){
      DoVoid(position, dz, current_matrix, trace);
    }
    DoSMove(current_matrix, dz, position, trace);
  }

  const Vec3 dx = (destination.x > position.x) ? unitX : -unitX;
  while (position.x != destination.x) {
    if(current_matrix(position+dx) == Full){
      DoVoid(position, dx, current_matrix, trace);
    }
    DoSMove(current_matrix, dx, position, trace);
  }

  const Vec3 dy = (destination.y > position.y) ? unitY : -unitY;
  while (position.y != destination.y) {
    if(current_matrix(position+dy) == Full){
      DoVoid(position, dy, current_matrix, trace);
    }
    DoSMove(current_matrix, dy, position, trace);
  }
}

Trace single_three_by_three_solver(Matrix& current_matrix, const Matrix& src_matrix, const Matrix& dst_matrix, const Vec3& lower_bound, const Vec3& upper_bound, Vec3& position){
  Trace trace;
  // calculate bounding box

  Vec3 lower = upper_bound - Vec3(1, 1, 1);
  Vec3 upper = lower_bound;

  //cout<<"bef"<<upper.x<<","<<upper.y<<","<<upper.z<<","<<lower.x<<","<<lower.y<<","<<lower.z<<endl;
  for (int x = lower_bound.x; x < upper_bound.x; ++x) {
    for (int y = lower_bound.y; y < upper_bound.y; ++y) {
      for (int z = lower_bound.z; z < upper_bound.z; ++z) {
	if (src_matrix(Vec3(x, y, z)) || dst_matrix(Vec3(x, y, z))) {
	  lower.x = std::min(lower.x, x);
	  lower.y = std::min(lower.y, y);
	  lower.z = std::min(lower.z, z);
	  upper.x = std::max(upper.x, x + 1);
	  upper.y = std::max(upper.y, y + 1);
	  upper.z = std::max(upper.z, z + 1);
	}
      }
    }
  }

  lower.x = lower_bound.x;
  upper.x = upper_bound.x;

  lower.z = lower_bound.z;
  upper.z = upper_bound.z;

  if (lower.x >= upper.x) {
    return {};
  }

  lower.x = min(lower.x, position.x);
  //cout<<"aft"<<upper.x<<","<<upper.y<<","<<upper.z<<","<<lower.x<<","<<lower.y<<","<<lower.z<<endl;
  // Move to the start position.
  naive_move(current_matrix, lower, position, trace);

  // ZigZag
  int x_direction = 1;
  int z_direction = 1;
  //cout<<"dig start"<<endl;
  for (;;) {
    //cout<<"digz : "<<position.x<<","<<position.y<<","<<position.z<<","<<z_direction<<endl;
    dig_to_z(dst_matrix, current_matrix, trace, position, z_direction, upper, lower);
    z_direction = -z_direction;
    const int xstep = std::min(3, std::max(x_direction * (lower.x - position.x), x_direction * (upper.x - 1 - position.x)));
    if(xstep > 0){
      //cout<<"digx : "<<position.x<<","<<position.y<<","<<position.z<<","<<x_direction<<endl;
      dig_to_x(dst_matrix, current_matrix, trace, position, x_direction, xstep);
    }else {
      const int ystep = std::min(3,  upper.y - position.y);
      if (ystep > 0) {
	x_direction = -x_direction;
	//cout<<"digy : "<<position.x<<","<<position.y<<","<<position.z<<","<<z_direction<<endl;
	dig_to_y(dst_matrix, current_matrix, trace, position, ystep);
      }else {
	break;
      }
    }
  }

  //cout<<"digdone"<<endl;
  if (dst_matrix(position) == Full) {
    // Move +y by one voxel to fill the last voxel.
    //cout<<"fill final"<<endl;
    DoSMove(current_matrix, Vec3(0, 1, 0 ), position, trace);
    DoFill(position, Vec3(0, -1, 0 ), current_matrix, trace);
  }

  //cout<<"single done"<<endl;
  return trace;
}

Trace three_by_three_solver_parallel(ProblemType problem_type, const Matrix& src_matrix, const Matrix& dst_matrix) {
    printf("three_by_three_solver\n");
    const int R = src_matrix.R;
    System system(R);
    Trace trace;
    Matrix current_matrix = src_matrix;
    // spread nanobots
    ASSERT(system.bots.size() == 1);

    // not a good hach
    int Nsq=1;
    while( (Nsq+1) * (Nsq+1) * 2 <= R && (Nsq+1) * (Nsq+1)  < system.bots.size() + system.bots[0].seeds.size()){
      ++Nsq;
    }
    const int N = Nsq * Nsq;
    cout<<"dotnum = "<<N<<endl;

    std::vector<int> boundaries_init;

    for (int i = 0; i <= N; ++i) {
      boundaries_init.push_back(R * i / N);
    }

    std::vector<int> boundaries;
    for (int i = 0; i <= Nsq; ++i) {
      boundaries.push_back(R * i / Nsq);
      //cout<<boundaries[i]<<",";
    }
    //cout<<"boundaries"<<boundaries.size()<<endl;

    ASSERT(system.bots[0].pos == Vec3(0, 0, 0));

    std::vector<Vec3> positions(N, system.bots[0].pos);

    // goto highest to avoid collision
    for(int i=0;i<R-1;++i){
      trace.push_back(CommandSMove{unitY});
      positions[0] += unitY;
    }
    //cout<<"move"<<positions[0].y<<endl;

    // spread on x-axis
    int num_active_nanobots = 1;
    while (positions[N - 1].x < boundaries_init[N - 1]) {
      const int index = num_active_nanobots - 1;
      for (int i = 0; i < index; ++i) {
	trace.push_back(CommandWait{});
      }
      if (positions[index].x == boundaries_init[index]) {
	trace.push_back(CommandFission{unitX, N - ++num_active_nanobots});
	positions[index + 1] = positions[index] + unitX;
      } else {
	trace.push_back(CommandSMove{unitX});
	positions[index] += unitX;
      }
    }

    //cout<<"spread end"<<endl;

    // filp should be after spread
    trace.push_back(CommandFlip{});
    for (int i = 1; i < N; ++i) {
      trace.push_back(CommandWait{});
    }

    int highest = 0;
    std::vector<Trace> traces;
    for (int i = 0; i < N; ++i) {
      //cout<<i<<endl;
      traces.push_back(single_three_by_three_solver(current_matrix, src_matrix, dst_matrix,
						    Vec3(boundaries[i / Nsq], 0, boundaries[i % Nsq]),
						    Vec3(boundaries[(i / Nsq)+1], R, boundaries[(i % Nsq) + 1]),
						    positions[i]));
        highest = std::max(highest, positions[i].y);
    }

    //cout<<"trace get"<<endl;
    std::size_t max_trace_size = 0;
    for (int i = 0; i < N; ++i) {
      naive_move(current_matrix, Vec3(positions[i].x, highest, positions[i].z),
		 positions[i], traces[i]);
      max_trace_size = std::max(max_trace_size, traces[i].size());
    }

    //cout<<"move to top"<<endl;
    // merge
    for (std::size_t t = 0; t < max_trace_size; ++t) {
      for (int i = 0; i < N; ++i) {
	if (t < traces[i].size()) {
	  trace.push_back(traces[i][t]);
	} else {
	  trace.push_back(CommandWait{});
	}
      }
    }

    //cout<<"merged"<<endl;
    trace.push_back(CommandFlip{}); // low.
    for (int i = 1; i < N; ++i) {
      trace.push_back(CommandWait{});
    }

    // go home
    for (int i = N - 1; i > 0; --i) {
      Trace trace_to_join;
      naive_move(current_matrix, positions[i - 1], positions[i], trace_to_join);
      const auto* cmd = boost::get<CommandSMove>(&trace_to_join.back());
      ASSERT(cmd);
      positions[i] -= cmd->lld;
      trace_to_join.pop_back();
      for (std::size_t t = 0; t < trace_to_join.size(); ++t) {
	for (int j = 0; j < i; ++j) {
	  trace.push_back(CommandWait{});
	}
            trace.push_back(trace_to_join[t]);
      }
      for (int j = 0; j < i - 1; ++j) {
	trace.push_back(CommandWait{});
      }
      trace.push_back(CommandFusionP{positions[i] - positions[i - 1]});
      trace.push_back(CommandFusionS{positions[i - 1] - positions[i]});
    }
    naive_move(current_matrix, Vec3(0, 0, 0), positions[0], trace);

    //cout<<"backed"<<endl;
    // finalize at the origin pos.
    trace.push_back(CommandHalt{});

    return trace;
}

REGISTER_ENGINE(3x3_solver_parallel, three_by_three_solver_parallel);
// vim: set si et sw=4 ts=4:


#include <typeinfo>
#include <fstream>

#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>

#include "crosslink.hpp"

using namespace std;
namespace ublas = boost::numeric::ublas;

Crosslink::Crosslink(int id, ublas::vector<double> r, ublas::matrix<int> conTabSelf, ublas::matrix<int> conTabOther):
	r_(3,0), f_(3,0),f1_(3,0),f2_(3,0)
	{
		id_ = id;
		r_ = r;

		conTabSelf_ = conTabSelf;
		conTabOther_ = conTabOther;
	}

void Crosslink::updatePosition(ublas::vector<double> r) {
	r_ = r;
}

ublas::vector<double> Crosslink::incrPosition(ublas::vector<double> r) {
	r_ += r;
	return r_;
}

void Crosslink::updateForce(ublas::vector<double> f) {
	f_ = f;
}

void Crosslink::updateFilForce(ublas::vector<double> f, int filId){
	if(filId==1){ f1_=f;}
	if(filId==2){ f2_=f;}
}

void Crosslink::incrForce(ublas::vector<double> f, int filId) {
	f_ += f;
	if(filId==1){f1_+=f;}
	else if(filId==2){f2_+=f;}
	else {cout << "FilId error"<< "\n";}

}

ublas::vector<double> Crosslink::getForce() {
	return f_;
}

ublas::vector<double> Crosslink::getFilForce(int filId){
	if(filId==1){ return f1_; }
	else if (filId==2){ return f2_; }
	else {cout << "FilId error2"<< "\n";}
}

ublas::vector<double> Crosslink::getPosition() {
	return r_;
}

int Crosslink::getId() {
	return id_;
}

ublas::matrix<int> Crosslink::getConTabSelf(){
	return conTabSelf_;
}

ublas::matrix<int> Crosslink::getConTabOther(){
	return conTabOther_;
}


#include <epoxy/gl.h>
#include <GLFW/glfw3.h>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <iostream>
#include <iomanip>

#include <shadertoy.hpp>
#include <shadertoy/utils/log.hpp>

#include <signal.h>

#include <png.hpp>

#include "gn_perf_config.hpp"
#include "stat_acc.hpp"
#include "gn_glfw.hpp"
#include "hash.hpp"

static volatile int sigint_signaled = 0;

void sigint_handler(int _) {
    sigint_signaled = 1;
}

namespace po = boost::program_options;
namespace fs = boost::filesystem;
using shadertoy::utils::log;
using shadertoy::gl::gl_call;

template <typename inttype> inttype topx(float value, float k = 1.f)
{
	float scaled_value = (((value * 2.f - 1.f) * k) / 2.f + .5f) * std::numeric_limits<inttype>::max();

	// Clamp
	if (scaled_value > std::numeric_limits<inttype>::max())
		scaled_value = static_cast<float>(std::numeric_limits<inttype>::max());
	else if (scaled_value < 0.0f)
		scaled_value = 0.0f;

	return static_cast<inttype>(scaled_value);
}

void write_output(const std::string &output_param, const stat_acc &time_ms, const std::shared_ptr<shadertoy::members::basic_member> &last_result, int width, int height, const std::string &include_stat, bool raw_output, const std::string &identifier)
{
    auto output_basename = (output_param == "auto") ? identifier : output_param;

    log::shadertoy()->info("Writing output at {}", output_basename);

    // Fetch from OpenGL
    std::vector<float> image_data(width * height * 4);

    // Get the texture
    auto texture = last_result->output();
    assert(texture);
    texture->get_image(0, GL_RGBA, GL_FLOAT, sizeof(float) * image_data.size(), image_data.data());

    // Write image
    png::image<png::rgba_pixel_16> image(width, height);
    for (int x = 0; x < width; ++x) {
        for (int y = 0; y < height; ++y) {
            float *pi = &image_data[4 * (y * width + x)];
            image.set_pixel(x, y, png::rgba_pixel_16(
                        topx<uint16_t>(pi[0]),
                        topx<uint16_t>(pi[1]),
                        topx<uint16_t>(pi[2]),
                        std::numeric_limits<uint16_t>::max()));
        }
    }

    image.write(output_basename + ".png");

    // Write details
    std::ofstream ofs((output_basename + ".txt").c_str());
    assert(!ofs.fail());

    // Print state identifier
    if (!raw_output)
    {
        ofs << "# " << identifier << std::endl;
    }

    bool output_header = false;
    if (include_stat.find("t_ms") != std::string::npos)
        ofs << time_ms.summary("", raw_output, output_header, "t_ms").c_str() << std::endl;
    if (include_stat.find("fps") != std::string::npos)
        ofs << time_ms.summary("", raw_output, output_header, "fps", [](auto x) { return 1.0e3 / x; }).c_str() << std::endl;
    if (include_stat.find("mpxps") != std::string::npos)
        ofs << time_ms.summary("", raw_output, output_header, "mpxps", [width, height](auto x) { return 1.0e-3 * width * height / x; }).c_str() << std::endl;
}

int main(int argc, char *argv[])
{
    int width, height, size;
    long long samples, warmup_samples;
    bool st_silent, all_silent, raw_output, sync_anyways, test_mode;
    std::string include_stat, output, lut_path;
    std::vector<std::string> defines;

    po::options_description gn_desc("Noise options");
    gn_desc.add_options()
        ("width", po::value(&width)->default_value(640), "Width of the rendering")
        ("height", po::value(&height)->default_value(480), "Height of the rendering")
        ("size,s", po::value(&size)->default_value(-1), "Size (overrides width and height) of the rendering")
        ("warmup,W", po::value(&warmup_samples)->default_value(-1), "Number of samples to warm-up the measurements")
        ("samples,n", po::value(&samples)->default_value(0), "Number of samples to collect for statistics")
        ("output,o", po::value(&output)->default_value(""), "Output path for the control frame")
        ("lut,l", po::value(&lut_path)->default_value(""), "LUT texture path")
        ("define,D", po::value(&defines)->multitoken()->composing(), "Preprocessor definitions for the shader\n"
         "The following values are supported: \n"
         "\t * SPLATS=n: number of splats per cell\n"
         "\t * F0=freq: frequency of the kernel\n"
         "\t - W0=angle: angle of the anisotropic kernel (defaults to pi/4)\n"
         "\t - TILE_SIZE=size: kernel diameter, in pixels (defaults to width/3)\n"
         "\t - RANDOM_SEED=r: random seed (defaults to 0)\n"
         "\t - DISP_SIZE=s: number of cells to look for contributing splats (defaults to 1)\n"
         "\t - KTRUNC: make kernel boundary C0\n"
         "\t - KSIN: use sin instead of cos for kernel\n"
         "\t - KSHOW: show kernel footprint as a disk\n"
         "\t - KHALF: use a half-size kernel, evaluate only 4 cells\n"
         "\t - KKAISER_BESSEL: use the Kaiser-Bessel window for the kernel\n"
         "\t - RANDOM_PHASE: use random phase kernel\n"
         "\t - WEIGHTS=weight: type of the random weights to use:\n"
         "\t   - WEIGHTS_UNIFORM: uniform [-1, 1] weights (default)\n"
         "\t   - WEIGHTS_BERNOULLI: Bernoulli {-1, 1} weights\n"
         "\t   - WEIGHTS_NONE: no weights\n"
         "\t - POINTS=type: type of the point distribution:\n"
         "\t   - POINTS_WHITE: white (Poisson) points (default)\n"
         "\t   - POINTS_STRATIFIED: stratified points\n"
         "\t   - POINTS_JITTERED: jittered rectangular grid points\n"
         "\t   - POINTS_HEX_JITTERED: jittered triangular grid points\n"
         "\t   - POINTS_GRID: regular rectangular grid points\n"
         "\t   - POINTS_HEX_GRID: regular triangular grid points\n"
         "\t - PRNG=type: type of the PRNG to use:\n"
         "\t   - PRNG_LCG: linear congruential generator (default)\n"
         "\t   - PRNG_XOROSHIRO: SIMD xoroshiro64** generator\n"
         "\t   - PRNG_HASH: integer hash generator (2^8 period)\n"
         "\t   - PRNG_XORSHIFT: xorshift tuple generator\n"
         "\t   - PRNG_NONE: constant generator for debugging\n"
         "\t - PRESET_BOOT: enable the boot preset, should be used with --lut\n");

    po::options_description desc("gn-perf " GN_PERF_VERSION " (" GN_PERF_BASE_DIR ")");
    desc.add_options()
        ("config,C", po::value<std::string>(), "Configuration file for the noise")
        ("quiet,q", po::bool_switch(&st_silent)->default_value(false), "Silence libshadertoy debug messages")
        ("very-quiet,Q", po::bool_switch(&all_silent)->default_value(false), "Silence everything but the final output")
        ("include-stat,I", po::value(&include_stat)->default_value("t_ms,fps,mpxps"), "Stats to include in the output")
        ("raw,r", po::bool_switch(&raw_output)->default_value(false), "Raw value output (no header nor size). Only applies to final output")
        ("sync,S", po::bool_switch(&sync_anyways)->default_value(false), "Force vsync even if measuring performance")
        ("test,T", po::bool_switch(&test_mode)->default_value(false), "TAP self-test mode")
        ("help,h", "Show this help message");

    desc.add(gn_desc);

    po::positional_options_description po_desc;
    po_desc.add("config", 1);

    po::variables_map vm;

    try
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(po_desc).run(), vm);

        if (vm.count("config") > 0)
        {
            auto config_name(vm["config"].as<std::string>());
            std::ifstream ifs(config_name.c_str());

            if (ifs.fail())
            {
                throw po::error("Failed to open config file");
            }

            po::store(po::parse_config_file(ifs, gn_desc), vm);
        }

        po::notify(vm);
    }
    catch (const po::error &ex)
    {
        std::cerr << ex.what() << std::endl;
        std::cerr << "See --help option for usage" << std::endl;
        return 1;
    }

    if (vm.count("help"))
    {
        std::cout << desc << std::endl;
        return 0;
    }

    if (size > 0)
    {
        width = height = size;
    }

    if (warmup_samples < 0 && samples != 0)
    {
        warmup_samples = 64;
    }
    else
    {
        warmup_samples = 0;
    }

    if (test_mode)
    {
        all_silent = true;
    }

    if (all_silent)
    {
        freopen("/dev/null", "w", stderr);
    }

    log::shadertoy()->set_level(st_silent ? spdlog::level::info : spdlog::level::debug);

    log::shadertoy()->info("gn-perf {}", GN_PERF_VERSION);
    log::shadertoy()->info("Base dir {}", GN_PERF_BASE_DIR);

    if (samples == 0)
        log::shadertoy()->info("No sample count specified, running at vsync for debug");
    else if (samples < 0)
        log::shadertoy()->info("About to collect enough samples so stddev% < {}e-2%", -samples);
    else
        log::shadertoy()->info("About to collect {} samples", samples);

#if HAS_NVML
    // We are using the first GPU anyways
    nvmlDevice_t device;
    bool nvml_enabled = false;

    if (nvmlInit() == NVML_SUCCESS) {
        nvml_enabled = true;
        log::shadertoy()->debug("Initialized NVML");
    } else {
        log::shadertoy()->warn("Failed to initialize NVML");
    }

    if (nvml_enabled && nvmlDeviceGetHandleByIndex(0, &device) == NVML_SUCCESS) {
        nvmlPstates_t pstate = NVML_PSTATE_UNKNOWN;
        nvmlDeviceGetPerformanceState(device, &pstate);

        char buf[NVML_DEVICE_NAME_BUFFER_SIZE];
        if (nvmlDeviceGetName(device, buf, sizeof(buf)) == NVML_SUCCESS) {
            log::shadertoy()->info("Found NVML device: {}, initial p-state: {}", buf, pstate);
        } else {
            log::shadertoy()->info("Found NVML device, initial p-state: {}", pstate);
        }
    } else {
        nvml_enabled = false;
        log::shadertoy()->warn("Device not found");
    }
#endif /* HAS_NVML */

    bool visible = samples == 0 || sync_anyways ? 1 : 0;
    return glfw_run(width, height, visible, [&](auto *window)
    {
        // Create the context and swap chain
        gn_perf_ctx ctx(width, height, defines, visible, lut_path);
        auto &context(ctx.context);
        auto &chain(ctx.chain);

        // Now render for 5s
        int frameCount = 0;
        double t = 0.;

        // Set the resize callback
        glfwSetWindowUserPointer(window, &ctx);

        stat_acc time_ms;

        fprintf(stderr, "%8s\t%10s\t%9s\t%13s\t%4s\t%4s\t%9s\n", "frame", "time_ms", "fps", "mpx_s", "wh_px", "ch_px", "stddevp");

        signal(SIGINT, sigint_handler);

        while (!glfwWindowShouldClose(window) && !sigint_signaled)
        {
            // Poll events
            glfwPollEvents();

            // Update uniforms
            context.state().get<shadertoy::iTime>() = t;
            context.state().get<shadertoy::iFrame>() = uhash(frameCount);

            // Set viewport
            // This is not necessary when the last pass is rendering to a
            // texture and it is followed by a screen_member, which calls
            // glViewport. In this example, we render directly to the
            // default framebuffer, so we need to set the viewport
            // ourselves.
            gl_call(glViewport, 0, 0, ctx.render_size.width, ctx.render_size.height);

            // Render the swap chain
            context.render(chain);

            // Buffer swapping
            glfwSwapBuffers(window);

            bool pstate_ok = true;
#if HAS_NVML
            // If we have a working NVML, only start measuring at P2 or higher
            nvmlPstates_t pstate = NVML_PSTATE_UNKNOWN;
            if (nvml_enabled && nvmlDeviceGetPerformanceState(device, &pstate) == NVML_SUCCESS) {
                pstate_ok = pstate <= NVML_PSTATE_2;
            }
#endif

            if (warmup_samples <= 0 && pstate_ok)
            {
                // Get the render time for the frame
                auto elapsed_time = ctx.image_buffer->elapsed_time();
                auto pixel_count = static_cast<double>(ctx.render_size.width * ctx.render_size.height);

                // 0 should not be measured by the driver
                if (elapsed_time != 0)
                    time_ms.sample(elapsed_time / 1e6);

                auto stddevp = time_ms.stddevp();
                if ((samples > 0 && time_ms.sample_count() == samples) ||
                    (samples < 0 && time_ms.sample_count() >= 16 && (stddevp * 1e4) < -samples))
                    glfwSetWindowShouldClose(window, 1);

                fprintf(stderr, "%8d\t%10lf\t%8.2lf\t%12.2lf\t%4d\t%4d\t%8.2lf\n",
                        frameCount,
                        elapsed_time / 1e6,
                        1.0e9 / elapsed_time,
                        1.0e3 * pixel_count / elapsed_time,
                        ctx.render_size.width,
                        ctx.render_size.height,
                        stddevp * 1e2);
            }
            else
            {
                warmup_samples--;
            }

            // Update time and framecount
            t = glfwGetTime();
            frameCount++;
        }

        fprintf(stderr, "\n");

        // Write output data
        if (!output.empty())
        {
            write_output(output, time_ms, ctx.chain.members().front(), ctx.render_size.width, ctx.render_size.height, include_stat, raw_output, ctx.identifier);
        }

        // TODO: actually test something
        const char *test_prefix = test_mode ? "# " : "";
        if (test_mode)
        {
            printf("1..1\nok 1\n");
        }

        // Print state identifier
        if (!raw_output)
        {
            printf("%s# %s\n", test_prefix, ctx.identifier.c_str());
        }

        bool output_header = false;
        if (include_stat.find("t_ms") != std::string::npos)
            printf("%s\n", time_ms.summary(test_prefix, raw_output, output_header, "t_ms").c_str());
        if (include_stat.find("fps") != std::string::npos)
            printf("%s\n", time_ms.summary(test_prefix, raw_output, output_header, "fps", [](auto x) { return 1.0e3 / x; }).c_str());
        if (include_stat.find("mpxps") != std::string::npos)
            printf("%s\n", time_ms.summary(test_prefix, raw_output, output_header, "mpxps", [&ctx](auto x) { return 1.0e-3 * ctx.render_size.width * ctx.render_size.height / x; }).c_str());

#if HAS_NVML
        nvmlShutdown();
#endif
    });
}

// vim: cino=


#include <boost/assert.hpp>
#include <iostream>
#include <fstream>

#include "nlohmann/json.hpp"

#include "MenuTree.hpp"
#include "MenuNode.hpp"
#include "utility/wrapper/sfVector2.hpp"
#include "utility/RC1DConverter.hpp"

namespace nemo
{

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

MenuTree::MenuTree(
	const XYPair&         pos,
	const XYPair&         dim,
	const RCPair&         row_by_col,
	const XYPair&         padding,
	const XYPair&         spacing,
	const TextBoxColors   box_colors,
	const TextBoxColors   entry_colors,
	const TextBoxColors   hover_colors,
	const FontProperties& font)

	: MenuNode(pos, dim, padding, box_colors, font)
	, spacing_(spacing)
	, rows_(row_by_col.r_)
	, cols_(row_by_col.c_)
	, entry_colors_(entry_colors)
	, cursor_(hover_colors)
{
	BOOST_ASSERT(spacing >= XYPair(XValue(0.f), YValue(0.f)));
	BOOST_ASSERT(row_by_col.r_ > 0);
	BOOST_ASSERT(row_by_col.c_ > 0);
}
///////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

std::shared_ptr<MenuNode>
MenuTree::add(const std::shared_ptr<MenuNode> child)
{
	const auto row_by_col = XYPair(XValue(int(cols_)), YValue(int(rows_)));
	const auto spaced_dim = getInnerSize() / row_by_col;
	const auto dim = spaced_dim - spacing_ * 2.f;
	child->setSize(dim);

	const auto rc_i = RC1DConverter(cols_).toRowColumn(children_.size());
	const auto rel_pos = spaced_dim
		* XYPair(XValue(int(rc_i.c_)), YValue(int(rc_i.r_)))
		+ spacing_;
	const auto abs_pos = getPosition() + rel_pos;
	child->setPosition(abs_pos);

	child->setColors(entry_colors_);
	children_.push_back(child);
	return shared_from_this();
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

void
MenuTree::drawIt(sf::RenderWindow& window)
const
{
	drawTextBox(window);

	for (auto c : children_) {
		c->drawIt(window);
	}
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

std::shared_ptr<MenuNode>
MenuTree::setCaption(const std::string& caption)
{
	constexpr auto vt_center = false;
	return makeCaption(caption, vt_center);
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

void
MenuTree::cursorUp()
noexcept
{
	moveCursor(Direction::Up);
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

void
MenuTree::cursorDown()
noexcept
{
	moveCursor(Direction::Down);
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

void
MenuTree::cursorLeft()
noexcept
{
	moveCursor(Direction::Left);
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

void
MenuTree::cursorRight()
noexcept
{
	moveCursor(Direction::Right);
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

void
MenuTree::moveCursor(const Direction dir)
noexcept
{
	if (children_.empty()) {
		// No menu options => no cursor => no movement.
		return;
	}

	auto& cursor_idx = cursor_.idx_;
	const auto last = static_cast<decltype(cursor_.idx_)>(children_.size()) - 1;

	switch (dir) {
		// Up/down changes the row.
		case Direction::Up: {
			const auto new_idx = cursor_idx - int(cols_);
			// The cursor should be able to wrap around the ends of the menu.
			cursor_idx = new_idx >= 0
				? new_idx
				: cursor_idx != 0
					? 0
					: last;
			break;
		}

		case Direction::Down: {
			const auto new_idx = cursor_idx + int(cols_);
			cursor_idx = new_idx <= last
				? new_idx
				: cursor_idx != last
					? last
					: 0;
			break;
		}

		// Left/right changes the column.
		case Direction::Right: {
			cursor_idx = cursor_idx < last ? cursor_idx + 1 : 0;
			break;
		}

		case Direction::Left: {
			cursor_idx = cursor_idx > 0 ? cursor_idx - 1 : last;
			break;
		}
	}
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

}