/******************************************************************************
* macro to draw density                                                       *
*                                                                             *
* Author: Dr. Christian Stratowa, Vienna, Austria.                            *
* Created: 02 Oct 2004                             Last modified: 21 Oct 2007 *
******************************************************************************/

///////////////////////////
//
// in new root session(s):
//   create file "plot.root"
//     > .L macroDrawDensity.C 
//     > CreateFile() 
//
// in new root session(s):
//   draw density
//     > .L macroDrawDensity.C 
//     > DrawDensity() 
//
//   or draw densities
//     > .L macroDrawDensity.C 
//     > DrawDensity("Tree1") 
//     > DrawDensity("Tree2") 
//     > DrawDensity("Tree3") 
//
///////////////////////////


#include "TBranch.h"
#include "TCanvas.h"
#include "TFile.h"
#include "TGraph.h"
#include "TH1F.h"
#include "TKey.h"
#include "TLeaf.h"
#include "TList.h"
#include "TMath.h"
#include "TMultiGraph.h"
#include "TRandom.h"
#include "TTree.h"

#ifndef __CINT__
#include <Riostream.h>
#endif

#include "float.h"


//______________________________________________________________________________
Double_t Mean(Int_t n, const Double_t *arr)
{
   if (n == 1) return arr[0];

   Double_t mean = 0.0;
   for (Int_t i=0; i<n; i++) mean += arr[i];

   return mean/n;
}//Mean

//______________________________________________________________________________
Double_t Var(Int_t n, const Double_t *arr, const Double_t mean)
{
   if (n == 1) return 0;

   Double_t var = 0.0;
   Double_t tmp = 0.0;
   for (Int_t i=0; i<n; i++) {
      tmp = arr[i] - mean;
      var += tmp * tmp;
   }//for_i

   return var / (n - 1);
}//Var

//______________________________________________________________________________
void MassDist(Int_t nx, Double_t *x, Double_t *w, Double_t xmin, Double_t xmax,
              Int_t ny, Double_t *y)
{
   Int_t    ixmin = 0;
   Int_t    ixmax = ny - 2;
   Double_t mass  = 0.0;
   Double_t delta = (xmax - xmin) / (ny - 1);
  
   for (Int_t i=0; i<ny; i++) y[i] = 0.0;

   for (Int_t i=0; i<nx; i++) mass += w[i];
   mass = 1.0/mass;

   for(Int_t i=0; i<nx; i++) {
      if (TMath::Finite(x[i])) {
         Double_t pos = (x[i] - xmin) / delta;
         Int_t    ix  = (Int_t)TMath::Floor(pos);
         Double_t fx  = pos - ix;

         if ((ixmin <= ix) && (ix <= ixmax)) {
            y[ix]     += w[i]*(1 - fx);
            y[ix + 1] += w[i]*fx;
         } else if (ix == -1) {
            y[0] += w[i]*fx;
         } else if (ix == ixmax + 1) {
            y[ix] += w[i]*(1 - fx);
         }//if
      }//if
   }//for_i
  
   for (Int_t i=0; i<ny; i++) y[i] *= mass;
}//MassDist

//______________________________________________________________________________
void TwiddleFactor4FFT(Int_t n, Int_t i, Double_t &tf_re, Double_t &tf_im)
{
   if (i ==0 ) {
      tf_re = 1;
      tf_im = 0;
   } else {
      tf_re =  TMath::Cos(2*TMath::Pi()*(Double_t)i / (Double_t)n);  
      tf_im = -TMath::Sin(2*TMath::Pi()*(Double_t)i / (Double_t)n); 
   }//if
}//TwiddleFactor4FFT

//______________________________________________________________________________
void TwiddleFactor4IFFT(Int_t n, Int_t i, Double_t &tf_re, Double_t &tf_im)
{
   if (i == 0) {
      tf_re = 1;
      tf_im = 0;
   } else {
      tf_re = TMath::Cos(2*TMath::Pi()*(Double_t)i / (Double_t)n);  
      tf_im = TMath::Sin(2*TMath::Pi()*(Double_t)i / (Double_t)n); 
   }//if
}//TwiddleFactor4IFFT

