// @(#)root/hist:$Id$

/** \class RooStats::HybridPlot
    \ingroup Roostats

This class provides the plots for the result of a study performed with the
HybridCalculatorOriginal class.

Authors: D. Piparo, G. Schott - Universitaet Karlsruhe

An example plot is available here:
http://www-ekp.physik.uni-karlsruhe.de/~schott/roostats/hybridplot_example.png
*/

#include "assert.h"
#include <cmath>
#include <iostream>
#include <map>

#include "RooStats/HybridPlot.h"
#include "TStyle.h"
#include "TF1.h"
#include "TAxis.h"
#include "TH1.h"
#include "TLine.h"
#include "TLegend.h"
#include "TFile.h"
#include "TVirtualPad.h"

#include <algorithm>

/// To build the THtml documentation
using namespace std;

ClassImp(RooStats::HybridPlot);

using namespace RooStats;

////////////////////////////////////////////////////////////////////////////////
/// HybridPlot constructor

HybridPlot::HybridPlot(const char* name,
                       const  char* title,
                       const std::vector<double> & sb_vals,
                       const std::vector<double> & b_vals,
                       double testStat_data,
                       int n_bins,
                       bool verbosity):
  TNamed(name,title),
  fSb_histo(NULL),
  fSb_histo_shaded(NULL),
  fB_histo(NULL),
  fB_histo_shaded(NULL),
  fData_testStat_line(0),
  fLegend(0),
  fPad(0),
  fVerbose(verbosity)
{
   int nToysSB = sb_vals.size();
   int nToysB = sb_vals.size();
   assert (nToysSB >0);
   assert (nToysB >0);

   // Get the max and the min of the plots
   double min = *std::min_element(sb_vals.begin(), sb_vals.end());
   double max = *std::max_element(sb_vals.begin(), sb_vals.end());

   double min_b = *std::min_element(b_vals.begin(), b_vals.end());
   double max_b = *std::max_element(b_vals.begin(), b_vals.end());


   if ( min_b < min) min = min_b;
   if ( max_b > max) max = max_b;

   if (testStat_data<min) min = testStat_data;
   if (testStat_data>max) max = testStat_data;

   min *= 1.1;
   max *= 1.1;

   // Build the histos

   fSb_histo = new TH1F ("SB_model",title,n_bins,min,max);
   fSb_histo->SetTitle(fSb_histo->GetTitle());
   fSb_histo->SetLineColor(kBlue);
   fSb_histo->GetXaxis()->SetTitle("test statistics");
   fSb_histo->SetLineWidth(2);

   fB_histo = new TH1F ("B_model",title,n_bins,min,max);
   fB_histo->SetTitle(fB_histo->GetTitle());
   fB_histo->SetLineColor(kRed);
   fB_histo->GetXaxis()->SetTitle("test statistics");
   fB_histo->SetLineWidth(2);

   for (int i=0;i<nToysSB;++i) fSb_histo->Fill(sb_vals[i]);
   for (int i=0;i<nToysB;++i) fB_histo->Fill(b_vals[i]);

   double histos_max_y = fSb_histo->GetMaximum();
   double line_hight = histos_max_y/nToysSB;
   if (histos_max_y<fB_histo->GetMaximum()) histos_max_y = fB_histo->GetMaximum()/nToysB;

   // Build the line of the measured -2lnQ
   fData_testStat_line = new TLine(testStat_data,0,testStat_data,line_hight);
   fData_testStat_line->SetLineWidth(3);
   fData_testStat_line->SetLineColor(kBlack);

   // The legend
   double golden_section = (std::sqrt(5.)-1)/2;

   fLegend = new TLegend(0.75,0.95-0.2*golden_section,0.95,0.95);
   TString title_leg="test statistics distributions ";
   title_leg+=sb_vals.size();
   title_leg+=" toys";
   fLegend->SetName(title_leg.Data());
   fLegend->AddEntry(fSb_histo,"SB toy datasets");
   fLegend->AddEntry(fB_histo,"B toy datasets");
   fLegend->AddEntry((TLine*)fData_testStat_line,"test statistics on data","L");
   fLegend->SetFillColor(0);
}

////////////////////////////////////////////////////////////////////////////////
/// destructor

HybridPlot::~HybridPlot()
{

  if (fSb_histo) delete fSb_histo;
  if (fB_histo) delete fB_histo;
  if (fSb_histo_shaded) delete fSb_histo_shaded;
  if (fB_histo_shaded) delete fB_histo_shaded;
  if (fData_testStat_line) delete fData_testStat_line;
  if (fLegend) delete fLegend;
}

////////////////////////////////////////////////////////////////////////////////
/// draw the S+B and B histogram in the current canvas

void HybridPlot::Draw(const char* )
{
   // We don't want the statistics of the histos
   gStyle->SetOptStat(0);

   // The histos

   if (fSb_histo->GetMaximum()>fB_histo->GetMaximum()){
      fSb_histo->DrawNormalized();
      fB_histo->DrawNormalized("same");
   }
   else{
      fB_histo->DrawNormalized();
      fSb_histo->DrawNormalized("same");
   }

   // Shaded
   fB_histo_shaded = (TH1F*)fB_histo->Clone("b_shaded");
   fB_histo_shaded->SetFillStyle(3005);
   fB_histo_shaded->SetFillColor(kRed);

   fSb_histo_shaded = (TH1F*)fSb_histo->Clone("sb_shaded");
   fSb_histo_shaded->SetFillStyle(3004);
   fSb_histo_shaded->SetFillColor(kBlue);

   // Empty the bins according to the data -2lnQ
   double data_m2lnq= fData_testStat_line->GetX1();
   for (int i=1;i<=fSb_histo->GetNbinsX();++i){
      if (fSb_histo->GetBinCenter(i)<data_m2lnq){
         fSb_histo_shaded->SetBinContent(i,0);
         fB_histo_shaded->SetBinContent(i,fB_histo->GetBinContent(i)/fB_histo->GetSumOfWeights());
      }
      else{
         fB_histo_shaded->SetBinContent(i,0);
         fSb_histo_shaded->SetBinContent(i,fSb_histo->GetBinContent(i)/fSb_histo->GetSumOfWeights());
      }
   }

   // Draw the shaded histos
   fB_histo_shaded->Draw("same");
   fSb_histo_shaded->Draw("same");

   // The line
   fData_testStat_line->Draw("same");

   // The legend
   fLegend->Draw("same");

   if (gPad) {
      gPad->SetName(GetName());
      gPad->SetTitle(GetTitle() );
   }

   fPad = gPad;

}

