//=============================================================================
/** 
 * @file modelBuilder.C
 * @brief Build an example model
 *
 * @author Tim-Philip Huecking, <huecking@hep.physik.uni-siegen.de>
 * 
 * Created:
 * @date 20-05-25 (TPH)
 *
 * Changed:
 *
 */
//=============================================================================

// system includes

// ROOT includes 
#include<RooRealVar.h>
#include<RooAddPdf.h>
#include<RooExponential.h>
#include<RooExtendPdf.h>
#include<RooGaussian.h>
#include<RooProdPdf.h>
#include<RooArgList.h>
#include<TBenchmark.h>
#include<RooRandom.h>
#include<TFile.h>

// local includes
#include "Config.C"

/**
 * @class ModelBuilder
 * @brief Build a model with 3 variables:
 *    - mass (bkg: Exponential, sig: Gaussian)
 *    - vgg (bkg: Exponential, sig: Exponential)
 *    - vexex (bkg: Gaussian, sig: Gaussian)
 * @Note This class became neccessary, becauese of #9
 *
 */
class ModelBuilder
{
   public:
      ModelBuilder (Configurator &cfg);
      virtual ~ModelBuilder ();
      /** @brief Build full model for our pseudo experiments */
      void buildFullModel(const bool setRndSeed=false);
      //
      // variables
      //
      
      // observables
     
      /** discriminating variable */
      unique_ptr<RooRealVar> m_mass{};
      /** control variable */
      unique_ptr<RooRealVar> m_vgg{};
      /** control variable */
      unique_ptr<RooRealVar> m_vexex{};

      // event numbers
   
      /** #Sig events to fit */
      unique_ptr<RooRealVar> m_nSigFit{};
      /** #Bkg events to fit */
      unique_ptr<RooRealVar> m_nBkgFit{};

      // mass
      
      //sig
      unique_ptr<RooRealVar> m_massSigMeanFit{};
      unique_ptr<RooRealVar> m_massSigStdDevFit{};
      unique_ptr<RooGaussian> m_pdfMassSig{};
      unique_ptr<RooExtendPdf> m_pdfExtMassSig{};
      //bkg
      unique_ptr<RooRealVar> m_massBkgLam{};
      unique_ptr<RooExponential> m_pdfMassBkg{};
      unique_ptr<RooExtendPdf> m_pdfExtMassBkg{};
      //sum
      shared_ptr<RooAddPdf> m_pdfMassSum{};

      // vgg
      
      //sig 
      unique_ptr<RooRealVar> m_vggSigMeanFit{};
      unique_ptr<RooRealVar> m_vggSigStdDevFit{};
      unique_ptr<RooGaussian> m_pdfVggSig{};
      unique_ptr<RooExtendPdf> m_pdfExtVggSig{};
      //bkg
      unique_ptr<RooRealVar> m_vggBkgMeanFit{};
      unique_ptr<RooRealVar> m_vggBkgStdDevFit{};
      unique_ptr<RooGaussian> m_pdfVggBkg{};
      unique_ptr<RooExtendPdf> m_pdfExtVggBkg{};
      //sum
      unique_ptr<RooAddPdf> m_pdfVggSum{};
      

      // product pdfs 

      //sig
      unique_ptr<RooProdPdf> m_pdfSigProd{};
      unique_ptr<RooExtendPdf> m_pdfExtSigProd{};
      //bkg
      unique_ptr<RooProdPdf> m_pdfBkgProd{};
      unique_ptr<RooExtendPdf> m_pdfExtBkgProd{};

      // full model
      unique_ptr<RooAddPdf> m_pdfFull{};


   private:
      //
      // functions
      //
      /** Build Mass signal pdf */
      void bldMassSig();
      /** Build Mass bkg pdf */
      void bldMassBkg();
      /** Build full mass pdf */
      void bldMass();
      /** Build vgg signal pdf */
      void bldVggSig();
      /** Build vgg bkg pdf */
      void bldVggBkg();
      /** Build full vgg pdf */
      void bldVgg();
      /** Build signal, bkg and full pdf */
      void bld_bkg_sig_full();
      /** Check if event numbers and obeservables have been defined
      * The build functions can not be called in arbitrary order 
      */
      bool CheckEventNumsAndObs();
      /** Remove constantness from a RooRealVar. This way we get a range from
      * -INF to +INFT. I do not another way to do that.
      */
      void unsetConst(unique_ptr<RooRealVar>& rrv);

      // 
      // variables
      //
      
      /** Indicate what is the status. Each model building function will
       * use it and do nothing if it is false
       */
      bool m_status = true;
      
