#include "TCanvas.h"
#include "TPad.h"
#include "TPaveStats.h"
#include "TH1.h"
#include "THStack.h"
#include "TStyle.h"
#include "TObjArray.h"
#include "TList.h"
#include "TMath.h"
#include "TUUID.h"
#include "TROOT.h"

TPaveStats * GetStats(TH1* h)
{ // This function returns the stats box of a histogram
  // if the histogram didn't have a stats box yet, it draws
  // the histogram to a temporary hidden canvas to generate one.
  // Fortunately the stats box persists even after the canvas is
  // destroyed.
  if(!h) return nullptr;
  TPaveStats * box = (TPaveStats *)h->FindObject("stats");
  if(!box)
    {
      gROOT->SetBatch(kTRUE);
      TUUID uniq;
      TCanvas cN(uniq.AsString(),uniq.AsString());
      cN.cd();
      h->Draw();
      gPad->Update();
      box = (TPaveStats *)h->FindObject("stats");
      gROOT->SetBatch(kFALSE);
    }
  return box;
}

TPaveStats * GetStats(TH1 & h)
{ // Same as the function above using a pointer, but it's a reference.
  // Note that it can't be a const ref, because .Draw modifies the TH1.
  TPaveStats * box = (TPaveStats *)h.FindObject("stats"); 
  if(!box)
    {
      gROOT->SetBatch(kTRUE);
      TUUID uniq;
      TCanvas cN(uniq.AsString(),uniq.AsString());
      cN.cd();
      h.Draw();
      gPad->Update();
      box = (TPaveStats *)h.FindObject("stats");
      gROOT->SetBatch(kFALSE);
    }
  return box;
}

TObjArray StatsArray(TH1 * h1 = nullptr,TH1 * h2 = nullptr,
		    TH1 * h3 = nullptr,TH1 * h4 = nullptr,
		    TH1 * h5 = nullptr,TH1 * h6 = nullptr)
{ // This function creates & returns a TObjArray filled with the
  // stat box pointers for up to six TH1s.  If they don't have stats
  // boxes, they are generated.  A more versatile thing could be done
  // with a variadic macro or template.
  TObjArray arr;
  if(h1) arr.Add(GetStats(h1)) ;
  if(h2) arr.Add(GetStats(h2)) ;
  if(h3) arr.Add(GetStats(h3)) ;
  if(h4) arr.Add(GetStats(h4)) ;
  if(h5) arr.Add(GetStats(h5)) ;
  if(h6) arr.Add(GetStats(h6)) ;
  return std::move(arr);
}

void tileStats(const TObjArray & boxes,Double_t X1NDC, Double_t X2NDC,Double_t Y1NDC, Double_t Y2NDC)
{ // Takes an array of stat box pointers and tiles them such 
  // that they fit in a box whose corner NDC coordinates are given.
  // The layout is Ceil(Sqrt(N)) x Ceil(N // Ceil(Sqrt(N))) boxes 
  // for N input boxes.
  const Int_t N = boxes.GetEntries();
  const Int_t n_across = TMath::CeilNint(TMath::Sqrt(N));
  const Int_t n_down = TMath::CeilNint((Double_t)N / (Double_t)n_across);
  const Double_t width = (X2NDC-X1NDC)/n_across;
  const Double_t height = (Y2NDC-Y1NDC)/n_down;

  Int_t across = 0;
  Int_t down = 0;
  for(Int_t i = 0;i<N;i++)
    {
      TPaveStats * tps = (TPaveStats*)(boxes[i]);
      if(!tps)
	{
	  boxes.Error("Error, contained a null pointer.","Aborting tileStats.");
	  return;
	}

      tps->SetX1NDC(X2NDC-(across+1)*width);
      tps->SetX2NDC(X2NDC-across*width);
      tps->SetY1NDC(Y2NDC-(down+1)*height);
      tps->SetY2NDC(Y2NDC-down*height);

      across++;
      if(across >= n_across)
	{
	  across = 0;
	  down++;
	}
    }
  return;
}

void tileStats(const TObjArray & boxes)
{ // Provides a default overload so the user doesn't have to
  // specify the stat box "bounding box".  The defauls are set
  // such that the individual stat boxes are the same size as a 
  // default stat box for a single histogram in a default canvas.
  const Int_t N = boxes.GetEntries();
  const Double_t defX1 = 7.80000016093254089e-01;
  const Double_t defX2 = 9.80000019073486328e-01;
  const Double_t defY1 = 7.15000011026859283e-01;
  const Double_t defY2 = 9.95000004768371582e-01;
  const Double_t defWidth = defX2-defX1;
  const Double_t defHeight = defY2-defY1;
  const Int_t n_across = TMath::CeilNint(TMath::Sqrt(N));
  const Int_t n_down = TMath::CeilNint((Double_t)N / (Double_t)n_across);

  tileStats(boxes,defX2-n_across*defWidth,defX2,defY2-n_down*defHeight,defY2);
}

void tileStats(TH1 * h1 = nullptr,TH1 * h2 = nullptr,
	       TH1 * h3 = nullptr,TH1 * h4 = nullptr,
	       TH1 * h5 = nullptr,TH1 * h6 = nullptr)
{ // Provides an overload for tileStats so that users can just put in
  // some TH1 pointers instead of having to actually create & pass around
  // a TObjArray of TPaveStats.
  tileStats(StatsArray(h1,h2,h3,h4,h5,h6));
}


void DemonstrateStats()
{
  // It's probably more illustrative if you copy & paste
  // these lines in the interpreter after doing
  // .L tileStats.C+
  gStyle->SetOptStat(1111111);

  TH1D * h1 = new TH1D("h1","h1",10,0,10);
  h1->Fill(1);
  TH1D * h2 = new TH1D("h2","h2",10,0,10);
  h2->Fill(2);

  tileStats(h1,h2);

  TCanvas * c1 = new TCanvas();
  c1->cd();

  // At this point, a user could do:
  // h1->Draw();
  // h2->Draw("sames");
  // But this is annoying to do if you have a lot
  // of histograms, and you can't do it cleanly in a loop.

  // A user could also do this:
  THStack * hs = new THStack("hs","hs");
  hs->Add(h1);
  hs->Add(h2);
  hs->Draw("nostack");
  // Which works, can be done in a loop, but for some reason
  // the stat boxes are drawn "under" the axis edges.  Could be 
  // fixed to work?
  
  // Ideally, there would be a new THStack-only option "tilestats"
  // which would automatically call the tileStats function for the 
  // collection of histograms in the THStack before drawing.  For 
  // "nostack", only the contained histograms would have stat boxes tiled.
  // For non-"nostack" drawing, the actual total stack histogram could
  // also be included.  This option would be incompatible with the "pads"
  // option.
}