//Renata Kopecna

#include <genlvlfit.hh>
#include <bu2kstarmumu_pdf.hh>
#include <fitter.hh>
#include <folder.hh>
#include <bu2kstarmumu_plotter.hh>
#include <paths.hh>
#include <event.hh>
#include <helpers.hh>
#include <design.hh>

#include <time.h>

#include <spdlog.h>

int genlvlfit(fcnc::options opts, bool fitPHSP, basic_params params){

    //Gen Lvl always 2017
    params.Run = 2;
    params.year = 2017;


    //yes, it is repetition of mcfit, but it is easier this way

    //--------------------------------
    //      Set constants
    //--------------------------------
    if(fitPHSP) spdlog::info("[MCFIT]\tFit PHSP genLvl."); //reference channel only
    else        spdlog::info("[MCFIT]\tFit signal genLvl MC.");//signal MC only

    const bool Blind  = false; //False as it is MC, HAS TO BE TRUE FOR DATA
    const bool UseBinnedFit = true; //Use bins in q2, if false, nBins is forced to be one
    const bool fixBMass = true; //Fix B mass to PDG value
    const bool plotPulls = true; //Do you wanna pull plots with or without pulls?

    //Scales for the fit parameters (meaning fit-ranges)
    double PprimeRangeScale = 10.0;
    double angleRange = +1.0;
    if (params.usePprime) angleRange *= PprimeRangeScale;
    double angleStepSize  = 0.05;

    if (!UseBinnedFit){
        spdlog::info("Running an unbinned fit");
        opts.TheQ2binsmin = get_TheQ2binsmin(1,false);
        opts.TheQ2binsmax = get_TheQ2binsmax(1,false);
    }
    const unsigned int nBins = opts.get_nQ2bins(); //unsigned because vector.size is unsigned
    spdlog::debug("Using {0:d} q2 bins.", nBins);

    //--------------------------------
    //      Set all options
    //--------------------------------
    opts.use_mkpi = false;
    opts.fit_mkpi = false;
    opts.crystalball = false;
    opts.swave = false;
    opts.weighted_fit = false; //Needed to plot the weights properly
    opts.use_event_norm = false; //Used for convoluting the acceptance into pdf

    opts.only_angles = true; //Do not fit the mass, just angles
    opts.only_Bmass = false; //Do not fit the angles, just mass
    opts.extended_ml = false;
    opts.flat_bkg = false;


    //TODO: possibly add polarity
    //--------------------------------
    //      Load data
    //--------------------------------
    spdlog::info("[MCFIT]\tLoading data...");
    int dataset = 5; //TODO: have a function for this
    if (fitPHSP) dataset = 4;
    std::vector<fcnc::event> events= fcnc::load_events(get_theFCNCpath(dataset,2), "Events", -1);

    UInt_t N_tot = events.size();
    if (N_tot==0){
        spdlog::error("Empty event vector!");
        return 404;
    }
    else spdlog::debug("Total number of used events:\t{0:d}", N_tot);

    //the rest can happily be arrays as nBins and nPDFs are well defined and it is 1D
    std::vector<Int_t> fitresults;

    //Initialize all needed things for the fit
    //current fitter, plotter, parameteres and pdfs:
    fcnc::fitter f(&opts);
    fcnc::folder fldr(&opts);
    fcnc::bu2kstarmumu_plotter * thePlotter;
    std::vector<fcnc::parameters*> theParams [nBins];
    std::vector<fcnc::pdf*> theProbs  [nBins];
    std::vector<std::vector<fcnc::event>*>  selection[nBins];

    //--------------------------------
    //          Read events
    //--------------------------------

    opts.run = params.Run; //Set proper run to options //TODO: check if really needed
    spdlog::debug("Run {0:d}", opts.run);
    opts.name = get_genLvlFit_label(fitPHSP, nBins, -1, params, opts.only_angles, opts.only_Bmass);
    opts.update_angle_ranges(); //Set angles in options back to defaults
    opts.update_efficiencies = false;

    //Set the label to MC
    opts.plot_label = "LHCb genLvl MC";

    //Loop over bins
    for(unsigned int b = 0; b < nBins; b++){
        opts.q2_min = opts.TheQ2binsmin.front();
        opts.q2_max = opts.TheQ2binsmax.back();

        //Create parameter set
        fcnc::bu2kstarmumu_parameters * leParameters = new fcnc::bu2kstarmumu_parameters(&opts);
        //create PDF
        fcnc::bu2kstarmumu_pdf * lePDF = new fcnc::bu2kstarmumu_pdf(&opts, leParameters);

        //Initialize basic parameters
        //define center of q2bin as effective q2
        leParameters->eff_q2.init_fixed( bin_center_q2(opts,b));
        leParameters->f_sig.init_fixed(1.0);

        if (!opts.only_angles){ //If not fitting the angles, do not init the mass parameters
            leParameters->init_mass_parameters(0,nBins,b,0.01);
            leParameters->m_b.init(PDGMASS_B_PLUS, B_MASS_LOW, B_MASS_HIGH, fixBMass ? 0.0: 1.0);
            leParameters->m_scale.init(1.0, 0.0, 2.0, 0.0);
        }


        //Angular observables
        leParameters->init_angular_parameters(nBins,b,angleStepSize, angleRange, false);

        //make sure all configured values are also the start_value:
        leParameters->take_current_as_start();

        theParams[b].push_back(leParameters);
        theProbs [b].push_back(lePDF);
        spdlog::info("Saved PDF and parameters!");

        //create vector with events according to the requested fits/pulls
        std::vector<fcnc::event> *leEvents= new std::vector<fcnc::event>;

        //Loop over events
        spdlog::debug("Loop over events");
        for(auto meas: events){ //WARN when more PDFs per run, this will fail!!!
            //crosscheck between mag up and mag down:
            if(params.polarity==1  && meas.magnet > 0) continue;
            if(params.polarity==-1 && meas.magnet < 0) continue;
            if(meas.q2 < opts.TheQ2binsmin.at(b) || meas.q2 > opts.TheQ2binsmax.at(b))  continue;
            if(meas.mkpi < opts.mkpi_min || meas.mkpi > opts.mkpi_max) continue;
            if(!opts.full_angular) fldr.fold(&meas);

            leEvents->push_back(meas);
        }

        //Update efficiencies  ONCE
        lePDF->update_cached_normalization(leParameters);
        lePDF->update_cached_efficiencies(leParameters, leEvents);

        spdlog::info("Finished selecting the events: {0:d}",leEvents->size());

        //save event vector in vector
        selection[b].push_back(leEvents);
        if ((selection[b]).at(0)->size() > 0){
            spdlog::info("[BIN{0:d}]\tDone!", b);
        }
        else{
            spdlog::critical("No events found for q2-bin={0:d}. Exit!",b);
            assert(0);
        }
    } //End loop over bins

    //Allocate the plotter
    thePlotter = new fcnc::bu2kstarmumu_plotter(&opts);


    //--------------------------------
    //          FIT
    //--------------------------------
    spdlog::info("[MCFIT]\tMC fit started.");

    //Measure the time for the fit:
    runTime timer = runTime();

    //Save the fit results
    std::vector<int> fit_results[nBins];
    std::vector<double> f_sigs[nBins];
    std::vector<double> f_sigserr[nBins];
    std::vector<UInt_t> evts_cntr[nBins];
    std::vector<fcnc::parameters*> vecParams;

    //fit all bins:
    for(unsigned int b = 0; b < nBins; b++){

        //Start the clock
        timer.start();
        time_t startTime = time(0);

        spdlog::info("[START]\tStart the fit for bin #{0:d}", b);

        //Int for fit status
        int fitresult = 0;

        //Delete the texFile
        std::string tag = get_genLvlFit_label(fitPHSP, nBins, b, params,
                                              opts.only_angles, opts.only_Bmass);
        clear_Latex_noteFile(latex_fitterFile(tag));

        spdlog::info("Running the fitter...");
        fitresult = f.fit(theProbs[b].at(0), theParams[b].at(0), selection[b].at(0),tag);
        fit_results[b].push_back(fitresult);


        //Stop the clock
        timer.stop(startTime);

        //Print the fit results
        spdlog::info("[BIN{0:d}]:\tFitresult: {1:d}", b, fitresult);

        //save signal fraction and event number for each bin and each pdf:
        f_sigs[b].push_back( ((fcnc::bu2kstarmumu_parameters *) theParams[b].at(0))->f_sig.get_value());
        f_sigserr[b].push_back( ((fcnc::bu2kstarmumu_parameters *) theParams[b].at(0))->f_sig.get_error());
        evts_cntr[b].push_back((selection[b]).at(0)->size());
        opts.q2_label = q2_label(opts.TheQ2binsmin.at(b), opts.TheQ2binsmax.at(b));
        if (!plotPulls) tag = tag+"_noPulls";


        std::vector<fcnc::bu2kstarmumu_pdf*> * prober = (std::vector<fcnc::bu2kstarmumu_pdf*> *)  & theProbs[b];
        std::vector<fcnc::bu2kstarmumu_parameters*> * paramser = (std::vector<fcnc::bu2kstarmumu_parameters*> *) & theParams[b];
        thePlotter->SetPulls(plotPulls);

        thePlotter->plot_added_pdfs(prober, paramser, &selection[b],
                                       get_genLvlFitPlot_path(fitPHSP),tag, false);

    } //end bin loop


    //--------------------------------
    //        Print & Save
    //--------------------------------

    //Print running time
    timer.print(nBins);

    //Print all fit results
    print_all_parameters(nBins, {0}, theParams, spdlog::level::debug);

    //Save the fit results
    for(unsigned int b = 0; b < nBins; b++){
            std::vector<fcnc::bu2kstarmumu_parameters*> * paramser = (std::vector<fcnc::bu2kstarmumu_parameters*> *) & theParams[b];
                paramser->at(0)->save_param_values(finalResult_genLvlMCfit_txt(fitPHSP, nBins, b, params, opts.only_angles, opts.only_Bmass));
    }

    //Print signal yield in the terminal and to a tex file
    if (!Blind) print_sig_yields(nBins, {0}, evts_cntr, f_sigs, f_sigserr);

    print_sig_yields_tex(get_genLvlFit_label(fitPHSP, nBins, -1, params,
                                             opts.only_angles, opts.only_Bmass),
                         nBins, {0}, &opts, evts_cntr, f_sigs, f_sigserr);


    //Save results to root file

    std::string results_file = final_result_name_genLvlMC(params, nBins, fitPHSP, opts.only_angles, opts.only_Bmass);
    save_results(results_file, nBins, {0}, fit_results, theParams, false, &opts);

    spdlog::info("[MCFIT]\tMC fit finished.");
    return 0;

}