      /** Config object */
      Configurator &m_cfg; 

      /** Temporary name for the next RooFit object */
      string m_tmpNam = "";

};


using namespace RooFit; 

ModelBuilder::ModelBuilder(Configurator &cfg)
: m_cfg{cfg}
{
}

ModelBuilder::~ModelBuilder()
{
}

void ModelBuilder::buildFullModel(const bool setRndSeed)
{

   // Set a random seed
   if(setRndSeed){RooRandom::randomGenerator()->SetSeed(0);}
   else{RooRandom::randomGenerator()->SetSeed(7);}

   // build observables
   m_mass.reset(new RooRealVar{ "mass", "mass - discriminating variable"
          , m_cfg.massRangeMin, m_cfg.massRangeMax, "a.u."});
   m_vgg.reset(new RooRealVar( "vgg"
         , "variable with gaussian bkg and sig - Control variable"
         ,  m_cfg.vggRangeMin, m_cfg.vggRangeMax , "a.u."));

   // build event numbers
   m_nSigFit.reset(new RooRealVar{"nSigFit"
         , "Number of signal events to be fitted", m_cfg.nSigExpVal }); 
   unsetConst(m_nSigFit);
   m_nBkgFit.reset(new RooRealVar{"nBkgFit"
         , "Number of bkg events to be fitted", m_cfg.nBkgExpVal}); 
   unsetConst(m_nBkgFit);
   
   // build pdfs
   bldMassSig();
   bldMassBkg();
   bldMass();
   bldVggSig();
   bldVggBkg();
   bldVgg();
   bld_bkg_sig_full();

   return;
}

//
// unsetConst
//
void ModelBuilder::unsetConst(unique_ptr<RooRealVar>& rrv)
{
   rrv->setConstant(kFALSE);
   return;
}


//
// CheckEventNumsAndObs
// 
bool ModelBuilder::CheckEventNumsAndObs()
{
   bool there = false;
   if(m_nSigFit and m_nBkgFit and m_mass and m_vgg)
   {
      return true;
   }
   prntErr("Wrong Usage of ModelBuilder. Necessary variables undefined"
         , "ModelBuilder::CheckEventNumsAndObs()");
   return there;
}

//
// ModelBuilder::bldMassSig()
//
void ModelBuilder::bldMassSig()
{
   // checks
   cout << "ModelBuilder::bldMassSig()" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   //build
   m_tmpNam = "massSigMeanFit";
   m_massSigMeanFit.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.massSigMeanTruth});
   unsetConst(m_massSigMeanFit);
   
   m_tmpNam = "massSigStdDevFit";
   m_massSigStdDevFit.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.massSigStdDevTruth});
   unsetConst(m_massSigStdDevFit);

   m_tmpNam = "pdfMassSig";
   m_pdfMassSig.reset(new RooGaussian{m_tmpNam.c_str(), m_tmpNam.c_str()
         , *m_mass.get(), *m_massSigMeanFit.get(), *m_massSigStdDevFit.get()});

   // make extended
   m_tmpNam = "pdfExtMassSig";
   m_pdfExtMassSig.reset(new RooExtendPdf{m_tmpNam.c_str()
         , m_tmpNam.c_str(), *m_pdfMassSig.get(), *m_nSigFit.get()});
   return;
}

//
// ModelBuilder::bldmassbkg()
//
void ModelBuilder::bldMassBkg()
{
   cout << "ModelBuilder::bldMassbkg" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   m_tmpNam = "massBkgLamFit";
   m_massBkgLam.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.massBkgLamTruth});
   unsetConst(m_massBkgLam);

   m_tmpNam = "pdfMassBkg";
   m_pdfMassBkg.reset(new RooExponential{m_tmpNam.c_str(), m_tmpNam.c_str()
         , *m_mass.get(), *m_massBkgLam.get()});

   // make extended
   m_tmpNam = "pdfExtMassBkg";
   m_pdfExtMassBkg.reset(new RooExtendPdf{m_tmpNam.c_str()
         , m_tmpNam.c_str(), *m_pdfMassBkg.get(), *m_nBkgFit.get()});
   return;
}

//
// ModelBuilder::bldMass()
//
void ModelBuilder::bldMass()
{
   cout << "ModelBuilder::bldMass()" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   m_tmpNam = "pdfMassSum";
   m_pdfMassSum.reset(new RooAddPdf{m_tmpNam.c_str(), m_tmpNam.c_str()
         , RooArgList(*m_pdfMassSig.get(), *m_pdfMassBkg.get())
         , RooArgList(*m_nSigFit.get(), *m_nBkgFit.get())});
   return;
}

