#include "YODA/BinnedAxis.h"
#include "YODA/FillableStorage.h"
#include "YODA/Binning.h"
#include "YODA/Utils/MetaUtils.h"
#include "YODA/Dbn.h"
#include <algorithm>
#include <functional>
#include <iostream>
#include <vector>

using namespace YODA;

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

using BinningT = Binning<DStrAxis, CDblAxis, DIntAxis>;

template <typename T>
using StorageT = BinnedStorage<T, std::string, double, int>;

template <size_t N, typename T>
using FillableT = FillableStorage<N, T, std::string, double, int>;

auto constructors() {
  BinningT binning({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});

  DStrAxis ax1({"test1", "test2"});
  CDblAxis ax2({1., 2., 3., 4.});
  DIntAxis ax3({1, 2, 3, 4, 5});

  auto binStorage1 = StorageT<double>(binning);
  auto binStorage2 = StorageT<double>({"test1", "test2"},
                                      {1., 2., 3., 4.},
                                      {1, 2, 3, 4, 5});
  auto binStorage3 = StorageT<double>(ax1, ax2, ax3);


  return CHECK_TEST_RES(true);
}

struct TestClass{
  TestClass(size_t k = 0) : counter(k) {}
  void func() {
    counter++;
    return;
  }
  size_t counter = 0;
};

auto adapterConstructors() {
  BinningT binning({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});

  FillableT<3, double>::FillAdapterT adapter1 = [](auto& storedNumber, ...) {
     auto test = storedNumber + 1;
     (void)test; /// @note To omit warning (unused variable)
  };
  auto adapter2 = [](auto& customObject, auto&& /*coords*/, auto /*weight*/, auto /*fraction*/) {
    customObject.func();
  };

  auto binStorage1 = FillableT<3, double>(binning); /// @note default adapter is used
  auto binStorage2 = FillableT<3, double>(binning, adapter1);
  auto binStorage3 = FillableT<3, TestClass>(binning, adapter2);

  binStorage3.fill({"test1", 1.5, 1}, std::make_index_sequence<3>());

  FillableT<3, TestClass>::BinType& bin = binStorage3.binAt({"test1", 1.5, 1});

  return CHECK_TEST_RES(bin.counter == 1);
}

auto set() {
  BinningT binning({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});

  auto binStorage = FillableT<3, TestClass>(binning, [](...){});

  binStorage.set({"test1", 1.5, 1}, TestClass(15));

  auto& bin = binStorage.binAt({"test1", 1.5, 1});

  return CHECK_TEST_RES(bin.counter == 15);
}

auto binsVectorWrapper() {
  auto binning = std::make_shared<BinningT>(
    BinningT({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5}));
  using BinT = Bin<3, int, BinningT>;
  std::vector<BinT> bins;

  bins.reserve(100);
  for(size_t i = 0; i < 100; i++){
    bins.push_back(BinT(i, *binning));
  }

  /// @note Must be sorted.
  std::vector<size_t> hiddenBins = {0, 9, 15, 87};

  auto vecWrapper = BinsVecWrapper<std::vector<BinT>>(bins, hiddenBins);

  for (auto& bin : vecWrapper) {
    for (auto& i : hiddenBins) {
      if (bin.index() == i)  return CHECK_TEST_RES(false);
    }
  }

  return CHECK_TEST_RES(true);
}

auto binSpaceCharacteristics() {
  BinningT binning({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});

  auto binStorage1 = FillableT<3, double>(binning);

  /// @note Querying bin corresponding to {"test2", {1., 2.}, 1}
  FillableT<3, double>::BinType& bin = binStorage1.bin(3);

  // bin.width<2>() // Won't compile: not a continueos axes.

  if(bin.width<1>() != 1.  ||
     bin.min<1>()   != 1.  ||
     bin.max<1>()   != 2.  ||
     bin.mid<1>()   != 1.5)
      return CHECK_TEST_RES(false);

  return CHECK_TEST_RES(true);
}

auto arithmeticWrapper() {
  auto wrappedVal1 = ArithmeticWrapper<double>(15.);
  auto wrappedVal2 = ArithmeticWrapper<double>(15.);

  if(wrappedVal1 + wrappedVal2 != 30.  ||
     wrappedVal1 - wrappedVal2 != 0.   ||
     wrappedVal1 * wrappedVal2 != 225. ||
     wrappedVal1 / wrappedVal2 != 1.   ||
     wrappedVal1 != wrappedVal2)
      return CHECK_TEST_RES(false);

  if(wrappedVal1 * 2 != 30.)
    return CHECK_TEST_RES(false);

  wrappedVal1 = wrappedVal1 * 2.;
  if(wrappedVal1 != 30.)
    return CHECK_TEST_RES(false);

  // wrappedVal1 % 10; // Won't compile. WrappedVal1 wraps double, for which there is no % operator.

  if(wrappedVal1 <  wrappedVal2 ||
     wrappedVal1 == wrappedVal2)
     return CHECK_TEST_RES(false);

  auto wrappedVal3 = ArithmeticWrapper<int>(15);

  wrappedVal3 *= 2;

  if(wrappedVal3 != 30)
    return CHECK_TEST_RES(false);

  wrappedVal3 %= 10;
  if(wrappedVal3 != 0)
    return CHECK_TEST_RES(false);

  BinningT binning({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});
  using BinT = Bin<3, int, BinningT>;

  auto bin = BinT(0, binning);
  bin += 5;

  if(bin != 5)
    return CHECK_TEST_RES(false);

  return CHECK_TEST_RES(true);
}

