// ROOT headers
#include <TROOT.h>
// User defined headers
#include "parser.h"
#include "pdfMaker.h"
#include "toolsStat.h"
#include "gof.h"
#include "deltaM.h"


// Max number of steps in amplitude
const Int_t aMax = 1;
// Max number of steps in phase
const Int_t pMax = 1;
// Amplitude step size
const Double_t aStep = 0.05;
// Phase step size
const Double_t pStep = 3.0;
// Global asymmetry [alpha=(1.0+gAsym)/2.0]
Double_t gAsym = 0;
//// Number of events to generate to obtain the distribution
const Int_t kMax = 100;

/////////////////////////////////////////////////////////////////////////////////////////
// MAIN
int main(int argc, char *argv[]){
  // Enable thread safety for ROOT
  ROOT::EnableThreadSafety();
  // Call the parser and set the required options
  Parser app(argc, argv);
  std::vector<std::string> resonance, channel, parameter;
  Int_t numThreads;
  //// Set options
  app.addOption(resonance, CmdlOptions::resonance);
  app.addOption(channel  , CmdlOptions::channel);
  //// Set flags
  app.addFlag(&CmdlOptions::distribution);
  //// Set variables
  app.addVariable(numThreads, CmdlOptions::thread);
  //// Set default values
  app.defOption(parameter, CmdlOptions::parameter);
  //// Now run everything
  app.run();
  // << "-------------------------------------------------------------------------" dummy;
  std::cout << "INFO:\tStudy of the parameters (a,\u03C6) correlation" << std::endl;
  // Initialize some variables
  std::string outPath, outName;
  std::string dirName;
  // Start the macro
  try {
    // Set the status
    StatusAndLog& status = StatusAndLog::getInstance();
    std::cout << "-------------------------------------------------------------------------" << std::endl;
    // Start producing the distributions
    std::string title = Tools::padRight("  GENERATING t_\u03B8 DISTRIBUTION", 59);
    status.header(title);
    // Skip the production of the t-Distribution if the 'distribution' flag has not been called
    if(CmdlOptions::distribution.value){
      // Loop on resonances and decay channels
      for(auto const &rs: resonance){
        for(auto const &ch: channel){
          status.info(std::string(rs+" "+CmdlOptions::channelToAlias(ch)));
          // Set the output File where the p(t;θ,...) will be saved
          outPath = ToolsDir::outFilePath();
          outName = ToolsDir::macroName+"_"+Tools::deleteChar(rs)+"_"+Tools::upperCase(CmdlOptions::channelToAlias(ch));
          TFile *outFile = ToolsROOT::openRootOutFile(outPath, outName, "RECREATE");
          // Set all the common plots
          status.update("Setting up the system");
          // Set the TTree to save
          std::map<std::string, TTree*> mapTDistr;
          //// Set the t-statistic variables
          Double_t t_a, t_p;
          // Set up the needed TTree to fill
          for(Int_t a=0; a<=aMax; a++){
            for(Int_t p=0; p<=pMax; p++){
              std::string tName = Tools::namePDF("tDistr", ch, rs, aStep*a, pStep*p);
              outFile->cd();
              TTree *tDistr = new TTree((tName).c_str(),"");
              mapTDistr.insert(std::pair<std::string, TTree*>(tName, tDistr));
              tDistr->Branch("t_a", &t_a ,"t_a/D");
              tDistr->Branch("t_p", &t_p ,"t_p/D");
            }
          }
          // Fit Δm and get the β value (β=S/(S+B)) to obtain the lambda value
          status.info("Fitting \u0394m distribution to extract \u03B2");
          // DeltaM fit_deltaM(ch, "Q");
          // Double_t beta = fit_deltaM.beta;
          Double_t beta = 1;
          status.info("\u03B2 = "+Tools::toStringPrec(beta, 3));
          //// Create the efficiency, the f(x) and the g_θ(x) plots, and the mask with the regions
          PDFMaker pdf;
          TH2D *effRun2 = pdf.getEff_Run2();
          TH2D *pdf_f   = pdf.getFPDF(effRun2, ch, rs, "");
          TH2D *gFunc_a = pdf.getGFunc(effRun2, ch, rs, parameter[0], "");
          TH2D *gFunc_p = pdf.getGFunc(effRun2, ch, rs, parameter[1], "");
          TH2D *pdfMask = pdf.getMask(ch);
          TH2D *regMask = ToolsROOT::makeRegionMask(effRun2, rs);
          // Get the number of events
          Int_t nEvents = ToolsData::getNumEvents(ch, CmdlOptions::source.defaultValues[0], {pdfMask, regMask}, "");
          // Get the data
          TH2D *data = ToolsData::getData(ch, CmdlOptions::source.defaultValues[0]);
          ToolsROOT::applyMask(data, pdfMask);
          ToolsROOT::applyMask(data, regMask);
          // Get the DeltaM sideband to be used as background PDF
          TH2D *pdf_b = ToolsData::getData(ch, CmdlOptions::source.defaultValues[0], "_sideband");
          ToolsROOT::applyMask(pdf_b, pdfMask);
          ToolsROOT::applyMask(pdf_b, regMask);
          pdf_b->Scale(1/pdf_b->Integral());
          //// Write the plots on the outFile
          outFile->cd();
          effRun2->Write();
          pdf_f->Write();
          gFunc_a->Write();
          gFunc_p->Write();
          pdfMask->Write();
          regMask->Write();
          data->Write();
          pdf_b->Write();
          // Iterate over amp and phi
          for(Int_t a=0; a<=aMax; a++){
            for(Int_t p=0; p<=pMax; p++){
              // Select the TTree
              std::string tName = Tools::namePDF("tDistr", ch, rs, aStep*a, pStep*p);
              TTree *tDistr = mapTDistr[tName];
              // Get the PDF to generate the data
              TH2D *pdf_p = pdf.getPDF(effRun2, ch, rs, aStep, pStep, a, p, "pdf_p");
              TH2D *pdf_n = pdf.getPDF(effRun2, ch, rs, aStep, pStep,-a,-p, "pdf_n");
              //// Apply the region mask
              ToolsROOT::applyMask(pdf_p, regMask);
              ToolsROOT::applyMask(pdf_n, regMask);
              pdf_p->Scale(1.0/pdf_p->Integral(),"nosw2");
              pdf_n->Scale(1.0/pdf_n->Integral(),"nosw2");
              // Generate the distribution
              status.info("PRODUCING: a = "+std::to_string(a)+"  p = "+std::to_string(p));
              TRandom3 *rand = new TRandom3();
              omp_set_num_threads(numThreads);
              ProgressBar bar(kMax, 50, numThreads);
              #pragma omp parallel for
              for(Int_t k=0; k<kMax; k++){
                // Print the loading bar
                bar.printBar();
                // Generate Dalitz plots based on the PDFs (copy the axis limits from the PDFs)
                // The number of events in the Dalitz plots is extracted from a Binomial distribution
                TH2D *randHisto_p = ToolsROOT::histCopyAxis(pdf_p, "randHisto_p_"+std::to_string(omp_get_thread_num()));
                TH2D *randHisto_n = ToolsROOT::histCopyAxis(pdf_n, "randHisto_n_"+std::to_string(omp_get_thread_num()));
                Int_t Np = rand->Binomial(nEvents,(1+gAsym)/2.0);
                Int_t Nn = nEvents-Np;
                randHisto_p->FillRandom(pdf_p, Np);
                randHisto_n->FillRandom(pdf_n, Nn);
                // Evaluate the statistic vector {t_a, t_p}
                //// First get the coefficient matrix A
                std::vector<TH2D *> gFuncVec = {gFunc_a, gFunc_p};
                TMatrixD *Ainv_Mx = ToolsStat::A_Mx_calculator(pdf_f, gFuncVec);
                //// Then calculate t
                std::vector<Double_t> tVec = ToolsStat::t_calculator(randHisto_p, randHisto_n, pdf_f, gFuncVec, nEvents, Ainv_Mx);
                // Save the t value in the TTree branch
                // Use omp critical to avoid race condition when filling the TTree  
                // The values of the variables must also be assigned at the end to avoid race condition
                #pragma omp critical
                {
                  t_a = tVec[0];
                  t_p = tVec[1];
                  tDistr->Fill();
                }
                // Close
                Ainv_Mx->Delete();
                randHisto_p->Delete();
                randHisto_n->Delete();
              }
              // Wait for all threads to reach this point before proceeding
              #pragma omp barrier
              rand->Delete();
              pdf_p->Delete();
              pdf_n->Delete();
            }
          }
          // Write the TTrees
          outFile->cd();
          for (const auto& pair : mapTDistr) {
            pair.second->Write();
            pair.second->Delete();
          }
          // Close
          data->Delete();
          pdf_b->Delete();
          effRun2->Delete();
          pdf_f->Delete();
          gFunc_a->Delete();
          gFunc_p->Delete();
          pdfMask->Delete();
          regMask->Delete();
          outFile->Close();
          outFile->Delete();
        }
      }
    }
    else {
      // If the 'distribution' flag has not been called, the creation of the distribution phase is skipped
      status.info("Skipped");
    }
    // Start plotting the distribution
    title = Tools::padRight("  PLOTTING", 59);
    status.header(title);
    status.info("WORK IN PROGRESS");
  }
  catch(std::string error){
    app.error = true;
    CmdlOptions::errorCatch(error);
    return 1;
  }
  return 0;
}
// END MAIN
/////////////////////////////////////////////////////////////////////////////////////////