    /*****************************************************************************
     * Project: RooFit                                                           *
     * Package: RooFitModels                                                     *
     *    File: $Id: Roo2DKeysPdf.cc,v 1.20 2005/06/20 15:51:05 wverkerke Exp $
     * Authors:                                                                  *
     *   AB, Adrian Bevan, Liverpool University, bevan@slac.stanford.edu         *
     *                                                                           *
     * Copyright (c) 2000-2005, Regents of the University of California,         *
     *                          Liverpool University,                            *
     *                          and Stanford University. All rights reserved.    *
     *                                                                           *
     * Redistribution and use in source and binary forms,                        *
     * with or without modification, are permitted according to the terms        *
     * listed in LICENSE (http://roofit.sourceforge.net/license.txt)             *
     *****************************************************************************/

    // -- CLASS DESCRIPTION [PDF] --
    //
    // Use of the PDF:
    //  If you generate the PDF using the same domain as you fit to, for data populated
    //  near the edges, there will be no compensating 'feed through from the other side
    //  of the boundary'.  There are the following options available to you to avoid or
    //  reduce edge effects in your fits:
    //	1) enlarge your fit domain to avoid having significant data at the boundaries
    //         of your fit.
    //	2) try using the 'm' option.  This mirrors all gaussuans within the domain at
    //         the domain boundary.  This will lead to an overestimate of the PDF near the
    //         boundary instead of the previously observed deficit.
    //      3) define the fit domain [the RooRealVars in the ctor] to be smaller than the
    //         domain which the data set ctor used.  This PDF is construced using the events
    //         in the data set, so if the boundaries are really giving you a problem and
    //         mirroring is not satisfactory solution to the problem, you have to take a hit
    //         in using a reduced data set to construct the PDF from.  Recall that you will
    //         be fitting to data in a different domain [and hence dataset] so there are no
    //         inconsistencies as a result [unless you are really limited by MC stats and
    //         use the same MC to generate the PDF and test the PDF's systematics].
    //
    //

    #include "RooFit.h"

    #include "Roo2DKeysPdf.h"
    #include "Roo2DKeysPdf.h"
    #include "RooRealVar.h"
    #include "TH2.h"
    #include "TFile.h"
    #include "TBranch.h"
    #include "TMath.h"

    //#include <math.h>

    ClassImp(Roo2DKeysPdf)

     Roo2DKeysPdf::Roo2DKeysPdf(const char *name, const char *title,
                           RooAbsReal& xx, RooAbsReal & yy, RooDataSet& data,  TString options, Double_t widthScaleFactor):
      RooAbsPdf(name,title),
      x("x", "x dimension",this, xx),
      y("y", "y dimension",this, yy)
    {
      setWidthScaleFactor(widthScaleFactor);
      loadDataSet(data, options);
    }

     Roo2DKeysPdf::Roo2DKeysPdf(const Roo2DKeysPdf & other, const char* name) :
      RooAbsPdf(other,name),
      x("x", this, other.x),
      y("y", this, other.y)
    {
      if(_verbosedebug) { cout << "Roo2DKeysPdf::Roo2DKeysPdf copy ctor" << endl; }

      _xMean   = other._xMean;
      _xSigma  = other._xSigma;
      _yMean   = other._yMean;
      _ySigma  = other._ySigma;
      _n       = other._n;

      _BandWidthType    = other._BandWidthType;
      _MirrorAtBoundary = other._MirrorAtBoundary;
      _widthScaleFactor = other._widthScaleFactor;

      _2pi     = other._2pi;
      _sqrt2pi = other._sqrt2pi;
      _nEvents = other._nEvents;
      _n16     = other._n16;
      _debug   = other._debug;
      _verbosedebug   = other._verbosedebug;
      _vverbosedebug  = other._vverbosedebug;

      _lox       = other._lox;
      _hix       = other._hix;
      _loy       = other._loy;
      _hiy       = other._hiy;
      _xoffset   = other._xoffset;
      _yoffset   = other._yoffset;

      _x  = new Double_t[_nEvents];
      _y  = new Double_t[_nEvents];
      _hx = new Double_t[_nEvents];
      _hy = new Double_t[_nEvents];

      //copy the data and bandwidths
      for(Int_t iEvt = 0; iEvt< _nEvents; iEvt++)
      {
        _x[iEvt]  = other._x[iEvt];
        _y[iEvt]  = other._y[iEvt];
        _hx[iEvt] = other._hx[iEvt];
        _hy[iEvt] = other._hy[iEvt];
      }
    }

     Roo2DKeysPdf::~Roo2DKeysPdf() {
      if(_verbosedebug) { cout << "Roo2DKeysPdf::Roo2KeysPdf dtor" << endl; }
        delete[] _x;
        delete[] _hx;
        delete[] _y;
        delete[] _hy;
    }

    /////////////////////////////////////////////////////////////////////////////
    // Load a new data set into the class instance.  If the calculation fails, //
    //    return 1                                                             //
    //    return 0 indicates a success                                         //
    /////////////////////////////////////////////////////////////////////////////
     Int_t Roo2DKeysPdf::loadDataSet(RooDataSet& data, TString options)
    {
      if(_verbosedebug) { cout << "Roo2DKeysPdf::loadDataSet" << endl; }

      setOptions(options);

      if(_verbosedebug) { cout << "Roo2DKeysPdf::loadDataSet(RooDataSet& data, TString options)" << endl; }

      _2pi       = 2.0*TMath::Pi() ;   //use pi from math.h
      _sqrt2pi   = sqrt(_2pi);
      _nEvents   = (Int_t)data.numEntries();
      if(_nEvents == 0)
      {
        cout << "ERROR:  Roo2DKeysPdf::loadDataSet The input data set is empty.  Unable to begin generating the PDF" << endl;
        return 1;
      }
      _n16       =  TMath::Power(_nEvents, -0.166666666); // = (4/[n(dim(R) + 2)])^1/(dim(R)+4); dim(R) = 2

      _lox       = x.min();
      _hix       = x.max();
      _loy       = y.min();
      _hiy       = y.max();

      _x  = new Double_t[_nEvents];
      _y  = new Double_t[_nEvents];
      _hx = new Double_t[_nEvents];
      _hy = new Double_t[_nEvents];

      Double_t x0 = 0.0;
      Double_t x1 = 0.0;
      Double_t x_2 = 0.0;
      Double_t y0 = 0.0;
      Double_t y1 = 0.0;
      Double_t y_2 = 0.0;

      //check that the data contain the variable we are interested in
      Int_t bad = 0;
      const RooAbsReal & xx = x.arg();
      const RooAbsReal & yy = y.arg();
      if(! (RooRealVar*)( (RooArgSet*)data.get(0) )->find( xx.GetName() ) )
      {
        cout << "Roo2DKeysPdf::Roo2DKeysPdf invalid RooAbsReal name: "<<xx.GetName()<<" not in the data set" <<endl;
        bad = 1;
      }
      if(! (RooRealVar*)( (RooArgSet*)data.get(0) )->find( yy.GetName() ) )
      {
        cout << "Roo2DKeysPdf::Roo2DKeysPdf invalid RooAbsReal name: "<<yy.GetName()<<" not in the data set" << endl;
        bad = 1;
      }
      if(bad)
      {
        cout << "Roo2DKeysPdf::Roo2DKeysPdf Unable to initilize object; incompatible RooDataSet doesn't contain"<<endl;
        cout << "                           all of the RooAbsReal arguments"<<endl;
        return 1;
      }

      //copy the data into local arrays
      const RooArgSet * values = data.get();
      const RooRealVar* X = ((RooRealVar*)(values->find(xx.GetName())) ) ;
      const RooRealVar* Y = ((RooRealVar*)(values->find(yy.GetName())) ) ;

      for (Int_t j=0;j<_nEvents;++j)
      {
        data.get(j) ;

        _x[j] = X->getVal() ;
        _y[j] = Y->getVal() ;

        x0+=1; x1+=_x[j]; x_2+=_x[j]*_x[j];
        y0+=1; y1+=_y[j]; y_2+=_y[j]*_y[j];
      }

      //==========================================//
      //calculate the mean and sigma for the data //
      //==========================================//
      if(_nEvents == 0)
      {
        cout << "Roo2DKeysPdf::Roo2DKeysPdf Empty data set was used; can't generate a PDF"<<endl;
      }

      _xMean  = x1/x0;
      _xSigma = sqrt(x_2/_nEvents-_xMean*_xMean);

      _yMean  = y1/y0;
      _ySigma = sqrt(y_2/_nEvents-_yMean*_yMean);

      _n = Double_t(1)/(_2pi*_nEvents*_xSigma*_ySigma);

      //calculate the PDF
      return calculateBandWidth(_BandWidthType);
    }

     void Roo2DKeysPdf::setOptions(TString options)
    {
      if(_verbosedebug) { cout << "Roo2DKeysPdf::setOptions" << endl; }

      options.ToLower();
      if( options.Contains("a") )   _BandWidthType    = 0;
      else                          _BandWidthType    = 1;
      if( options.Contains("n") )   _BandWidthType    = 1;
      else                          _BandWidthType    = 0;
      if( options.Contains("m") )   _MirrorAtBoundary = 1;
      else                          _MirrorAtBoundary = 0;
      if( options.Contains("d") )   _debug            = 1;
      else                          _debug            = 0;
      if( options.Contains("v") )   { _debug         = 1; _verbosedebug = 1; }
      else                            _verbosedebug  = 0;
      if( options.Contains("vv") )  { _vverbosedebug = 1; }
      else                            _vverbosedebug = 0;

      if( _debug )
      {
        cout << "Roo2DKeysPdf::setOptions(TString options)    options = "<< options << endl;
        cout << "\t_BandWidthType    = " << _BandWidthType    << endl;
        cout << "\t_MirrorAtBoundary = " << _MirrorAtBoundary << endl;
        cout << "\t_debug            = " << _debug            << endl;
        cout << "\t_verbosedebug     = " << _verbosedebug     << endl;
        cout << "\t_vverbosedebug    = " << _vverbosedebug    << endl;
      }
    }

     void Roo2DKeysPdf::getOptions(void) const
    {
      cout << "Roo2DKeysPdf::getOptions(void)" << endl;
      cout << "\t_BandWidthType                           = " << _BandWidthType    << endl;
      cout << "\t_MirrorAtBoundary                        = " << _MirrorAtBoundary << endl;
      cout << "\t_debug                                   = " << _debug            << endl;
      cout << "\t_verbosedebug                            = " << _verbosedebug     << endl;
      cout << "\t_vverbosedebug                           = " << _vverbosedebug    << endl;
    }

    //=====================================================//
    // calculate the kernal bandwith for x & y             //
    // & Calculate the probability look up table _p[i][j]  //
    //=====================================================//
     Int_t Roo2DKeysPdf::calculateBandWidth(Int_t kernel)
    {
      if(_verbosedebug) { cout << "Roo2DKeysPdf::calculateBandWidth(Int_t kernel)" << endl; }
      if(kernel != -999)
      {
        _BandWidthType = kernel;
      }

      Double_t h = 0.0;

      Double_t sigSum       = _xSigma*_xSigma + _ySigma*_ySigma;
      Double_t sqrtSum      = sqrt( sigSum );
      Double_t sigProd      = _ySigma*_xSigma;
      if(sigProd != 0.0)  h = _n16*sqrt( sigSum/sigProd );
      if(sqrtSum == 0)
      {
        cout << "Roo2DKeysPdf::calculateBandWidth The sqr(variance sum) == 0.0. " << " Your dataset represents a delta function."<<endl;
        return 1;
      }

      Double_t hXSigma = h * _xSigma;
      Double_t hYSigma = h * _ySigma;
      Double_t xhmin   = hXSigma * sqrt(2.)/10;  //smallest anticipated bandwidth
      Double_t yhmin   = hYSigma * sqrt(2.)/10;

      //////////////////////////////////////
      //calculate bandwidths from the data//
      //////////////////////////////////////
      if(_BandWidthType == 1)  //calculate a trivial bandwith
      {
        cout << "Roo2DKeysPdf::calculateBandWidth Using a normal bandwith (same for a given dimension) based on"<<endl;
        cout << "                                 h_j = n^{-1/6}*sigma_j for the j^th dimension and n events * "<<_widthScaleFactor<<endl;
        Double_t hxGaussian = _n16 * _xSigma * _widthScaleFactor;
        Double_t hyGaussian = _n16 * _ySigma * _widthScaleFactor;
        for(Int_t j=0;j<_nEvents;++j)
        {
          _hx[j] = hxGaussian;
          _hy[j] = hyGaussian;
          if(_hx[j]<xhmin) _hx[j] = xhmin;
          if(_hy[j]<yhmin) _hy[j] = yhmin;
         }
      }
      else //use an adaptive bandwith to reduce the dependance on global data distribution
      {
        cout << "Roo2DKeysPdf::calculateBandWidth Using an adaptive bandwith (in general different for all events) [default]"<<endl;
        cout << "                                 scaled by a factor of "<<_widthScaleFactor<<endl;
        Double_t xnorm   = h * TMath::Power(_xSigma/sqrtSum, 1.5) * _widthScaleFactor;
        Double_t ynorm   = h * TMath::Power(_ySigma/sqrtSum, 1.5) * _widthScaleFactor;
        for(Int_t j=0;j<_nEvents;++j)
        {
          Double_t f_ti =  TMath::Power( g(_x[j], _x, hXSigma, _y[j], _y, hYSigma), -0.25 ) ;
          _hx[j] = xnorm * f_ti;
          _hy[j] = ynorm * f_ti;
          if(_hx[j]<xhmin) _hx[j] = xhmin;
          if(_hy[j]<yhmin) _hy[j] = yhmin;
        }
      }

      return 0;
    }

    //=======================================================================================//
    // evaluate the kernal estimation for x,y, interpolating between the points if necessary //
    //=======================================================================================//
     Double_t Roo2DKeysPdf::evaluate() const
    {
      // use the cacheing intrinsic in RFC to bypass the grid and remove
      // the grid and extrapolation approximation in the kernel estimation method
      //implementation - cheers Wouter :)
      if(_vverbosedebug) { cout << "Roo2DKeysPdf::evaluate()" << endl; }
      return evaluateFull(x,y);
    }

    /////////////////////////////////////////////////////////
    // Evaluate the sum of the product of the 2D kernels   //
    // for use in calculating the fixed kernel estimate, f //
    // given the bandwiths _hx[j] and _hy[j]               //
    /////////////////////////////////////////////////////////
    // _n is calculated once in the constructor
     Double_t Roo2DKeysPdf::evaluateFull(Double_t thisX, Double_t thisY) const
    {
      if( _vverbosedebug ) { cout << "Roo2DKeysPdf::evaluateFull()" << endl; }

      Double_t f=0.0;

      Double_t rx2, ry2, zx, zy;
      if( _MirrorAtBoundary )
      {
        for (Int_t j = 0; j < _nEvents; ++j)
        {
          rx2 = 0.0; ry2 = 0.0; zx = 0.0; zy = 0.0;
          if(_hx[j] != 0.0) rx2 = (thisX - _x[j])/_hx[j];
          if(_hy[j] != 0.0) ry2 = (thisY - _y[j])/_hy[j];

          if(_hx[j] != 0.0) zx = exp(-0.5*rx2*rx2)/_hx[j];
          if(_hy[j] != 0.0) zy = exp(-0.5*ry2*ry2)/_hy[j];

          zx += highBoundaryCorrection(thisX, _hx[j], x.max(), _x[j])
     	 +   lowBoundaryCorrection(thisX, _hx[j], x.min(), _x[j]);
          zy += highBoundaryCorrection(thisY, _hy[j], y.max(), _y[j])
     	 +   lowBoundaryCorrection(thisY, _hy[j], y.min(), _y[j]);
          f += zy * zx;
          //      f += _n * zy * zx; // ooops this is a normalisation factor :(
        }
      }
      else
      {
        for (Int_t j = 0; j < _nEvents; ++j)
        {
          rx2 = 0.0; ry2 = 0.0; zx = 0.0; zy = 0.0;
          if(_hx[j] != 0.0) rx2 = (thisX - _x[j])/_hx[j];
          if(_hy[j] != 0.0) ry2 = (thisY - _y[j])/_hy[j];

          if(_hx[j] != 0.0) zx = exp(-0.5*rx2*rx2)/_hx[j];
          if(_hy[j] != 0.0) zy = exp(-0.5*ry2*ry2)/_hy[j];
          f += zy * zx;
          //      f += _n * zy * zx; // ooops this is a normalisation factor  :(
        }
      }
      return f;
    }

    // Apply the mirror at boundary correction to a dimension given the space position to evaluate
    // at (thisVar), the bandwidth at this position (thisH), the boundary (high/low) and the
    // value of the data kernal that this correction is being applied to  tVar (i.e. the _x[ix] etc.)
     Double_t Roo2DKeysPdf::highBoundaryCorrection(Double_t thisVar, Double_t thisH, Double_t high, Double_t tVar) const
    {
      if(_vverbosedebug) { cout << "Roo2DKeysPdf::highBoundaryCorrection" << endl; }

      if(thisH == 0.0) return 0.0;
      Double_t correction = (thisVar + tVar - 2.0* high )/thisH;
      return exp(-0.5*correction*correction)/thisH;
    }

     Double_t Roo2DKeysPdf::lowBoundaryCorrection(Double_t thisVar, Double_t thisH, Double_t low, Double_t tVar) const
    {
      if(_vverbosedebug) { cout << "Roo2DKeysPdf::lowBoundaryCorrection" << endl; }

      if(thisH == 0.0) return 0.0;
      Double_t correction = (thisVar + tVar - 2.0* low )/thisH;
      return exp(-0.5*correction*correction)/thisH;
    }

    //==========================================================================================//
    // calculate f(t_i) for the bandwidths                                                      //
    //                                                                                          //
    // g = 1/(Nevt * sigma_j * sqrt2pi)*sum_{all evts}{prod d K[ exp{-(xd - ti)/sigma_jd^2} ]}  //
    //                                                                                          //
    //==========================================================================================//
     Double_t Roo2DKeysPdf::g(Double_t varMean1, Double_t * _var1, Double_t sigma1, Double_t varMean2, Double_t * _var2, Double_t sigma2) const
    {
      if((_nEvents == 0.0) || (sigma1 == 0.0) || (sigma2 == 0)) return 0.0;

      Double_t c1 = -1.0/(2.0*sigma1*sigma1);
      Double_t c2 = -1.0/(2.0*sigma2*sigma2);
      Double_t d  = 4.0*c1*c2  /(_sqrt2pi*_nEvents);
      Double_t z  = 0.0;

      for (Int_t i = 0; i < _nEvents; ++i)
      {
        Double_t r1 =  _var1[i] - varMean1;
        Double_t r2 =  _var2[i] - varMean2;
        z          += exp( c1 * r1*r1 ) * exp( c2 * r2*r2 );
      }
      z = z*d;
      return z;
    }

     Int_t Roo2DKeysPdf::getBandWidthType() const
    {
      if(_BandWidthType == 1)  cout << "The Bandwidth Type selected is Trivial" << endl;
      else                     cout << "The Bandwidth Type selected is Adaptive" << endl;

      return _BandWidthType;
    }

     Double_t Roo2DKeysPdf::getMean(const char * axis) const
    {
      if((axis == x.GetName()) || (axis == "x") || (axis == "X"))      return _xMean;
      else if((axis == y.GetName()) || (axis == "y") || (axis == "Y")) return _yMean;
      else
      {
        cout << "Roo2DKeysPdf::getMean unknown axis "<<axis<<endl;
      }
      return 0.0;
    }

     Double_t Roo2DKeysPdf::getSigma(const char * axis) const
    {
      if((axis == x.GetName()) || (axis == "x") || (axis == "X"))      return _xSigma;
      else if((axis == y.GetName()) || (axis == "y") || (axis == "Y")) return _ySigma;
      else
      {
        cout << "Roo2DKeysPdf::getSigma unknown axis "<<axis<<endl;
      }
      return 0.0;
    }

     void Roo2DKeysPdf::writeToFile(char * outputFile, const char * name) const
    {
      TString histName = name;
      histName        += "_hist";
      TString nName    = name;
      nName           += "_Ntuple";
      writeHistToFile( outputFile,    histName);
      writeNTupleToFile( outputFile,  nName);
    }

    // plot the PDf as a histogram and save to file
    // so that it can be loaded in as a Roo2DHist Pdf in the future to
    // save on calculation time
     void Roo2DKeysPdf::writeHistToFile(char * outputFile, const char * histName) const
    {
      TFile * file = 0;
      cout << "Roo2DKeysPdf::writeHistToFile This member function is temporarily disabled" <<endl;
      //make sure that any existing file is not over written
      file = new TFile(outputFile, "UPDATE");
      if (!file)
      {
        cout << "Roo2DKeysPdf::writeHistToFile unable to open file "<< outputFile <<endl;
        return;
      }

      const RooAbsReal & xx = x.arg();
      const RooAbsReal & yy = y.arg();
      RooArgSet values( RooArgList( xx, yy ));
      RooRealVar * xArg = ((RooRealVar*)(values.find(xx.GetName())) ) ;
      RooRealVar * yArg = ((RooRealVar*)(values.find(yy.GetName())) ) ;

      TH2F * hist = (TH2F*)xArg->createHistogram("hist", *yArg);
      hist = (TH2F*)this->fillHistogram(hist, RooArgList(*xArg, *yArg) );
      hist->SetName(histName);

      file->Write();
      file->Close();
    }

    // save the data and calculated bandwidths to file
    // as a record of what produced the PDF and to give a reduced
    // data set in order to facilitate re-calculation in the future
     void Roo2DKeysPdf::writeNTupleToFile(char * outputFile, const char * name) const
    {
      TFile * file = 0;

      //make sure that any existing file is not over written
      file = new TFile(outputFile, "UPDATE");
      if (!file)
      {
        cout << "Roo2DKeysPdf::writeNTupleToFile unable to open file "<< outputFile <<endl;
        return;
      }
      RooAbsReal & xArg = (RooAbsReal&)x.arg();
      RooAbsReal & yArg = (RooAbsReal&)y.arg();

      Double_t theX, theY, hx/*, hy*/;
      TString label = name;
      label += " the source data for 2D Keys PDF";
      TTree * _theTree =  new TTree(name, label);
      if(!_theTree) { cout << "Unable to get a TTree for output" << endl; return; }
      _theTree->SetAutoSave(1000000000);  // autosave when 1 Gbyte written

      //name the TBranches the same as the RooAbsReal's
      const char * xname = xArg.GetName();
      const char * yname = yArg.GetName();
      if(xname == "") xname = "x";
      if(yname == "") yname = "y";

      _theTree->Branch(xname, &theX, " x/D");
      _theTree->Branch(yname, &theY, " y/D");
      _theTree->Branch("hx",  &hx,   " hx/D");
      _theTree->Branch("hy",  &hx,  " hy/D");

      for(Int_t iEvt = 0; iEvt < _nEvents; iEvt++)
      {
        theX = _x[iEvt];
        theY = _y[iEvt];
        hx   = _hx[iEvt];
        hx   = _hy[iEvt];
        _theTree->Fill();
      }
      file->Write();
      file->Close();
    }

    /////////////////////////////////////////////////////
    // print out _p[_nPoints][_nPoints] indicating the //
    // domain limits                                   //
    /////////////////////////////////////////////////////
     void Roo2DKeysPdf::PrintInfo(ostream & out) const
    {
      out << "Roo2DKeysPDF instance domain information:"<<endl;
      out << "\tX_min          = " << _lox <<endl;
      out << "\tX_max          = " << _hix <<endl;
      out << "\tY_min          = " << _loy <<endl;
      out << "\tY_max          = " << _hiy <<endl;

      out << "Data information:" << endl;
      out << "\t<x>             = " << _xMean <<endl;
      out << "\tsigma(x)       = " << _xSigma <<endl;
      out << "\t<y>             = " << _yMean <<endl;
      out << "\tsigma(y)       = " << _ySigma <<endl;

      out << "END of info for Roo2DKeys pdf instance"<< endl;
    }