EWP-BplusToKstMuMu-AngAna/Code/FCNCFitter/bu2kstarmumu.cc

/**
 * @file bu2kstarmumu.cc
 * @author Christoph Langenbruch, David Gerick, Renata Kopecna
 * @date 2020-30-11
 *
 * Main control file for the Bu -> K*+mumu analysis
 */

#include <sys/stat.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cstdio>

//Yes, I know that <> is not ideal for compilation time,
//but I was lazy to always put the correct path there at the begining
//when files were moved around rather often
#include <parameters.hh>
#include <parse.hh>
#include <toystudy.hh>
#include <angularcorr.hh>
#include <massfit.hh>
#include <mainfit.hh>
#include <backgroundfit.hh>
#include <mcfit.hh>
#include <momfit.hh>
#include <genlvlfit.hh>
#include <toysfit.hh>
#include <help.hh>
#include <constants.hh>
#include <generatetoys.hh>
#include <helpers.hh>
#include <event.hh>
#include <design.hh>
#include <paths.hh>

#include "sources/Scripts/RunningScripts.hh"

#include <TROOT.h>
#include <colors.hh>


int main ( int argc, char *argv[] )
{
    //Set gROOT and gStyle
    gROOT->SetStyle("Plain");
    set_gStyle();
    //Initiliaze pretty colors
    myColorScheme::init();
    //gStyle->SetPaintTextFormat("4.1f m");
    //gStyle->SetCanvasPreferGL(true);

    ///set the fitting options
    fcnc::options opts;
    opts = fcnc::options("-1");    

    //how many events should be taken per Test/ToyFit?
    UInt_t nMCEvents   = 350;


    //Retun errors:
    //1: wrong options
    //5: Inexplicable stuff
    //404: not found

    basic_params pars = basic_params();
    basic_actions acts = basic_actions();
    //The pars and acts are a bit wonky at the moment, hopefully as I progress this will get better
    if (parseOpts(argc, argv, pars, acts, nMCEvents) == 1) return 1;
    if (checkOpts(pars,acts) == 1) return 1;

    //Set opts.folding according to the parsed options
    //default should be -1
    opts.folding = pars.folding;

    ////////////////////////////////////////////
    // [START]    GLOBAL OPTIONS
    ////////////////////////////////////////////

    //Use only DTF, could be removed
    opts.DTF = true;

    //SET OUTPUT DESIGN
    reset_spdlog();
    set_spdlog_level(pars.verbosity);

    //Make ROOT shut up
    gStyle -> SetOptStat(0);
    gROOT->SetBatch(kTRUE);
    gErrorIgnoreLevel = kWarning;
    //gErrorIgnoreLevel = kError;       

    opts.ncores = NCORES;
    if(opts.ncores > 1) spdlog::info("[CPU]\tUse parallelisation on {0:d} CPU cores", opts.ncores);

    //FIXED to true, since we have flatQ2 phase-space MC
    opts.IsFlatQ2          = true;

    //get individual angular acceptance corrections for every year
    opts.angacccorrperyear = false;
    //Decide in what format everything should be saved in
    opts.write_eps = true;
    opts.write_C   = false;
    opts.write_jpg = false;
    opts.write_pdf = true;

    //Set the default path for plots_folder, somehow it doesn't want to accept PLOTS_PATH in the options constructor, possibly fix TODO
    opts.plot_folder = PLOTS_PATH;

    //fit pdfs over complete phi range [-pi, +pi]
    opts.full_angular      = opts.folding == -1;
    //use a two tailed CrystalBall function instead of a Gaussian
    opts.twotailedcrystalball = DOUBLE_CB;
    opts.crystalball = !DOUBLE_CB;
    //No exponential, but flat background model
    opts.flat_bkg          = false;

    //////////////////////////
    // HARDCODED FUNCTIONS: //
    //////////////////////////

    //default is NOT Kshort channel:
    opts.KS = false;

    //Set 'globally' fl and afb to (not) be used
    opts.fit_fl = true;
    opts.fit_afb = true;

    //Set 'globaly' to use lambda in the background fit
    //fit exponential with exp(-lambda*x) instead of exp(-x/tau), it has better precision
    opts.fit_lambda = true;

    //Assymetries to be included in the fit itself
    opts.fit_asymmetries = false; //you don't have enough stats for this to be true

    /////////////////////////////////////
    // Set the order of the polynomial //
    /////////////////////////////////////
    opts.bkg_order_costhetal = 2; //TODO: put as a const into constants.cc and modify it accordingly in options.cc
    opts.bkg_order_costhetak = 5;

    ////////////////////////
    // ANGULAR RESOLUTION //
    ////////////////////////

    //Set the order of legendre polynomial
    opts.eff_order_costhetal = ORDER_COSTHETAL;
    opts.eff_order_costhetak = ORDER_COSTHETAK;
    opts.eff_order_phi       = ORDER_PHI;
    opts.eff_order_q2        = ORDER_Q2;

    std::vector<std::vector<double>> resolution = get_resolution();

    //Have angular correction per each Run (false)
    opts.angacccorrbothruns = false;

    //Don't use angular acceptance convoluted in the likelihood (false)
    opts.use_angular_acc = false;

    ////////////////////////////////////////////
    //             GLOBAL OPTIONS             //
    ////////////////////////////////////////////

    //TODO fix this (actually nice, but can be done elsewhere)
    //maximum range of considered B mass:
    opts.m_low    = B_MASS_LOW;
    opts.m_high   = B_MASS_HIGH;
    //signal window of B mass:
    opts.m_min    = PDGMASS_B-B_MASS_TIGHT_WINDOW;
    opts.m_max    = PDGMASS_B+B_MASS_TIGHT_WINDOW;

    //If MC, use MC weights//TODO: check for data
    opts.multiply_eff = pars.dataset >0;

    //TODO move to paths
    std::string angularsuffix = opts.full_angular ? "full_angular" : "folding"+std::to_string(opts.folding);


    //generate vector with years for every run option, including just one year
    std::vector<Int_t>years = (pars.Run == 0 ? std::vector<Int_t>{pars.year} : get_years(pars.Run, pars.dataset==1, pars.dataset==2));
    spdlog::debug("Using years: " + convert_vector_to_string(years));

    //---------------------------------------------
    //    Put stuff from pars and acts into opts
    //---------------------------------------------
    //yes it is dumb but I cannot be bothered to change the whole options
    opts.verbose = pars.verbosity;
    opts.run = pars.Run; //This is not ideal when fitting per year,
                         //From now on passing the pars as well
    opts.useMC = pars.dataset > 0;
    opts.fit_pprimes = pars.usePprime;

    //-------------------------------
    //    Get corresponding Q2 bins
    //-------------------------------
    if (pars.reference) pars.nBins = 1; //For reference, always use 1 bin
    opts.TheQ2binsmin = get_TheQ2binsmin(pars.nBins, pars.reference);
    opts.TheQ2binsmax = get_TheQ2binsmax(pars.nBins, pars.reference);

    spdlog::debug("Using {0:d} q2 bins.", opts.TheQ2binsmax.size());

    if (spdlog::default_logger_raw()->level()==spdlog::level::info){
        spdlog::info("Using the following binning scheme:");
        for(UInt_t b = 0; b < opts.TheQ2binsmin.size(); b++){
            assert(opts.TheQ2binsmin.at(b) != 0.0);
            assert(opts.TheQ2binsmax.at(b) != 0.0);
            assert(opts.TheQ2binsmin.at(b) < opts.TheQ2binsmax.at(b));
            std::cout << std::fixed << std::setprecision(2) << "\tbin" << b+1 << ": [" << opts.TheQ2binsmin.at(b) << "-" << opts.TheQ2binsmax.at(b) << "] GeV^2/c^4" << std::endl;
        }
    }

    //-------------------------------
    //    Set ranges of the angles
    //-------------------------------
    //Not ideal to have it as an opts and as a constant, but this is how it is going to be now
    opts.reset_angle_ranges();
    //The ranges might change in the folding, hence I will keep the opts.angle_min/max for this case


    //Check if something is wrong; if any task is assigned, raise this flag
    bool isItDoingAnything_flag = false;
    int returnCode = 0;

    //-------------------------------
    //     create branches for FCNCfitter
    //-------------------------------
    if(acts.convert){ //these booleans are only for running, used only here
        //maybe put it into a new class //TODO
        isItDoingAnything_flag = true;
        return convert_tuples(pars.dataset, get_theFCNCpath(pars.dataset, opts.run), opts,years);
    }

    //-------------------------------
    //     scan angular acceptance max. order of legendre
    //-------------------------------
    if(acts.angCorrScan){
        isItDoingAnything_flag = true;
        bool quickTest = false;
        bool test_4times1D = false;
        bool checkSignificance = false;
        bool runMinuit = false;
        bool checkFactorization = false;
        bool do3Dmoments = false;
        const bool assumePhiEven  = IS_PHI_EVEN;
        spdlog::debug("Scanning angular acceptance corrections...");
        returnCode =  scan_max_order_angular_ccorrection(opts, assumePhiEven, quickTest, test_4times1D,
                                                  checkSignificance, runMinuit, checkFactorization, do3Dmoments);
        if (returnCode!=0) return returnCode;
    }

    //-------------------------------
    //     fit angular acceptance
    //-------------------------------
    if (acts.angCorr){
        if (pars.testInt > -1) sanityCheck_MC(opts);
        else{
            isItDoingAnything_flag = true;
            bool testSaving = false;
            spdlog::debug("Calculating angular acceptance correction parameters...");
            returnCode = get_angular_acceptance(opts, testSaving);
            if (returnCode!=0) return returnCode;
        }
    }

    //--------------------------------
    //     FIT MASS
    //--------------------------------
    if(acts.fitType == 2){
        isItDoingAnything_flag = true;
        bool splitRuns = true;
        massfit(opts, pars.reference, splitRuns, pars);
    }

    //--------------------------------
    //     MC FIT
    //--------------------------------
    if(acts.fitType == 4){
        isItDoingAnything_flag = true;
        opts.only_angles = false;//(nBins == 8); //Do not fit the mass, just angles //Apply only when fitting B0 binning
        //Reference MC is broken (K* is PHSP), so when fitting the reference, fit only mass
        opts.only_Bmass  = pars.reference; //Do not fit the angles, just mass
        opts.initSM     = !pars.reference; //Set to true only for reference
        returnCode = mcfit_4D(opts, pars.reference, pars.dataset==3, pars);
        //yes, fitRef could be included in params, but I like to keep my options open :)
        if (returnCode!=0) return returnCode;
    }

    //--------------------------------
    //     COMPLETE FIT
    //--------------------------------
    if(acts.fitType == 1){
        isItDoingAnything_flag = true;
        bool fitToy = pars.testInt > -1; //If test, then do fit toys
        if (pars.folding == 5){ //We do not need this here, we can run the code 5 times, no?
            for (int f = 0; f < 5; f++){
                pars.folding = f;
                opts.folding = f;
                spdlog::debug("Fitting with folding={0:d}", pars.folding);
                mainfit(opts, pars, pars.reference, fitToy, pars.likelyhood, pars.FeldCous);
            }
        }
        else {
            mainfit(opts, pars, pars.reference, fitToy, pars.likelyhood, pars.FeldCous);
        }
    }

    //--------------------------------
    //     BACKGROUND ONLY FIT
    //--------------------------------
    if(acts.fitType == 0){
        isItDoingAnything_flag = true;

        //Clear the bkgFit file
        clear_Latex_noteFile("_bkgFit");

        //Set options
        //One can also fit only upper/lower mass sideband, to set this use -index
        bool LowMassFit = (pars.index == 1);
        bool HighMassFit = (pars.index == 2);
        bool Use2DAngularBins = false;
        bool fitKpiMass = true;

        //Run the background only fit
        int bkgfit = 0;
        if (pars.folding == 5){ //We do not need this here, we can run the code 5 times, no?
            for (int f = 0; f < 5; f++){
                pars.folding = f;
                opts.folding = f;
                bkgfit =+ backgroundfit(opts, pars.reference,
                                    LowMassFit, HighMassFit, fitKpiMass,
                                  Use2DAngularBins, pars);
            }
        }
        else bkgfit = backgroundfit(opts, pars.reference,
                                    LowMassFit, HighMassFit, fitKpiMass,
                                    Use2DAngularBins, pars);
        if (bkgfit != 0) return bkgfit;
    }

    //--------------------------------
    //     start GENERATOR LEVEL MC FIT
    //--------------------------------
    if(acts.fitType == 5){
        isItDoingAnything_flag = true;
        opts.initSM = pars.nBins != 8; //SM for 8 bins not implemented
        genlvlfit(opts, pars.dataset==4, pars);
    }

    //--------------------------------
    //     method of moment fit
    //--------------------------------
    //TODO
    if(acts.fitType == 3){
        isItDoingAnything_flag = true;
        bool blind = true;
        bool Fit2bins = false, Fit1bin = false, FitAllbins = false; //TODO: remove
        momfit(opts, pars.nBins==0, pars.dataset==1, pars.dataset==0, blind,
                Fit1bin, Fit2bins, FitAllbins);
    }

    //--------------------------------
    //     BEGIN angular resolution
    //--------------------------------

    if(acts.angRes){
        isItDoingAnything_flag = true;       
        get_angular_resolution(opts, get_years(opts.run, true, false));
    }

    //-------------------------------
    //     generate Toy events
    //-------------------------------

    if (acts.genToys){
        isItDoingAnything_flag = true;
        spdlog::debug("Generating and fitting toys");

        bool onlySig = true;
        bool onlyBkg = false;
        if (pars.folding == 5){ //We do not need this here, we can run the code 5 times, no?
            for (int f = 0; f < 5; f++){
                pars.folding = f;
                opts.folding = f;
                spdlog::debug("Generating and fitting a toy sample with folding={0:d}", pars.folding);
                toysfit(opts, pars.nBins, pars.reference, pars, onlySig, onlyBkg);
            }
        }        
        else {
            if (pars.testInt > -1)saveToys(pars, opts, pars.dataset>0, false); //if MC generate without bkg
            else{
                if(pars.folding != -1)spdlog::debug("Generating and fitting a toy sample with folding={0:d}", pars.folding);
                toysfit(opts, pars.nBins, pars.reference, pars, onlySig, onlyBkg);
            }
        }
    }

    //-------------------------------------------
    //     Run whatever script you decide to run
    //-------------------------------------------
    //Yes, this is tedious, but I don't want to hardcode everything just so root is happy
    //nor I wanna compile everything, so here we are
    if(acts.script){
        isItDoingAnything_flag = true;
        int status = runWhatever(pars, acts);       //FILL IN RunningScripts.cc BY HAND
        if(status)return status;
    }


    if (!isItDoingAnything_flag) spdlog::warn("Nothing to do! Exitting.");
    else                         spdlog::info("Done with everything. Exit.");
    return 0;
}