//______________________________________________________________________________
void FFT(Int_t n, Double_t *f_re, Double_t *f_im)
{
   Int_t    baseE, baseO, blocks, points, points2;
   Double_t even_re, even_im, odd_re, odd_im, tf_re, tf_im;

   blocks = 1;
   points = 1 << n;

   for (Int_t i=0; i<n; i++) {
      points2 = points >> 1;
      baseE = 0;

      for (Int_t j=0; j<blocks; j++) {
         baseO = baseE + points2;

         for (Int_t k=0; k<points2; k++) {
            even_re = f_re[baseE + k] + f_re[baseO + k]; 
            even_im = f_im[baseE + k] + f_im[baseO + k];

            TwiddleFactor4FFT(points, k, tf_re, tf_im);

            odd_re = (f_re[baseE + k] - f_re[baseO + k])*tf_re
                   - (f_im[baseE + k] - f_im[baseO + k])*tf_im;
            odd_im = (f_re[baseE + k] - f_re[baseO + k])*tf_im
                   + (f_im[baseE + k] - f_im[baseO + k])*tf_re; 

            f_re[baseE + k] = even_re;
            f_im[baseE + k] = even_im;
            f_re[baseO + k] = odd_re;
            f_im[baseO + k] = odd_im;
         }//for_k

         baseE = baseE + points;
      }//for_j 
                   
      blocks = blocks << 1; 
      points = points >> 1;
   }//for_i
}//FFT

//______________________________________________________________________________
void IFFT(Int_t n, Double_t *f_re, Double_t *f_im)
{
   Int_t    blocks, points, points2, baseB, baseT;
   Double_t top_re, top_im, bot_re, bot_im, tf_re, tf_im;

   blocks = 1 << (n-1);
   points = 2;  
   for (Int_t i=0; i<n; i++) {
      points2 = points >> 1;
      baseT = 0;

      for (Int_t j=0; j<blocks; j++) {
         baseB = baseT + points2;

         for (Int_t k=0; k<points2; k++) {
            top_re = f_re[baseT + k];
            top_im = f_im[baseT + k];
	
            TwiddleFactor4IFFT(points, k, tf_re, tf_im);

            bot_re = f_re[baseB + k]*tf_re - f_im[baseB + k]*tf_im;
            bot_im = f_re[baseB + k]*tf_im + f_im[baseB + k]*tf_re;

            f_re[baseT + k] = top_re + bot_re;
            f_im[baseT + k] = top_im + bot_im;
            f_re[baseB + k] = top_re - bot_re;   
            f_im[baseB + k] = top_im - bot_im; 
         }//for_k  
  
         baseT = baseT + points; 
      }//for_j

      blocks = blocks >> 1;
      points = points << 1;
   }//for_i
}//IFFT

//______________________________________________________________________________
Int_t FFTDensityConvolve(Int_t n, Double_t *x_re, Double_t *y_re)
{
   Int_t err = 0;

   // to stop rounding problems 
   Int_t nlog2 = (Int_t)(TMath::Log((Double_t)n) / TMath::Log(2.0) + 0.5);

// Init local arrays
   Double_t *x_im = 0;
   Double_t *y_im = 0;
   Double_t *c_re = 0;
   Double_t *c_im = 0;

// Create local arrays
   if (!(x_im = new Double_t[n])) {err = 1; goto cleanup;}
   if (!(y_im = new Double_t[n])) {err = 1; goto cleanup;}
   if (!(c_re = new Double_t[n])) {err = 1; goto cleanup;}
   if (!(c_im = new Double_t[n])) {err = 1; goto cleanup;}

   FFT(nlog2, y_re, y_im);
   FFT(nlog2, x_re, x_im);
  
   for (Int_t i=0; i<n; i++){
      c_re[i] = y_re[i]*x_re[i] + y_im[i]*x_im[i];
      c_im[i] = y_im[i]*x_re[i] - y_re[i]*x_im[i];
   }//for_i
  
   IFFT(nlog2, c_re, c_im);

   for (Int_t i=0; i<n; i++){
      x_re[i] = c_re[i];
   }//for_i

cleanup:
   delete [] c_im;
   delete [] c_re;
   delete [] y_im;
   delete [] x_im;

   return err;
}//FFTDensityConvolve