auto dbnAdapter() {
  BinningT binning({"test1", "test2"}, {1., 2., 3., 4.}, {1, 2, 3, 4, 5});

  auto binStorage = FillableT<3, Dbn<3>>(binning);

  binStorage.fill({"test1", 1.4, 1}, std::make_index_sequence<3>(), 2.0, 1.0);
  binStorage.fill({"test1", 1.6, 1}, std::make_index_sequence<3>(), 2.0, 1.0);

  auto& bin = binStorage.binAt({"test1", 1.4, 1});

  return CHECK_TEST_RES(bin.numEntries() == 2);
}

auto getBins() {
  using BinStorageT = BinnedStorage<double, std::string, double>;
  BinStorageT binnedStorage1({"test1", "test2"}, { 1.0, 2.0, 3.0});

  std::vector<size_t> overflowIndices = {0, 1, 2, 3, 6, 9, 10, 11};

  auto binsVec1 = binnedStorage1.bins();

  for(auto& bin : binsVec1) {
    for(auto& i : overflowIndices) {
      if(bin.index() == i) return CHECK_TEST_RES(false);
    }
  }

  /// @note Double curly braces to construct temporary axes, because
  /// there is no constructor defined to work with arguments
  /// of std::initializer_list<T> and std::initializer_list<std::initializer_list<T>>
  /// type. Intermediate axis is created.
  /*BinStorageT binnedStorage2({{"test1", "test2"}}, {{{1., 2.}, {3., 4.}}});

  std::vector<size_t> hiddenBinIndices = {4, 5}; /// {0,1} + 2*2

  auto binsVec2 = binnedStorage2.bins(true);

  for(auto& bin : binsVec2) {
    for(auto& i : hiddenBinIndices) {
      if(bin.index() == i) return CHECK_TEST_RES(false);
    }
  }*/

  return CHECK_TEST_RES(true);
}

auto mergeBins() {
  using BinStorageT = BinnedStorage<double, std::string, double>;
  BinStorageT binnedStorage1({"test1", "test2"}, {1., 2., 3.});
  BinStorageT binnedStorage2({"test1", "test2"}, {1., 2., 3., 4., 5., 6.});

  // binnedStorage.mergeBins<0>({1,2}); /// Won't compile. Discrete axes are not mergable.
  binnedStorage1.mergeBins<1>({1,2});

  binnedStorage2.mergeBins<1>({0,4});

  if(binnedStorage1.numBins() != 2 ||
     binnedStorage2.numBins() != 2)
    return CHECK_TEST_RES(false);

  return CHECK_TEST_RES(true);
}

auto binnedStorageOperators() {
  using BinStorageT = BinnedStorage<double, std::string, double>;
  BinStorageT binnedStorage1({"test1", "test2"}, {1., 2., 3.});
  BinStorageT binnedStorage2({"test1", "test2"}, {1., 2., 3., 4., 5., 6.});

  binnedStorage2.mergeBins<1>({3,6});

  if(binnedStorage1 != binnedStorage2)
    return CHECK_TEST_RES(false);

  for(auto& bin : binnedStorage1.bins()) {
    bin += 1;
  }

  for(auto& bin : binnedStorage2.bins()) {
    bin += 1;
  }

  /*binnedStorage1 += binnedStorage2;

  for(auto& bin : binnedStorage1.bins()) {
    if(bin != 2) return CHECK_TEST_RES(false);
  }

  binnedStorage2 -= binnedStorage1;

  for(auto& bin : binnedStorage2.bins()) {
    if(bin != -1) return CHECK_TEST_RES(false);
  }*/

  return CHECK_TEST_RES(true);
}

auto testBinnedStorage() {
  return CHECK_TEST_RES(
    (constructors()            == EXIT_SUCCESS) &&
    (adapterConstructors()     == EXIT_SUCCESS) &&
    (set()                     == EXIT_SUCCESS) &&
    (binsVectorWrapper()       == EXIT_SUCCESS) &&
    (binSpaceCharacteristics() == EXIT_SUCCESS) &&
    (arithmeticWrapper()       == EXIT_SUCCESS) &&
    (dbnAdapter()              == EXIT_SUCCESS) &&
    (getBins()                 == EXIT_SUCCESS) &&
    (mergeBins()               == EXIT_SUCCESS) &&
    (binnedStorageOperators()  == EXIT_SUCCESS));
}

int main() {
  int rtn = EXIT_SUCCESS;

  rtn = testBinnedStorage();

  return rtn;
}

