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

using namespace YODA;

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

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


/// @brief Edge function retrieves edge corresponding to the index.
/// Indices are generated in construction time, for each value of std::vector<T>
auto dAxisEdge() {
  StrVec v1 = {"test1", "test2"};

  Axis<std::string> axis(v1);

  return CHECK_TEST_RES(axis.edge(2).compare("test2") == 0);
}

auto dAxisSize() {
  StrVec v1 = {"test1", "test2"};

  Axis<std::string> axis(v1);

  return CHECK_TEST_RES(axis.size() == 3);
}

/// @brief Axis index function returns index of element in the axis' hashmap
auto dAxisIndex() {
  StrVec v1 = {"test1", "test2"};

  Axis<std::string> axis(v1);

  return CHECK_TEST_RES(axis.index("test2") == 2) +
         CHECK_TEST_RES(axis.index("test3") == 0);
}

auto dAxisGetEdges() {
  StrVec v1 = {"test1", "test2"};

  Axis<std::string> axis(v1);

  const vec<std::string>& edges = axis.edges();
  vec<std::string> edgesCopy = axis.edges();

  std::sort(v1.begin(), v1.end());
  std::sort(edgesCopy.begin(), edgesCopy.end());

  return CHECK_TEST_RES(edgesCopy == v1) +
         CHECK_TEST_RES(edges == v1);
}

auto dAxisSameEdges() {
  StrVec v1 = {"test1", "test2"};
  StrVec v2 = {"test1"};

  Axis<std::string> axis1(v1);
  Axis<std::string> axis2(v1);
  Axis<std::string> axis3(v2);

  return CHECK_TEST_RES(axis1.hasSameEdges(axis2) && !axis1.hasSameEdges(axis3));
}

/// @brief Axis sharedEdges function return intersection of Axis' edges
auto dAxisSharedAxis() {
  StrVec v1 = {"test1", "test2"};
  StrVec v2 = {"test1", "test3", "test2"};

  StrVec expected = {"test1", "test2"};

  Axis<std::string> axis1(v1);
  Axis<std::string> axis2(v2);

  StrVec sEdges = axis1.sharedEdges(axis2);

  std::sort(sEdges.begin(), sEdges.end());
  std::sort(expected.begin(), expected.end());

  return CHECK_TEST_RES(sEdges == expected);
}

auto dAxisSubsetEdges() {
  StrVec v1 = {"test1", "test3", "test2"};
  StrVec v2 = {"test1", "test2"};

  Axis<std::string> axis1(v1);
  Axis<std::string> axis2(v2);

  return CHECK_TEST_RES(axis1.isSubsetEdges(axis2) && !axis2.isSubsetEdges(axis1));
}

auto testDAxis() {
  return CHECK_TEST_RES(
    (dAxisEdge()         == EXIT_SUCCESS) &&
    (dAxisSize()         == EXIT_SUCCESS) &&
    (dAxisIndex()        == EXIT_SUCCESS) &&
    (dAxisGetEdges()     == EXIT_SUCCESS) &&
    (dAxisSameEdges()    == EXIT_SUCCESS) &&
    (dAxisSharedAxis()   == EXIT_SUCCESS) &&
    (dAxisSubsetEdges()  == EXIT_SUCCESS));
}

auto cAxisEdge() {
  vec<double> v1 = { 3.45678, 7.123124, 512.2152515 };

  Axis<double> axis(v1);

  /// @note axis.edge[0] should be equal to -inf
  return CHECK_TEST_RES(fuzzyEquals(axis.edge(1), v1[0]));
}

auto cAxisSize() {
  vec<double> v1 = { 3.45678, 7.123124, 512.2152515 };

  Axis<double> axis(v1);

  /// v1.size() + 1: (+-inf edges) - amount of bins between edges
  return CHECK_TEST_RES(axis.size() == 4);
}

auto cAxisIndex() {
  vec<double> v1 = { 3.45678, 7.123124, 512.2152515 };

  Axis<double> axis(v1);

  return CHECK_TEST_RES(axis.index(8) == 2) +
         CHECK_TEST_RES(axis.index(0) == 0) +
         CHECK_TEST_RES(axis.index(512.21525150001) == 3);
}

auto cAxisIndexTrivial() {
  Axis<double> axis({0.0});

  return CHECK_TEST_RES(axis.numBins(true) == 2) +
         CHECK_TEST_RES(axis.numBins() == 0) +
         CHECK_TEST_RES(axis.index(-1.0) == 0) +
         CHECK_TEST_RES(axis.index( 1.0) == 1);
}