//______________________________________________________________________________
void Kernelize(Int_t n, Double_t *x, Double_t bw, const char *kernel)
{
   Double_t a, ax;

   if (strcmp(kernel, "gaussian") == 0) {
   // Gaussian Kernel
      for (Int_t i=0; i<n; i++) {
         x[i] = TMath::Gaus(x[i], 0, bw, kTRUE);
      }//for_i
   } else if (strcmp(kernel, "epanechnikov") == 0) {
   // Epanechnikov Kernel
      a = bw * TMath::Sqrt(5.0);
      for (Int_t i=0; i<n; i++) {
         ax = TMath::Abs(x[i]);
         if (ax < a) {
            x[i] = 3.0/(4.0*a)*(1.0 - (ax/a)*(ax/a));
         } else {
            x[i] = 0.0;
         }//if
      }//for_i
   } else if (strcmp(kernel, "rectangular") == 0) {
   // Rectangular Kernel
      a = bw * TMath::Sqrt(3.0);
      for (Int_t i=0; i<n; i++) {
         ax = TMath::Abs(x[i]);
         if (ax < a) {
            x[i] = 0.5 / a;
         } else {
            x[i] = 0.0;
         }//if
      }//for_i
   } else if (strcmp(kernel, "triangular") == 0) {
   // Triangular Kernel
      a = bw * TMath::Sqrt(6.0);
      for (Int_t i=0; i<n; i++) {
         ax = TMath::Abs(x[i]);
         if (ax < a) {
            x[i] = (1.0 - ax/a) / a;
         } else {
            x[i] = 0.0;
         }//if
      }//for_i
   } else if (strcmp(kernel, "biweight") == 0) {
   // Biweight Kernel
      a = bw * TMath::Sqrt(7.0);
      for (Int_t i=0; i<n; i++) {
         ax = TMath::Abs(x[i]);
         if (ax < a) {
            x[i] = 15.0/16.0 * (1.0 - (ax/a)*(ax/a)) * (1.0 - (ax/a)*(ax/a)) / a;
         } else {
            x[i] = 0.0;
         }//if
      }//for_i
   } else if (strcmp(kernel, "cosine") == 0) {
   // Cosine Kernel
      a = bw / TMath::Sqrt(1.0/3.0 - 2.0/(TMath::Pi()*TMath::Pi()));
      for (Int_t i=0; i<n; i++) {
         ax = TMath::Abs(x[i]);
         if (ax < a) {
            x[i] = (1.0 + TMath::Cos(TMath::Pi()*x[i]/a)) / (2*a);
         } else {
            x[i] = 0.0;
         }//if
      }//for_i
   } else if (strcmp(kernel, "optcosine") == 0) {
   // Optcosine Kernel
      a = bw / TMath::Sqrt(1.0 - 8.0/(TMath::Pi()*TMath::Pi()));
      for (Int_t i=0; i<n; i++) {
         ax = TMath::Abs(x[i]);
         if (ax < a) {
            x[i] = (TMath::Pi()/4.0)*TMath::Cos(TMath::Pi()*x[i]/(2*a)) / a;
         } else {
            x[i] = 0.0;
         }//if
      }//for_i
   }//if
}//Kernelize

//______________________________________________________________________________
Double_t Bandwidth(Int_t n, Double_t *x, Double_t iqr)
{
   Double_t hi, lo;
  
//   hi = TMath::RMS(n, x); // (x-m)/n not (x-m)/(n-1) =>is not stdev!!
   hi = TMath::Sqrt(Var(n, x, Mean(n, x)));
  
   if (hi > iqr) lo = iqr/1.34;
   else          lo = hi;

   if (lo == 0) {
      if      (hi != 0)               lo = hi;
      else if (TMath::Abs(x[1]) != 0) lo = TMath::Abs(x[1]);
      else                            lo = 1.0;
   }//if

   return (0.9*lo*TMath::Power((Double_t)n, -0.2));
}//Bandwidth 

