
#include "YODA/Counter.h"
#include "YODA/BinnedDbn.h"
#include "YODA/BinnedEstimate.h"
#include "YODA/Histo.h"
#include "YODA/Profile.h"
#include "YODA/WriterYODA.h"
#include "YODA/WriterFLAT.h"
#include "YODA/Utils/TestHelpers.h"
#include <string>
#include <iostream>

using namespace YODA;

using MyBinningT = Binning<Axis<double>, Axis<int>, Axis<std::string>>;


auto testBin(const MyBinningT& binning) {

  std::cout << "# Running " << __func__ << std::endl;

  Bin<3, double, MyBinningT> bin1(85, binning);
  std::cout << "My index is " << bin1.index();
  std::cout << " and my differential volume is " << bin1.dVol();
  std::cout << std::endl;
  std::cout << "Along the continuous x-axis my (lower,upper) edges are ";
  std::cout << "(" << bin1.xMin() << "," << bin1.xMax() << ")";
  std::cout << std::endl;

  BinBase<double, MyBinningT> bin2(85, binning);
  std::cout << "My index is " << bin2.index();
  std::cout << " and my differential volume is " << bin2.dVol();
  std::cout << std::endl;
  std::cout << "Along the continuous x-axis my (lower,upper) edges are ";
  std::cout << "(" << bin2.min<0>() << "," << bin2.max<0>() << ")";
  std::cout << std::endl;

  assert(bin1.width<0>() == bin1.xWidth());

  return CHECK_TEST_RES(true);
}


auto testBinnedStorage(const MyBinningT& binning) {

  std::cout << "# Running " << __func__ << std::endl;

  // BinnedStorage<ContentType, AxisType1, AxisType2, ...>
  BinnedStorage<std::string, double, int, std::string> bs(binning);

  std::cout << "This BinnedStorage has dimension " << bs.dim();
  std::cout << " which is 1 for the content and " << bs.binDim();
  std::cout << " for the axes." << std::endl;

  std::cout << "The number of visible bins is " << bs.numBins();
  std::cout << std::endl;
  std::cout << "The number of bins including overflows is " << bs.numBins(true);
  std::cout << std::endl;

  FillableStorage<1, double, std::string> fs (Axis<std::string>({"A", "B", "C"}));
  std::cout << "This FillableStorage has fill dimension " << fs.fillDim() << std::endl;
  std::cout << "The number of visible bins is " << fs.numBins();
  std::cout << std::endl;
  std::cout << "The number of bins including overflows is " << fs.numBins(true);
  std::cout << std::endl;

  fs.fill({"A"});
  fs.fill({"B"},           10); // optional weight
  fs.fill({"C"},          100); // optional weight
  fs.fill({"not an edge"}, -1); // "other"flow
  for (const auto& b : fs.bins(true)) {
    std::cout << "Bin " << b.index() << " has content " << b.raw() << std::endl;
  }

  return CHECK_TEST_RES(true);
}


auto testBinning(const Axis<double>& caxis,
                 const Axis<int>& daxis1,
                 const Axis<std::string>& daxis2) {

  std::cout << "# Running " << __func__ << std::endl;

  MyBinningT binning(caxis, daxis1, daxis2);
  std::cout << "Binning object has " << binning.dim();
  std::cout << " dimensions and " << binning.numBins(true);
  std::cout << " number of bins in total" << std::endl;

  std::cout << "Global index given local indices (1, 5, 2) is ";
  std::cout << binning.localToGlobalIndex({1,5,2}) << std::endl;
  std::cout << "Global index given coordinates (0.5, 1, B) is ";
  std::cout << binning.globalIndexAt({0.5,1,"B"}) << std::endl;

  return CHECK_TEST_RES(
    (testBin(binning)           == EXIT_SUCCESS) &&
    (testBinnedStorage(binning) == EXIT_SUCCESS));
}