//
// ModelBuilder::bldVggSig
//
void ModelBuilder::bldVggSig()
{
   cout << "ModelBuilder::bldVggSig" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   //build
   m_tmpNam = "vggSigMeanFit";
   m_vggSigMeanFit.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.vggSigMeanTruth});
   unsetConst(m_vggSigMeanFit);
   
   m_tmpNam = "vggSigStdDevFit";
   m_vggSigStdDevFit.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.vggSigStdDevTruth});
   unsetConst(m_vggSigStdDevFit);

   m_tmpNam = "pdfVggSig";
   m_pdfVggSig.reset(new RooGaussian{m_tmpNam.c_str(), m_tmpNam.c_str()
         , *m_vgg.get(), *m_vggSigMeanFit.get(), *m_vggSigStdDevFit.get()});
   
   // make extended
   m_tmpNam = "pdfExtVggSig";
   m_pdfExtVggSig.reset(new RooExtendPdf{m_tmpNam.c_str()
         , m_tmpNam.c_str(), *m_pdfVggSig.get(), *m_nSigFit.get()});
   return;
}

//
// ModelBuilder::bldVggBkg()
//
void ModelBuilder::bldVggBkg()
{
   cout << "ModelBuilder::bldVggBkg" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   //build
   m_tmpNam = "vggBkgMeanFit";
   m_vggBkgMeanFit.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.vggBkgMeanTruth});
   unsetConst(m_vggBkgMeanFit);
   
   m_tmpNam = "vggBkgStdDevFit";
   m_vggBkgStdDevFit.reset(new RooRealVar{m_tmpNam.c_str(), m_tmpNam.c_str()
         , m_cfg.vggBkgStdDevTruth});
   unsetConst(m_vggBkgStdDevFit);

   m_tmpNam = "pdfVggBkg";
   m_pdfVggBkg.reset(new RooGaussian{m_tmpNam.c_str(), m_tmpNam.c_str()
         , *m_vgg.get(), *m_vggBkgMeanFit.get(), *m_vggBkgStdDevFit.get()});
   
   // make extended
   m_tmpNam = "pdfExtVggBkg";
   m_pdfExtVggBkg.reset(new RooExtendPdf{m_tmpNam.c_str()
         , m_tmpNam.c_str(), *m_pdfVggBkg.get(), *m_nBkgFit.get()});
   return;
}

//
// ModelBuilder::bldVgg
//
void ModelBuilder::bldVgg()
{
   cout << "ModelBuilder::bldVgg" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   m_tmpNam = "pdfVggSum";
   m_pdfVggSum.reset(new RooAddPdf{m_tmpNam.c_str(), m_tmpNam.c_str()
         , RooArgList(*m_pdfVggSig.get(), *m_pdfVggBkg.get())
         , RooArgList(*m_nSigFit.get(), *m_nBkgFit.get())});
   return;
}


//
// ModelBuilder::bld_bkg_sig_full
//
void ModelBuilder::bld_bkg_sig_full()
{
   cout << "ModelBuilder::bld_bkg_sig_full" << endl;
   CheckEventNumsAndObs();
   if(not m_status){return;}
   
   // signal product
   m_tmpNam = "pdfSigProd";
   m_pdfSigProd.reset(new RooProdPdf{m_tmpNam.c_str(), m_tmpNam.c_str()
         , RooArgList(*m_pdfMassSig.get(), *m_pdfVggSig.get()
               )});
   // make extended
   m_tmpNam = "pdfExtSigProd";
   m_pdfExtSigProd.reset(new RooExtendPdf{m_tmpNam.c_str()
         , m_tmpNam.c_str(), *m_pdfSigProd.get(), *m_nSigFit.get()});

   // bkg product
   m_tmpNam = "pdfBkgProd";
   m_pdfBkgProd.reset(new RooProdPdf{m_tmpNam.c_str(), m_tmpNam.c_str()
         , RooArgList(*m_pdfMassBkg.get(), *m_pdfVggBkg.get()
               )});
   // make extended
   m_tmpNam = "pdfExtBkgProd";
   m_pdfExtBkgProd.reset(new RooExtendPdf{m_tmpNam.c_str(), m_tmpNam.c_str()
         , *m_pdfBkgProd.get(), *m_nBkgFit});

   // Full pdf
   m_tmpNam = "pdfFull";
   m_pdfFull.reset(new RooAddPdf{m_tmpNam.c_str(), m_tmpNam.c_str()
         , RooArgList(*m_pdfSigProd.get(), *m_pdfBkgProd.get())
         , RooArgList(*m_nSigFit.get(), *m_nBkgFit.get())});
   return;
}