//______________________________________________________________________________
void LinearInterpolate(Double_t *xin, Double_t *yin,
                       Int_t nout, Double_t *xout, Double_t *yout)
{
   Int_t i, j, ij;

   for (Int_t k=0 ; k<nout; k++) {
      i = 0;
      j = nout - 1;
    
      if(xout[k] < xin[i]) {yout[k] = yin[0];      continue;}
      if(xout[k] > xin[j]) {yout[k] = yin[nout-1]; continue;}
 
      // find the correct interval by bisection
      while (i < j - 1) {  // xin[i] <= xout[k] <= xin[j]
         ij = (i + j)/2;   // i+1 <= ij <= j-1
         if (xout[k] < xin[ij]) j = ij;
         else                   i = ij; // still i < j
      }//while
  
      if (xout[k] == xin[j]) {yout[k] = yin[j]; continue;}
      if (xout[k] == xin[i]) {yout[k] = yin[i]; continue;}
  
      yout[k] = yin[i] + (yin[j] - yin[i])*((xout[k] - xin[i])/(xin[j] - xin[i]));
   }//for_k
}//LinearInterpolate

//______________________________________________________________________________
Int_t Density(Int_t n, Double_t *x, Double_t *w, Int_t nout,  Double_t *xout,
              Double_t *yout, const char *kernel = "epanechnikov")
{
   Int_t err   = 0;
   Int_t nout2 = 2*nout;
   Double_t lo, hi, iqr, bw, from, to;

// Init local arrays
   Int_t    *indx = 0;
   Double_t *sort = 0;
   Double_t *xin  = 0;
   Double_t *yin  = 0;
   Double_t *yord = 0;

// Create local arrays
   if (!(indx = new Int_t[n]))        {err = 1; goto cleanup;}
   if (!(sort = new Double_t[n]))     {err = 1; goto cleanup;}
   if (!(xin  = new Double_t[nout]))  {err = 1; goto cleanup;}
   if (!(yin  = new Double_t[nout2])) {err = 1; goto cleanup;}
   if (!(yord = new Double_t[nout2])) {err = 1; goto cleanup;}

// Sort x to get lo, hi and interquartile range
//??   TMath::Sort(n, x, indx);
   TMath::Sort(n, x, indx, kFALSE);
   for (Int_t i=0; i<n; i++) sort[i] = x[indx[i]];
  
   lo  = sort[0];
   hi  = sort[n-1];
   iqr = sort[(Int_t)(0.75*n + 0.5)] - sort[(Int_t)(0.25*n + 0.5)];

// Bandwidth  
   bw = Bandwidth(n, x, iqr);
  
   lo = lo - 7*bw;
   hi = hi + 7*bw;

   for (Int_t i=0; i<=nout; i++) {
      yin[i] = (Double_t)i/(Double_t)(nout2 - 1)*2*(hi - lo);
   }//for_i

   for (Int_t i=nout+1; i<nout2; i++) {
      yin[i] = -yin[nout2 - i];
   }//for_i

   Kernelize(nout2, yin, bw, kernel);

   MassDist(n, x, w, lo, hi, nout, yord);

   err = FFTDensityConvolve(nout2, yin, yord);
   if (err != 0) goto cleanup;

   // corrections to get correct output range 
   to   = hi - 4*bw;
   from = lo + 4*bw;

   for (Int_t i=0; i<nout; i++) {
      xin[i]  = (Double_t)i / (Double_t)(nout - 1)*(hi - lo)   + lo;
      xout[i] = (Double_t)i / (Double_t)(nout - 1)*(to - from) + from;
   }//for_i

   for (Int_t i=0; i<nout; i++) {
      yin[i] = yin[i]/nout2;
   }//for_i

   LinearInterpolate(xin, yin, nout, xout, yout);

cleanup:
   delete [] xin;
   delete [] yord;
   delete [] yin;
   delete [] sort;
   delete [] indx;

   return err;
}//Density

//______________________________________________________________________________
void TestDensity(const char *kernel = "epanechnikov")
{
   Int_t npts = 512;
   Int_t nrow = 20;
   Double_t x[] = {-20,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,20};

   Double_t *weights = new Double_t[nrow];
   Double_t *dens_x  = new Double_t[npts];
   Double_t *dens_y  = new Double_t[npts];

   for (Int_t j=0; j<nrow; j++) {
      weights[j] = 1.0;
   }
  
   Density(nrow, x, weights, npts, dens_x, dens_y, kernel);

   TCanvas *fCanvas = new TCanvas("test", "test", 10, 10, 400, 400);
   TGraph  *graph   = new TGraph(npts, dens_x, dens_y);
   graph->Draw("ACP");
   fCanvas->Update();

//   delete [] dens_y;
//   delete [] dens_x;
   delete [] weights;
}//TestDensity