auto testBinnedAxis() {

  std::cout << "# Running " << __func__ << std::endl;

  const auto caxis = Axis<double>{{0.1, 2.3, 4.5}};

  std::cout << "Continuous axis with " << caxis.edges().size() << " floating-point edges:" << std::endl;
  for (auto edge : caxis.edges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;

  std::cout << "This axis has " << caxis.numBins(true) << " bins in total." << std::endl;

  const auto daxis1 = Axis<int>{{-3, -2, -1, 0, 1, 2, 3}};
  const auto daxis2 = Axis<std::string>{{"A", "B", "C"}};

  std::cout << "Discrete axis with " << daxis1.edges().size() << " integral edges:" << std::endl;
  for (auto edge : daxis1.edges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;

  std::cout << "Discrete axis with " << daxis2.edges().size() << " string edges:" << std::endl;
  for (auto edge : daxis2.edges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;

  std::cout << "This axis has " << daxis1.numBins(true) << " bins in total." << std::endl;
  std::cout << "This axis has " << daxis2.numBins(true) << " bins in total." << std::endl;

  return CHECK_TEST_RES((testBinning(caxis, daxis1, daxis2) == EXIT_SUCCESS));
}


auto testBinMasking(Histo2D& h2) {

  std::cout << "# Running " << __func__ << std::endl;

  std::cout << "# bins before: " << h2.numBins(true);
  std::cout << "  isMasked? " << h2.isMasked(10) << std::endl;
  h2.maskBin(10);
  std::cout << "# bins after:  " << h2.numBins(true);
  std::cout << "  isMasked? " << h2.isMasked(10) << std::endl;

  std::cout << "# bins excluding masked bins: ";
  std::cout << h2.numBins(true) << std::endl;
  std::cout << "# bins including masked bins: ";
  std::cout << h2.numBins(true, true) << std::endl;

  for (size_t i : h2.maskedBins()) {
    std::cout << "\t" << i;
  }
  std::cout << std::endl;

  return CHECK_TEST_RES(true);
}


auto testRebinning(Histo2D& h2) {

  std::cout << "# Running " << __func__ << std::endl;

  std::cout << "nBins before: " << h2.numBinsX();
  std::cout << "  " << h2.numBinsY() << std::endl;
  std::cout << "x-edges before:" << std::endl;
  for (double edge : h2.xEdges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;
  std::cout << "y-edges before:" << std::endl;
  for (double edge : h2.yEdges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;
  std::cout << "bin values before:" << std::endl;
  for (size_t idxY = 0; idxY < h2.numBinsY(); ++idxY) {
    for (size_t idxX = 0; idxX < h2.numBinsX(); ++idxX) {
      std::cout << "\t" << h2.bin(idxX+1, idxY+1).sumW();
    }
    std::cout << std::endl;
  }
  std::cout << std::endl;

  h2.rebinX(2);
  h2.rebinY(2);
  std::cout << "nBins after: " << h2.numBinsX();
  std::cout << "  " << h2.numBinsY() << std::endl;
  std::cout << "x-edges after:" << std::endl;
  for (double edge : h2.xEdges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;
  std::cout << "y-edges after:" << std::endl;
  for (double edge : h2.yEdges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;
  std::cout << "bin values after:" << std::endl;
  for (size_t idxY = 0; idxY < h2.numBinsY(); ++idxY) {
    for (size_t idxX = 0; idxX < h2.numBinsX(); ++idxX) {
      std::cout << "\t" << h2.bin(idxX+1, idxY+1).sumW();
    }
    std::cout << std::endl;
  }
  std::cout << std::endl;

  Histo1D h1(std::vector<double>{-3, -1, 0, 2, 10, 42});
  std::cout << "x-edges before:" << std::endl;
  for (double edge : h1.xEdges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;
  h1.rebinX({-3., 2., 10.});
  std::cout << "x-edges after:" << std::endl;
  for (double edge : h1.xEdges()) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;

  return CHECK_TEST_RES(true);
}


auto testEdgesAndWidths() {

  std::cout << "# Running " << __func__ << std::endl;

  Histo1D h1(std::vector<double>{-3, -1, 0, 2, 10, 42});
  std::cout << "H1 edges along x-axis:" << std::endl;
  for (auto edge : h1.xEdges(true)) {
    std::cout << "\t" << edge;
  }
  std::cout << std::endl;

  std::cout << "H1 bin widths along x-axis:" << std::endl;
  for (auto width : h1.xWidths(true)) {
    std::cout << "\t" << width;
  }
  std::cout << std::endl;

  return CHECK_TEST_RES(true);
}


auto testNanTreatment() {

  std::cout << "# Running " << __func__ << std::endl;

  const double nanval = std::numeric_limits<double>::quiet_NaN();
  Histo1D h1(5, 0.0, 1.0);
  h1.fill(1.2); h1.fill(3.4);
  std::cout << "# entries: " << h1.numEntries() << std::endl;
  h1.fill(nanval, 10);
  std::cout << "# entries: " << h1.numEntries() << std::endl;
  std::cout << "NaN count: " << h1.nanCount() << std::endl;
  std::cout << "NaN sumW:  " << h1.nanSumW()  << std::endl;
  std::cout << "NaN sumW2: " << h1.nanSumW2() << std::endl;

  return CHECK_TEST_RES(true);
}


auto testTypeReductions() {

  std::cout << "# Running " << __func__ << std::endl;

  Profile1D p1(5, 0, 1);
  Histo1D h1 = p1.mkHisto();
  assert(h1 == Histo1D(5, 0, 1));

  {
    /// Case 1: BinnedProfile(N+1)D -> BinnedHistoND
    /// The binning is reduced by one dimension,
    /// and the Dbn is reduced by two dimensions.
    BinnedHisto2D<int,std::string> h2D({1,2,3}, {"A", "B", "C"});
    std::vector<BinnedHisto<std::string>> h1d_slices = h2D.mkHistos<0>();
    std::vector<BinnedHisto<int>> h1d_slices2 = h2D.mkHistos<1>();
    BinnedHisto1D<int> h1D_i1 = h2D.mkMarginalHisto<1>();
    BinnedHisto1D<std::string> h1D_s1 = h2D.template mkMarginalHisto<0>();
    std::cout << h2D.fillDim() << " -> " << h1D_i1.fillDim();
    std::cout << " , " << h1D_s1.fillDim() << std::endl;

    /// Case 2: BinnedHisto(N+1)D -> BinnedHistoND
    /// Both the Dbn and the binning are reduced by one dimension.
    BinnedProfile2D<int,std::string> p2D({1,2,3}, {"A", "B", "C"});
    std::vector<BinnedProfile<std::string>> p1d_slices1 = p2D.mkProfiles<0>();
    std::vector<BinnedProfile<int>> p1d_slices2 = p2D.mkProfiles<1>();
    BinnedHisto1D<int> h1D_i2 = p2D.mkMarginalHisto<1>();
    BinnedHisto1D<std::string> h1D_s2 = p2D.mkMarginalHisto<0>();
    std::cout << p2D.fillDim() << " -> " << h1D_i2.fillDim();
    std::cout << " , " << h1D_s2.fillDim() << std::endl;
  }

  {
    /// Case 1: BinnedHisto(N+1)D -> BinnedProfileND
    /// The Dbn dim is the same, but the binning is reduced by one dimension.
    BinnedHisto2D<int,std::string> h2D({1,2,3}, {"A", "B", "C"});
    BinnedProfile1D<std::string> p1D_s1 = h2D.mkMarginalProfile<0>();
    BinnedProfile1D<int> p1D_i1 = h2D.mkMarginalProfile<1>();
    std::cout << h2D.fillDim() << " -> " << p1D_i1.fillDim();
    std::cout << " , " << p1D_s1.fillDim() << std::endl;

    /// Case 2: BinnedProfile(N+1)D -> BinnedProfileND
    /// Both the Dbn and the binning are reduced by one dimension.
    BinnedProfile2D<int,std::string> p2D({1,2,3}, {"A", "B", "C"});
    BinnedProfile1D<int> p1D_i2 = p2D.mkMarginalProfile<1>();
    BinnedProfile1D<std::string> p1D_s2 = p2D.mkMarginalProfile<0>();
    std::cout << p2D.fillDim() << " -> " << p1D_i2.fillDim();
    std::cout << " , " << p1D_s2.fillDim() << std::endl;
  }

  Estimate1D e1 = h1.mkEstimate();
  Estimate1D e2 = p1.mkEstimate();
  assert(e1 == e2);

  BinnedHisto2D<int,std::string> h2({1,2,3}, {"A", "B", "C"});
  BinnedEstimate2D<int,std::string> e2D = h2.mkEstimate();
  std::vector<BinnedEstimate1D<std::string>> e1D_s = h2.mkEstimates<0>(); // slice along int axis
  std::vector<BinnedEstimate1D<int>> e1D_i = h2.mkEstimates<1>(); // slice along string axis
  std::cout << e2D.type() << std::endl;
  for (const auto& ao : e1D_s) {
    std::cout << "\t" << ao.type();
  }
  std::cout << std::endl;
  for (const auto& ao : e1D_i) {
    std::cout << "\t" << ao.type();
  }
  std::cout << std::endl;

  Scatter2D s1 = h1.mkScatter();
  Scatter2D s2 = p1.mkScatter();

  return CHECK_TEST_RES(true);
}


auto testBinnedDbn() {

  std::cout << "# Running " << __func__ << std::endl;

  {
    Histo1D h1(10, 0, 100); // 10 bins between 0 and 100
    std::vector<double> edges = {0, 10, 20, 30, 40, 50};
    Histo1D h2(edges);
    BinnedHisto2D<int, std::string> h3({ 1, 2, 3 }, { "A", "B", "C" });
  }

  Histo1D h1(5, 0.0, 1.0);
  h1.fill(0.2);
  Profile1D p1(5, 0.0, 1.0);
  p1.fill(0.2, 3.5);

  size_t nbinsX = 4, nbinsY = 6;
  double lowerX = 0, lowerY = 0;
  double upperX = 4, upperY = 6;
  Histo2D h2(nbinsX, lowerX, upperX,
             nbinsY, lowerY, upperY);
  std::cout << "H2 has fillDim = " << h2.fillDim()
            << " and binDim = " << h2.binDim() << std::endl;
  std::cout << "H2 has " << h2.numBinsX() << " x "
            << h2.numBinsY() << " = "
            << h2.numBins() << " visible bins." << std::endl;
  std::cout << "Including masked bins, H2 has ";
  std::cout << h2.numBins(true) << " bins."  << std::endl;

  double w = 0;
  for (auto& b : h2.bins()) {
    h2.fill(b.xMid(), b.yMid(), ++w); // uses global bin index
  }

  for (size_t i = 0; i < h2.numBins(true); ++i) {
    assert (i == h2.bin(i).index());
  }

  std::cout << "H2 bins using local indices:" << std::endl;
  for (size_t idxY = 0; idxY < h2.numBinsY(); ++idxY) {
    for (size_t idxX = 0; idxX < h2.numBinsX(); ++idxX) {
      // indices along any given axis are 0-indexed and include under/overflows
      // offset by 1 to account for this:
      std::cout << "\t(" << idxX+1 << "," << idxY+1 << ") = ";
      std::cout << h2.bin(idxX+1, idxY+1).sumW();
    }
    std::cout << std::endl;
  }
  std::cout << std::endl;

  std::cout << "H2 bins using local indices + under/overflows:" << std::endl;
  for (size_t idxY = 0; idxY < h2.numBinsY(true); ++idxY) {
    for (size_t idxX = 0; idxX < h2.numBinsX(true); ++idxX) {
      std::cout << "\t(" << idxX << "," << idxY << ") = ";
      std::cout << h2.bin(idxX, idxY).sumW();
    }
    std::cout << std::endl;
  }
  std::cout << std::endl;

  return CHECK_TEST_RES(
    (testBinMasking(h2)   == EXIT_SUCCESS) &&
    (testRebinning(h2)    == EXIT_SUCCESS) &&
    (testEdgesAndWidths() == EXIT_SUCCESS) &&
    (testNanTreatment()   == EXIT_SUCCESS) &&
    (testTypeReductions() == EXIT_SUCCESS));
}


auto testBinnedEstimate() {

  std::cout << "# Running " << __func__ << std::endl;

  Estimate1D e1(5, 0, 1);
  const std::string typeA("uncor,typeA");
  for (auto& b : e1.bins()) {
    const double v = b.index()+1;
    b.setVal(v); b.setErr(0.1*v, typeA);
  }
  const auto covM = e1.covarianceMatrix();
  for (size_t i = 0; i < covM.size(); ++i) {
    for (size_t j = 0; j < covM[i].size(); ++j) {
      std::cout << "\t(" << i << "," << j << ")=";
      std::cout << covM[i][j];
    }
    std::cout << std::endl;
  }
  std::cout << std::endl;

  Estimate2D e2(5, 0, 1, 5, 0, 1);
  std::vector<Estimate1D> e1Ds = e2.mkEstimates<1>();

  Estimate1D e3(5, 0, 1);
  Scatter2D s1 = e3.mkScatter();

  return CHECK_TEST_RES(true);
}


auto testCounter() {

  std::cout << "# Running " << __func__ << std::endl;

  Counter c;
  c.fill();  c.fill(10);
  std::cout << "# entries:      " << c.numEntries() << std::endl;
  std::cout << "Sum of weights: " << c.sumW() << std::endl;

  return CHECK_TEST_RES(true);
}


auto testEstimate() {

  std::cout << "# Running " << __func__ << std::endl;
  Estimate e1;
  e1.setVal(42); // set central value
  e1.setErr({-4,5}); // set (down, up) error pair
  std::cout << e1._toString() << std::endl;
  std::cout << "relError()=" << e1.relErr().first;
  std::cout << ", " << e1.relErr().second << std::endl;

  const std::string typeA("uncor,typeA");
  Estimate e2(2, {-2,3}, typeA);
  std::cout << "err()=(" << e2.err(typeA).first << ",";
  std::cout << e2.err(typeA).second << ")" << std::endl;
  std::cout << "(down,up)=(" << e2.errDown(typeA);
  std::cout << "," << e2.errUp(typeA) << ")" << std::endl;

  const std::string typeC("uncor,typeC");
  e2.setErr({1,-1}, typeC);
  std::cout << "(down,up)=(" << e2.errDown(typeC) << ",";
  std::cout << e2.errUp(typeC) << ")" << std::endl;
  std::cout << "(neg,pos)=(" << e2.errNeg(typeC) << ",";
  std::cout << e2.errPos(typeC) << ")" << std::endl;

  const std::string typeD("uncor,typeD");
  e2.setErr({0.5,0.8}, typeD);
  std::cout << "(down,up)=(" << e2.errDown(typeD) << ",";
  std::cout << e2.errUp(typeD) << ")" << std::endl;
  std::cout << "(neg,pos)=(" << e2.errNeg(typeD) << ",";
  std::cout << e2.errPos(typeD) << ")" << std::endl;

  e2.setErr({-2,2}, "cor,typeB");
  std::cout << "quadSum()=(" << e2.quadSum().first << ", ";
  std::cout << e2.quadSum().second << ")" << std::endl;

  std::cout << "totalErr()=(" << e2.totalErrNeg() << ", ";
  std::cout << e2.totalErrPos() << ")" << std::endl;

  e2.setErr({-4,5});
  std::cout << "totalErr()=(" << e2.totalErrNeg() << ", ";
  std::cout << e2.totalErrPos() << ")" << std::endl;

  Estimate e3(2, {-4,3}, typeA);
  Estimate e4(1, {-3,4}, typeA);
  std::cout << (e3+e4)._toString() << std::endl;
  std::cout << (e3-e4)._toString() << std::endl;

  const std::string typeB("cor,typeB");
  Estimate e5(2, {-4,3}, typeB);
  Estimate e6(1, {-3,4}, typeB);
  std::cout << (e5+e6)._toString() << std::endl;
  std::cout << (e5-e6)._toString() << std::endl;

  std::cout << (e3/e4)._toString() << std::endl; // (dR/R) ** 2 = (dN/N) ** + (dD/D) **2
  std::cout << (e5/e6)._toString() << std::endl; // (R+dR) = (N+dN)/(D+dD)

  return CHECK_TEST_RES(true);
}


auto testPoint() {

  std::cout << "# Running " << __func__ << std::endl;

  Point1D p1(1.0, 0.1);
  Point1D p2(1.0, 0.1, 0.2);
  Point1D p3(1.0, {0.3, 0.4});
  Point2D p4(1.0, 2.0, {0.3, 0.4}, {0.5, 0.6});
  std::cout << p1.val() << " - " << p1.errs()[0].first;
  std::cout << " + " << p1.errs()[0].second << std::endl;
  std::cout << p2.val() << " - " << p2.errs()[0].first;
  std::cout << " + " << p2.errs()[0].second << std::endl;
  std::cout << p3.val() << " - " << p3.errs()[0].first;
  std::cout << " + " << p3.errs()[0].second << std::endl;
  std::cout << p4.val(0) << " - " << p4.errs()[0].first;
  std::cout << " + " << p4.errs()[0].second << "  and  ";
  std::cout << p4.val(1) << " - " << p4.errs()[1].first;
  std::cout << " + " << p4.errs()[1].second << std::endl;

  return CHECK_TEST_RES(true);
}


auto testScatter() {

  std::cout << "# Running " << __func__ << std::endl;

  Scatter2D s;
  s.addPoint({1,2}, {{0.1,0.2},{0.3,0.4}});
  std::cout << s.numPoints() << std::endl;
  for (const auto& p : s.points()) {
    std::cout << p.val(0) << " - " << p.errs()[0].first;
    std::cout << " + " << p.errs()[0].second << "  and  ";
    std::cout << p.val(1) << " - " << p.errs()[1].first;
    std::cout << " + " << p.errs()[1].second << std::endl;
  }

  return CHECK_TEST_RES(true);
}


auto testWriter() {

  std::cout << "# Running " << __func__ << std::endl;

  Counter c("C");
  c.fill(10);
  Histo1D h1(2,0,1, "H1D_d");
  h1.fill(0.1, 10);
  h1.fill(0.7,  7);
  BinnedHisto1D<std::string> h2({"A"}, "H1D_s");
  h2.fill(std::string("A"), 3);
  h2.fill(std::string("not A"), 5);
  YODA::WriterYODA::write(std::cout, c);
  YODA::WriterYODA::write(std::cout, h1);
  YODA::WriterYODA::write(std::cout, h2);

  std::vector<AnalysisObject*> aos = { &c, &h1, &h2 };
  //YODA::Writer& aoWriter = YODA::WriterYODA::create();
  //aoWriter.write("output.yoda", aos);

  YODA::WriterFLAT::write(std::cout, c);
  YODA::WriterFLAT::write(std::cout, h1);
  YODA::WriterFLAT::write(std::cout, h2);

  return CHECK_TEST_RES(true);
}


int main() {

  int res = CHECK_TEST_RES(
    (testBinnedAxis()     == EXIT_SUCCESS) &&
    (testBinnedDbn()      == EXIT_SUCCESS) &&
    (testBinnedEstimate() == EXIT_SUCCESS) &&
    (testCounter()        == EXIT_SUCCESS) &&
    (testEstimate()       == EXIT_SUCCESS) &&
    (testPoint()          == EXIT_SUCCESS) &&
    (testScatter()        == EXIT_SUCCESS) &&
    (testWriter()         == EXIT_SUCCESS));
  std::cout << "# Exiting" << std::endl;
  return res;

}

