#ifndef __JCOMPAREHISTOGRAMS__JTESTKOLMOGOROV_2D__
#define __JCOMPAREHISTOGRAMS__JTESTKOLMOGOROV_2D__

#include <istream>
#include <ostream>

#include "JCompareHistograms/JResultTitle.hh"
#include "JCompareHistograms/JTest_t.hh"

#include "TMath.h"
#include "TH2.h"


/**
 * \author rgruiz, bjung
 */
namespace JCOMPAREHISTOGRAMS {}
namespace JPP { using namespace JCOMPAREHISTOGRAMS; }

namespace JCOMPAREHISTOGRAMS {

  /**
   * Implementation of the Kolmogorov test for 2D histograms.\n
   * This class is derived from the abstract class JTest_t(). For a general description of the implementation of this and other tests derived from JTest_t(), see its documentation.\n
   * This test compares two 2D histograms. If the parameter `slice` equals x, X y or Y, the histograms are sliced along the corresponding axis, and the Kolmogorov test is applied to each slice.\n 
   * If `slice` equals n or N, the histograms are not sliced, and JKolmogorovTest_2D() is applied.\n 
   * The input parameter `threshold`, is used to evaluate whether the test is passed or failed for each slice or for the full 2D distribution.\n
   * The parameter `threshold` should therefore be a real value between 0 and 1.
   */
  class JTestKolmogorov_2D: 
    public JTest_t
  { 
  public:
    
    /**
     * Default constructor.
     */    
    JTestKolmogorov_2D() :
      JTest_t("KS_2D", "p-Value(KS)")
    {}

    
    /**
     * Applies Kolmogorov test for two ROOT TH2 histograms.
     *
     * \param o1 First histogram
     * \param o2 Second histogram
     */
    void test(const TObject* o1, const TObject* o2) override
    {
      using namespace std;
      using namespace JPP;

      const TH2* h1 = dynamic_cast<const TH2*>(o1);
      const TH2* h2 = dynamic_cast<const TH2*>(o2);

      if (h1 == NULL || h2 == NULL) {
	THROW(JValueOutOfRange, "JTestKolmogorov_2D::test(): Could not cast given TObjects to TH2.");
      }
    
      const int n1x = h1->GetNbinsX();
      const int n2x = h2->GetNbinsX();
      const int n1y = h1->GetNbinsY();
      const int n2y = h2->GetNbinsY();

      if(n1x != n2x || n1y != n2y)
	THROW(JValueOutOfRange, "JTestKolmogorov_2D::test(): Histograms with different bining. The objects: " <<
	      h1->GetName() << " and " << h2->GetName() << " can not be compared." << endl);
     
      if(h1->Integral() == 0 || h2->Integral() == 0) {
	THROW(JValueOutOfRange, "JTestKolmogorov_2D::test(): Empty histogram: " <<
	      h1->GetName() << " and " << h2->GetName() << " can not be compared." << endl);
      }

      const double s1 = 1./h1->Integral();
      const double s2 = 1./h2->Integral();

      TH2* h3 = (TH2*) h1->Clone(h1->GetName() == h2->GetName() ?
				 MAKE_CSTRING(h1->GetName()                            << "_" << testName) :
				 MAKE_CSTRING(h1->GetName() << "_VS_" << h2->GetName() << "_" << testName));

      double ew1, ew2, w1 = 0, w2 = 0;

      for (int i = 1; i <= n1x; ++i) {
	for (int j = 1; j <= n1y; ++j) {
	  ew1 = h1->GetBinError(i,j);
	  ew2 = h2->GetBinError(i,j);
	  w1 += ew1*ew1;
	  w2 += ew2*ew2;
	}
      }

      bool afunc1 = false;
      bool afunc2 = false;
    
      double esum1 = 0, esum2 = 0;
      
      if (w1 > 0) {
	esum1 = 1./s1/s1/w1;
      } else {
	afunc1 = true;
      }
      
      if (w2 > 0) {
	esum2 = 1./s2/s2/w2;
      } else {
	afunc2 = true;
      }
      
      if (afunc2 && afunc1) {
	THROW(JValueOutOfRange, "JTestKolmogorov_2D::test(): Errors are zero for both histograms");
      }
      
      double c1 = 0, c2 = 0;
      
      double dmax1 = 0;
      
      for (int i=1 ; i<=n1x ; ++i){
	for (int j=1 ; j<=n1y ; ++j){
	  
	  c1 += s1*h1->GetBinContent(i,j);
	  c2 += s2*h2->GetBinContent(i,j);
	  
	  double d = TMath::Abs(c1-c2)*TMath::Sqrt(esum1*esum2/(esum1+esum2));
	  
	  dmax1 = TMath::Max(dmax1,TMath::Abs(c1-c2));
	  
	  h3->Fill(i,j,d);
	}
      }

      c1 = 0, c2 = 0;
    
      double dmax2 = 0;
      
      for (int j=1 ; j<=n1y ; ++j){
	for (int i=1 ; i<=n1x ; ++i){
	  
	  c1 += s1*h1->GetBinContent(i,j);
	  c2 += s2*h2->GetBinContent(i,j);
	  
	  double d = TMath::Abs(c1-c2)*TMath::Sqrt(esum1*esum2/(esum1+esum2));
	  
	  dmax2 = TMath::Max(dmax2,TMath::Abs(c1-c2));
	  
	  h3->Fill(i,j,d);
	}
      }
      
      double dmax = 0.5*(dmax1+dmax2);
      
      double z;

      if  (afunc1) {
	z = dmax*TMath::Sqrt(esum2);
      } else if (afunc2) {
	z = dmax*TMath::Sqrt(esum1);
      } else {
	z = dmax*TMath::Sqrt(esum1*esum2/(esum1+esum2));
      }
    
      const double pValue = TMath::KolmogorovProb(z);

      for (int i=1 ; i<=n1x ; ++i) {
	for (int j=1 ; j<=n1y ; ++j) {
	  h3->SetBinContent(i,j,TMath::KolmogorovProb(0.5 * h3->GetBinContent(i,j)));
	}
      }

      const bool passed = (pValue > threshold);

      const JResultTitle title(testName, resultType, passed, pValue);
      
      h3->SetTitle(title.getTitle().c_str());
      h3->GetZaxis()->SetTitle(resultType.c_str());      
    
      const JTestResult r (testName,
			   JRootObjectID(MAKE_STRING(h1->GetDirectory()->GetPath() << h1->GetName())),
			   JRootObjectID(MAKE_STRING(h2->GetDirectory()->GetPath() << h1->GetName())),
			   resultType, pValue, threshold, h3, passed);
      
      this->push_back(r);
    }

  
    /**
     * Read test parameters from input.
     *
     * \param  in            input stream
     * \return               input stream
     */
    std::istream& read(std::istream& in) override
    {
      using namespace JPP;
      
      in >> threshold;

      if (threshold < 0.0 || threshold > 1.0) {
	THROW(JValueOutOfRange, "JTestKolmogorov_2D::read(): Invalid threshold value " << threshold);
      }
      
      return in >> threshold;
    }
    
  
  private:
    
    double threshold;  //!< threshold p-value to decide if test is passed.
  };
}
  
#endif