#include "YODA/BinnedAxis.h"
#include "YODA/Binning.h"
#include "YODA/Utils/MetaUtils.h"
#include <iostream>
#include <string>
#include <algorithm>

using namespace YODA;

template<typename T>
using vec = std::vector<T>;
using StrVec = vec<std::string>;

using DStrAxis = Axis<std::string>;
using DIntAxis = Axis<int>;
using CDblAxis = Axis<double>;

auto areEqualVecs = [](const std::vector<size_t>& lhs, const std::vector<size_t>& rhs) {
  return lhs.size() == rhs.size() && std::equal(lhs.begin(), lhs.end(), rhs.begin());
};


auto bConstr() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  DStrAxis ax1(v1);
  CDblAxis ax2(v2);
  DIntAxis ax3(v3);

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);
  Binning<DStrAxis, CDblAxis, DIntAxis> binning2(
    {"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});
  Binning<DStrAxis, CDblAxis, DIntAxis> binning3(ax1, ax2, ax3);

  using IndexArr = std::array<size_t, 3>;
  IndexArr axesSizes{3, 5, 6};

  return CHECK_TEST_RES(binning1.isCompatible(binning2) &&
                        binning2.isCompatible(binning1) &&
                        binning2.isCompatible(binning3) &&
                        binning1._getAxesSizes() == axesSizes);
}

auto blocalToGlobalIndex() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  // The mapping goes like global = xLocal + yLocal * XLen + zLocal * YLen * XLen
  return CHECK_TEST_RES((binning1.localToGlobalIndex({0,0,0}) == 0)  && // 0 + 0*3 + 0*5*3
                        (binning1.localToGlobalIndex({1,4,4}) == 73) && // 1 + 4*3 + 4*5*3
                        (binning1.localToGlobalIndex({1,2,3}) == 52));  // 1 + 2*3 + 3*5*3
}

auto bglobalToLocalIndices() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  return CHECK_TEST_RES((binning1.globalToLocalIndices(0) == std::array<size_t, 3>{0,0,0})  &&
                        (binning1.globalToLocalIndices(73) == std::array<size_t, 3>{1,4,4}) &&
                        (binning1.globalToLocalIndices(52) == std::array<size_t, 3>{1,2,3}));
}

auto blocalIndicesAt() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  using IndexArr = std::array<size_t, 3>;

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  return CHECK_TEST_RES((binning1.localIndicesAt({"test2", 2.0001, 5}) == IndexArr{2,2,5})  &&
                        (binning1.localIndicesAt({"test1", -5, 1}) == IndexArr{1,0,1}) &&
                        (binning1.localIndicesAt({"test2", 5, 5}) == IndexArr{2,4,5}));
}

auto bglobalIndex() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  // The mapping goes like global = xLocal + yLocal * XLen + zLocal * YLen * XLen
  return CHECK_TEST_RES((binning1.globalIndexAt({"test2", 2.0001, 5}) == 83)  && // 2 + 2*3 + 5*5*3
                        (binning1.globalIndexAt({"test1",  -5, 1})    == 16)  && // 1 + 0*3 + 1*5*3
                        (binning1.globalIndexAt({"test1", -25, 1})    == 16)  && // 1 + 0*3 + 1*5*3
                        (binning1.globalIndexAt({"test2",  5,  5})    == 89)  && // 2 + 4*3 + 5*5*3
                        (binning1.globalIndexAt({"test2",  54, 5})    == 89));   // 2 + 4*3 + 5*5*3
}

auto bisCompatible() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);
  Binning<DStrAxis, CDblAxis, DIntAxis> binning2(v1, v2, v3);

  return CHECK_TEST_RES(binning1.isCompatible(binning2) &&
                        binning2.isCompatible(binning1));
}

auto bgetAxis() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  DStrAxis ax1(v1);
  CDblAxis ax2(v2);
  DIntAxis ax3(v3);

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(ax1, ax2, ax3);

  return CHECK_TEST_RES(binning1.axis<0>().hasSameEdges(ax1) &&
                        binning1.axis<1>().hasSameEdges(ax2) &&
                        binning1.axis<2>().hasSameEdges(ax3));
}

auto bgetAxesSizes() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  using IndexArr = std::array<size_t, 3>;

  return CHECK_TEST_RES((binning1._getAxesSizes() == IndexArr{3, 5, 6}));
}

auto bdim() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  return CHECK_TEST_RES((binning1.dim() == 3));
}

auto bgetNBins() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning1(v1, v2, v3);

  return CHECK_TEST_RES((binning1.numBins() == 90)); // 3*5*6
}