////////////////////////////////////////////////////////////////////////////////
/// All the objects are written to rootfile

void HybridPlot::DumpToFile (const char* RootFileName, const char* options)
{

   TFile ofile(RootFileName,options);
   ofile.cd();

   // The histos
   fSb_histo->Write();
   fB_histo->Write();

   // The shaded histos
   if (fB_histo_shaded!=NULL && fSb_histo_shaded!=NULL){
      fB_histo_shaded->Write();
      fSb_histo_shaded->Write();
   }

   // The line
   fData_testStat_line->Write("Measured test statistics line tag");

   // The legend
   fLegend->Write();

   ofile.Close();

}

////////////////////////////////////////////////////////////////////////////////

void HybridPlot::DumpToImage(const char * filename) {
   if (!fPad) {
      Error("HybridPlot","Hybrid plot has not yet been drawn ");
      return;
   }
   fPad->Print(filename);
}

////////////////////////////////////////////////////////////////////////////////
/// Perform 2 times a gaussian fit to fetch the center of the histo.
/// To get the second fit range get an interval that tries to keep into account
/// the skewness of the distribution.

double HybridPlot::GetHistoCenter(TH1* histo_orig, double n_rms, bool display_result){
   // Get the center of the histo

   TString optfit = "Q0";
   if (display_result) optfit = "Q";

   TH1F* histo = (TH1F*)histo_orig->Clone();

   // get the histo x extremes
   double x_min = histo->GetXaxis()->GetXmin();
   double x_max = histo->GetXaxis()->GetXmax();

   // First fit!

   TF1* gaus = new TF1("mygaus", "gaus", x_min, x_max);

   gaus->SetParameter("Constant",histo->GetEntries());
   gaus->SetParameter("Mean",histo->GetMean());
   gaus->SetParameter("Sigma",histo->GetRMS());

   histo->Fit(gaus,optfit);

   // Second fit!
   double sigma = gaus->GetParameter("Sigma");
   double mean = gaus->GetParameter("Mean");

   delete gaus;

   std::cout << "Center is 1st pass = " << mean << std::endl;

   double skewness = histo->GetSkewness();

   x_min = mean - n_rms*sigma - sigma*skewness/2;
   x_max = mean + n_rms*sigma - sigma*skewness/2;;

   TF1* gaus2 = new TF1("mygaus2", "gaus", x_min, x_max);
   gaus2->SetParameter("Mean",mean);

   // second fit : likelihood fit
   optfit += "L";
   histo->Fit(gaus2,optfit,"", x_min, x_max);


   double center = gaus2->GetParameter("Mean");

   if (display_result) {
      histo->Draw();
      gaus2->Draw("same");
   }
   else {
      delete histo;
   }
   delete gaus2;

   return center;


}

////////////////////////////////////////////////////////////////////////////////
/// We let an horizontal bar go down and we stop when we have the integral
/// equal to the desired one.

double* HybridPlot::GetHistoPvals (TH1* histo, double percentage){

   if (percentage>1){
      std::cerr << "Percentage greater or equal to 1!\n";
      return NULL;
   }

   // Get the integral of the histo
   double* h_integral=histo->GetIntegral();

   // Start the iteration
   std::map<int,int> extremes_map;

   for (int i=0;i<histo->GetNbinsX();++i){
      for (int j=0;j<histo->GetNbinsX();++j){
         double integral = h_integral[j]-h_integral[i];
         if (integral>percentage){
            extremes_map[i]=j;
            break;
         }
      }
   }

   // Now select the couple of extremes which have the lower bin content diff
   std::map<int,int>::iterator it;
   int a,b;
   double left_bin_center(0.),right_bin_center(0.);
   double diff=10e40;
   double current_diff;
   for (it = extremes_map.begin();it != extremes_map.end();++it){
      a=it->first;
      b=it->second;
      current_diff=std::fabs(histo->GetBinContent(a)-histo->GetBinContent(b));
      if (current_diff<diff){
         //std::cout << "a=" << a << " b=" << b << std::endl;
         diff=current_diff;
         left_bin_center=histo->GetBinCenter(a);
         right_bin_center=histo->GetBinCenter(b);
      }
   }

   double* d = new double[2];
   d[0]=left_bin_center;
   d[1]=right_bin_center;
   return d;
}

////////////////////////////////////////////////////////////////////////////////
/// Get the median of an histogram.

double HybridPlot::GetMedian(TH1* histo){

   //int xbin_median;
   Double_t* integral = histo->GetIntegral();
   int median_i = 0;
   for (int j=0;j<histo->GetNbinsX(); j++)
      if (integral[j]<0.5)
         median_i = j;

   double median_x =
      histo->GetBinCenter(median_i)+
      (histo->GetBinCenter(median_i+1)-
       histo->GetBinCenter(median_i))*
      (0.5-integral[median_i])/(integral[median_i+1]-integral[median_i]);

   return median_x;
}