auto cAxisGetEdges() {
  vec<double> v1 = { 3.45678, 7.123124, 512.2152515 };

  Axis<double> axis(v1);
  auto edges = axis.edges();

  return CHECK_TEST_RES(std::equal(
    std::next(edges.begin()), std::prev(edges.end()), v1.begin(), fuzzyEqComp));
}
namespace YODA {
auto cAxisIsSorted() {
  vec<double> v1 = { 512.2152515, 7.123124, 3.45678 };

  Axis<double> axis(v1);
  auto edges = axis.edges();

  std::sort(v1.begin(), v1.end());

  return CHECK_TEST_RES(std::equal(
    std::next(edges.begin()), std::prev(edges.end()), v1.begin(), fuzzyEqComp));
  return 0;
}
}


auto cAxisSameEdges() {
  vec<double> v1 = { 512.2152515, 7.123124, 3.45678 };
  vec<double> v2 = { 512.2152515 };

  Axis<double> axis1(v1);
  Axis<double> axis2(v1);
  Axis<double> axis3(v2);

  return CHECK_TEST_RES(axis1.hasSameEdges(axis2) && !axis1.hasSameEdges(axis3));
}

auto cAxisSharedAxis() {
  vec<double> v1 = { 512.2152515, 7.123124, 3.45678 };
  vec<double> v2 = { 512.2152515, 3.45678 };

  vec<double> expected = { 3.45678, 512.2152515 };

  Axis<double> axis1(v1);
  Axis<double> axis2(v2);

  vec<double> sEdges = axis1.sharedEdges(axis2);

  /// since std::set_intersection preserves the order of equal elements, and
  /// all edges are sorted on construction stage, it is not necessary to sort
  /// vectors before comparison
  return CHECK_TEST_RES(std::equal(std::next(sEdges.begin()),
                                   std::prev(sEdges.end()),
                                   expected.begin(),
                                   fuzzyEqComp));
}

auto cAxisSubsetEdges() {
  vec<double> v1 = { 512.2152515, 7.123124, 3.45678 };
  vec<double> v2 = { 512.2152515, 3.45678 };

  Axis<double> axis1(v1);
  Axis<double> axis2(v2);

  return CHECK_TEST_RES(axis1.isSubsetEdges(axis2) && !axis2.isSubsetEdges(axis1));
}

auto cAxisConstructors() {
  vec<double> v1 = { 3.45678, 7.123124, 512.2152515 };
  vec<double> v2 = { 0.1, 5.0, 6.0, 10.0, 11.0, 15.0 }; /// Masked bins (gaps) at indices 2 and 4
  vec<double> v3 = { 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0 };
  vec<std::pair<double, double>> edgePairs = {{v2[0], v2[1]}, {v2[2], v2[3]}, {v2[4], v2[5]}};
  vec<size_t> maskedBins = {2, 4};


  Axis<double> axis1(v1);
  Axis<double> axis2(edgePairs);
  Axis<double> axis3(10, 0.0, 10.0);

  const vec<double>& edges1 = axis1.edges();
  const vec<double>& edges2 = axis2.edges();
  const vec<double>& edges3 = axis3.edges();

  auto compEdges = [](const vec<double>& expected, const vec<double>& edges){
    return std::equal(std::next(edges.begin()),
                      std::prev(edges.end()),
                      expected.begin());
  };

  bool testMaskedIndices = areEqualVecs(axis2.maskedBins(), maskedBins);

  return CHECK_TEST_RES(compEdges(v1, edges1) &&
                        compEdges(v2, edges2) &&
                        compEdges(v3, edges3) &&
                        testMaskedIndices);
}

auto testCAxis() {
  return CHECK_TEST_RES(
    (cAxisEdge()         == EXIT_SUCCESS) &&
    (cAxisSize()         == EXIT_SUCCESS) &&
    (cAxisIndex()        == EXIT_SUCCESS) &&
    (cAxisIndexTrivial() == EXIT_SUCCESS) &&
    (cAxisGetEdges()     == EXIT_SUCCESS) &&
    (cAxisIsSorted()     == EXIT_SUCCESS) &&
    (cAxisSameEdges()    == EXIT_SUCCESS) &&
    (cAxisSharedAxis()   == EXIT_SUCCESS) &&
    (cAxisSubsetEdges()  == EXIT_SUCCESS) &&
    (cAxisConstructors() == EXIT_SUCCESS));
}

int main() {
  int rtn = EXIT_SUCCESS;

  rtn = testDAxis() + testCAxis();

  return rtn;
}