auto bcalcSliceSize() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning(v1, v2, v3);

  return CHECK_TEST_RES(
    binning.calcSliceSize(0) == 30 && // (3+2)*(5+1)
    binning.calcSliceSize(1) == 18 && // (2+1)*(5+1)
    binning.calcSliceSize(2) == 15);  // (2+1)*(3+2)
}

auto bsliceBinning() {
  vec<double> v1 = {1., 2.};
  vec<double> v2 = {1., 2.};
  vec<double> v3 = {1., 2.};

  Binning<CDblAxis, CDblAxis, CDblAxis> binning(v1, v2, v3);

  /// @note Bin index calculates as follows:
  /// x + y*width + z*width*height.
  /// Slice will be performed at first axis, at bin 0, thus
  /// bin indices will be calculated as 0 + y*width + z*width*height.
  vec<size_t> expectedIndices1 = {0,  3,   6,  /// + 0*width*height
                                  9,  12, 15,  /// + 1*width*height
                                  18, 21, 24}; /// + 2*width*height

  vec<size_t> sliceIndices = binning.sliceIndices(0, 0);

  /// x + y*width + 2*width*height
  vec<size_t> expectedIndices2 = {18, 19, 20,  /// + 0*width
                                  21, 22, 23,  /// + 1*width
                                  24, 25, 26}; /// + 2*width

  vec<size_t> expectedIndices3;

  auto appendVec = [](vec<size_t>& lhs, vec<size_t> rhs) {
    lhs.insert(lhs.end(), rhs.begin(), rhs.end());
  };

  appendVec(expectedIndices3, expectedIndices1);
  appendVec(expectedIndices3, expectedIndices2);

  vec<size_t> slicedIndices = binning.sliceIndices({{0, {0}}, {2, {2}}});

  return CHECK_TEST_RES(areEqualVecs(sliceIndices,  expectedIndices1) &&
                        areEqualVecs(slicedIndices, expectedIndices3));
}

auto boverflowBinIndices() {
  StrVec v1 = {"test1", "test2"};
  vec<double> v2 = {1., 2., 3., 4.};
  vec<int> v3 = {1, 2, 3, 4, 5};

  Binning<DStrAxis, CDblAxis, DIntAxis> binning(v1, v2, v3);

  vec<size_t> expectedIndices = {
     0,  1,  2,     // x + 0*3 + 0*5*3
     3,  4,  5,     // x + 1*3 + 0*5*3
     6,  7,  8,     // x + 2*3 + 0*5*3
     9, 10, 11,     // x + 3*3 + 0*5*3
    12, 13, 14,     // x + 4*3 + 0*5*3
    15, 16, 17,     // x + 0*3 + 1*5*3
    18, 21, 24,     // 0 + y*3 + 1*5*3
    27, 28, 29,     // x + 4*3 + 1*5*3
    30, 31, 32,     // x + 0*3 + 2*5*3
    33, 36, 39,     // 0 + y*3 + 2*5*3
    42, 43, 44,     // x + 4*3 + 2*5*3
    45, 46, 47,     // x + 0*3 + 3*5*3
    48, 51, 54,     // 0 + y*3 + 3*5*3
    57, 58, 59,     // x + 4*3 + 3*5*3
    60, 61, 62,     // x + 0*3 + 4*5*3
    63, 66, 69,     // 0 + y*3 + 4*5*3
    72, 73, 74,     // x + 4*3 + 4*5*3
    75, 76, 77,     // x + 0*3 + 5*5*3
    78, 81, 84,     // 0 + y*3 + 5*5*3
    87, 88, 89,     // x + 4*3 + 5*5*3
  };

  return CHECK_TEST_RES(areEqualVecs(binning.calcOverflowBinsIndices(), expectedIndices));
}

auto testBinning() {
  return CHECK_TEST_RES(
    (bConstr()                == EXIT_SUCCESS) &&
    (blocalToGlobalIndex()    == EXIT_SUCCESS) &&
    (bglobalToLocalIndices()  == EXIT_SUCCESS) &&
    (blocalIndicesAt()        == EXIT_SUCCESS) &&
    (bglobalIndex()           == EXIT_SUCCESS) &&
    (bisCompatible()          == EXIT_SUCCESS) &&
    (bgetAxis()               == EXIT_SUCCESS) &&
    (bgetAxesSizes()          == EXIT_SUCCESS) &&
    (bdim()                   == EXIT_SUCCESS) &&
    (bgetNBins()              == EXIT_SUCCESS) &&
    (bcalcSliceSize()         == EXIT_SUCCESS) &&
    (bsliceBinning()          == EXIT_SUCCESS) &&
    (boverflowBinIndices()    == EXIT_SUCCESS));
}

int main() {
  int rtn = EXIT_SUCCESS;

  rtn = testBinning();

  return rtn;
}