//______________________________________________________________________________
void CreateFile()
{
// create "Plot.root"
   TFile f("Plot.root","recreate");

// Fill Tree1 (adopted from macro tree1.C)
   TTree t1("Tree1","a simple Tree with simple variables");
   Float_t px, py, pz;
   Double_t random;
   Int_t ev;
   t1.Branch("px",&px,"px/F");
   t1.Branch("py",&py,"py/F");
   t1.Branch("pz",&pz,"pz/F");
   t1.Branch("random",&random,"random/D");
   t1.Branch("ev",&ev,"ev/I");
      // fill the tree
   for (Int_t i=0;i<10000;i++) {
      gRandom->Rannor(px,py);
      px = 11*px;
      py = 13*py;
      pz = px*px + py*py + 300;
//      random = gRandom->Rndm();
      random = gRandom->Gaus(1.6);
      ev = i;
      t1.Fill();
   }
   // write tree to file
  t1.Write();

// Fill Tree2
   TTree t2("Tree2","a simple Tree with simple variables");
   t2.Branch("px",&px,"px/F");
   t2.Branch("py",&py,"py/F");
   t2.Branch("pz",&pz,"pz/F");
   t2.Branch("random",&random,"random/D");
   t2.Branch("ev",&ev,"ev/I");
      // fill the tree
   for (Int_t i=0;i<10000;i++) {
      gRandom->Rannor(px,py);
      px = 12*px;
      py = 12*py;
      pz = px*px + 20*py*py + 200;
//      random = gRandom->Rndm();
      random = gRandom->Gaus(1.1);
      ev = i;
      t2.Fill();
   }
   // write tree to file
  t2.Write();

// Fill Tree3
   TTree t3("Tree3","a simple Tree with simple variables");
   t3.Branch("px",&px,"px/F");
   t3.Branch("py",&py,"py/F");
   t3.Branch("pz",&pz,"pz/F");
   t3.Branch("random",&random,"random/D");
   t3.Branch("ev",&ev,"ev/I");
      // fill the tree
   for (Int_t i=0;i<10000;i++) {
      gRandom->Rannor(px,py);
      px = 11*px;
      py = 13*py;
      pz = px*px + 30*py*py;
//      random = gRandom->Rndm();
      random = gRandom->Gaus(-1.3);
      ev = i;
      t3.Fill();
   }
   // write tree to file
  t3.Write();
}//CreateFile

//______________________________________________________________________________
void DrawDensity(const char *treename = "*", const char *leafname = "random",
                 Option_t *opt="L", const char *kernel = "epanechnikov",
                 Int_t npts = 512)
{
// Open file
   TFile *file = TFile::Open("Plot.root", "READ");
   if (!file) return;

   TTree   *tree  = 0;
   TLeaf   *leaf  = 0;
   TBranch *brch  = 0;

// Add trees to list
   TList *fTrees = new TList();
   if (strcmp(treename,"*") == 0) {
      TKey *key = 0;
      TIter next(file->GetListOfKeys());
      while ((key = (TKey*)next())) {
         tree = (TTree*)file->Get(key->GetName());
         fTrees->Add(tree);
      }//while
   } else {
      tree = (TTree*)file->Get(treename);
      fTrees->Add(tree);
   }//if

   Int_t numtrees = fTrees->GetSize();
   Int_t entries  = (Int_t)(tree->GetEntries());
cout << "numtrees= " << numtrees << endl;
cout << "entries=  " << entries << endl;

// New canvas and multigraph
   TCanvas     *fCanvas = new TCanvas("test", "test", 10, 10, 400, 400);
   TMultiGraph *mgraph  = new TMultiGraph();

   TH1F   *frame = 0;
   TGraph *graph = 0;

   TString title = "";
   TString titleX = "";
   TString titleY = "";

   Double_t fMinX =  DBL_MAX;
   Double_t fMinY =  DBL_MAX;
   Double_t fMaxX = -DBL_MAX;
   Double_t fMaxY = -DBL_MAX;
   Double_t minX  =  DBL_MAX;
   Double_t minY  =  DBL_MAX;
   Double_t maxX  = -DBL_MAX;
   Double_t maxY  = -DBL_MAX;
   Double_t value =  0;

   Double_t *index = 0;
   Double_t *arr   = 0;
   Double_t *wght  = 0;
   Double_t *xden  = 0;
   Double_t *yden  = 0;
   if (!(index = new Double_t[entries])) {goto cleanup;} 
   if (!(arr   = new Double_t[entries])) {goto cleanup;}  
   if (!(wght  = new Double_t[entries])) {goto cleanup;} 
   if (!(xden  = new Double_t[npts]))    {goto cleanup;} 
   if (!(yden  = new Double_t[npts]))    {goto cleanup;} 

   for (Int_t k=0; k<entries; k++) wght[k] = 1.0; 
   for (Int_t k=0; k<npts;    k++) xden[k] = yden[k] = 0;

// Loop over trees and fill graphs
   for (Int_t i=0; i<numtrees; i++) {
      tree = (TTree*)(fTrees->At(i));
cout << "tree= " << tree->GetName() << endl;
      if (!(leaf = tree->FindLeaf(leafname))) {
         cerr << "Error: Leaf <" << leafname << "> not found." << endl;
         goto cleanup;
      }//if
      if (!(brch = leaf->GetBranch())) {goto cleanup;}

   // BEGIN fill arrays
//      FillArrays(entries, brch, leaf, index, arr);
      for (Int_t k=0; k<entries; k++) index[k] = k + 1;

      // min/max for index, i.e. x-axis
      fMinX = index[0];
      fMaxX = index[entries-1];
      // min/max for y-axis
      fMinY = DBL_MAX;
      fMaxY = -DBL_MAX;

      for (Int_t k=0; k<entries; k++) { 
         brch->GetEntry(k);
         arr[k] = leaf->GetValue();
         fMinY  = (fMinY < arr[k]) ? fMinY : arr[k];
         fMaxY  = (fMaxY > arr[k]) ? fMaxY : arr[k];
      }//for
   // END fill arrays
cout << "arr: fMinX= " << fMinX << "   fMinY= " << fMinY << "   fMaxX= " << fMaxX << "   fMaxY= " << fMaxY << endl;

      // density
      Density(entries, arr, wght, npts, xden, yden, kernel);

      // fill graphs
      graph = new TGraph(npts, xden, yden);

      value = TMath::MinElement(npts, xden);
      minX = (minX > value) ? value : minX;
      value = TMath::MinElement(npts, yden);
      minY = (minY > value) ? value : minY;
      value = TMath::MaxElement(npts, xden);
      maxX = (maxX < value) ? value : maxX;
      value = TMath::MaxElement(npts, yden);
      maxY = (maxY < value) ? value : maxY;
cout << "den: minX= " << minX << "   minY= " << minY << "   maxX= " << maxX << "   maxY= " << maxY << endl;

      mgraph->Add(graph, opt);
   }//for_i

// Set titles
   title = "Leaf <" + TString(leafname) + "> for <" ;
   title += numtrees;
   title += "> trees";
   titleX = TString(leafname);
   titleY = TString("Density");

// Draw frame
   frame = gPad->DrawFrame(minX - 0.2*minX, minY - 0.2*minY, maxX + 0.2*maxX, maxY + 0.2*maxY);
cout << "pad: minX= " << minX << "   minY= " << minY << "   maxX= " << maxX << "   maxY= " << maxY << endl;

   frame->SetTitle(title);
   frame->SetXTitle(titleX);
   frame->SetYTitle(titleY);
   frame->GetXaxis()->CenterTitle(kTRUE);
   frame->GetYaxis()->CenterTitle(kTRUE);

// Draw multigraph
   mgraph->Draw(opt); 

// Cleanup
cleanup:
   if (index) {delete [] index; index = 0;}
   if (arr)   {delete [] arr;   arr   = 0;}
   if (yden)  {delete [] yden;  yden  = 0;}
   if (xden)  {delete [] xden;  xden  = 0;}
   if (wght)  {delete [] wght;  wght  = 0;}
   if(fTrees) {fTrees->Delete(); delete fTrees; fTrees = 0;}
}//DrawDensity