//Renata Kopecna

#include "mcfit.hh"

#include <event.hh>
#include <fitter.hh>
#include <bu2kstarmumu_plotter.hh>
#include <folder.hh>
#include <paths.hh>
#include <design.hh>
#include <helpers.hh>
#include <spdlog.h>

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

    //--------------------------------
    //      Sanity checks
    //--------------------------------

    if (params.Run != 1 && params.Run != 2 && params.Run != 12){
        spdlog::error("Wrong run selected! Please use either 1, 2 or 12 for fititng the MC!");
        return 1;
    }
    //--------------------------------
    //      Set constants
    //--------------------------------
    if(fitReference) spdlog::info("[MCFIT]\tFit reference channel: J/psi K*+"); //reference channel only
    else             spdlog::info("[MCFIT]\tFit signal 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 SimultaneousFit  = true; //Simultanously fit all PDFs
    const bool fixBMass = false; //Fix B mass to PDG value
    const bool plotPulls = true; //Do you wanna pull plots with or without pulls?

    int prescale = 1; //The fraction of events to be kept
    //Eg if prescale = 3, only a third of events will be kept
    //This is not EXACTLY a third though, but a third from all the possible events,
    //as this is taken from the loaded events, not from the selected events
    //In an extreme case when eg prescale is 2, only MagDown is to be fitted and every second event
    //is MagDown, all events would be still kept.
    //To use all the events, set it to 1

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

    if (!UseBinnedFit){
        spdlog::info("Running an unbinned fit");
        opts.TheQ2binsmin = get_TheQ2binsmin(1,fitReference);
        opts.TheQ2binsmax = get_TheQ2binsmax(1,fitReference);
    }
    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; //5D fit
    opts.fit_mkpi = false; //4D+1 fit

    opts.full_angular = (params.folding ==-1);
    opts.folding = params.folding;
    //TODO: this needs to be fixed, it is written like for idiots
    if(opts.fit_mkpi || opts.use_mkpi){
        opts.only_Bmass = false;
        opts.only_mass2DFit = fitReference;
    }
    bool dontFitAngles = fitReference;

    opts.swave = false;
    opts.multiply_eff = true; //MC weights need to be taken into account
    opts.weighted_fit = true; //Needed to plot the weights properly
    opts.use_event_norm = false; //Used for convoluting the acceptance into pdf

    //Set the plot folder
    opts.plot_folder = get_MCfitPlot_path(fitReference, fitPHSP);

    //Set the available PDFs: this depends on the selected Run for now, so if only one is selected, run over one pdf, if both, use two
    std::vector<UInt_t> pdf_idx;
    if (params.Run == 1 || params.Run == 12) pdf_idx.push_back(1);
    if (params.Run == 2 || params.Run == 12) pdf_idx.push_back(2);

    //--------------------------------
    //      Load data
    //--------------------------------
    spdlog::info("[MCFIT]\tLoading data...");
    std::vector<std::vector<fcnc::event>>events;

    if (opts.run == 1 || opts.run == 12){
        events.push_back(fcnc::load_events(get_theFCNCpath(MC_dataset(fitReference,fitPHSP),1), "Events", -1));
    }
    if (opts.run == 2 || opts.run == 12){
        events.push_back(fcnc::load_events(get_theFCNCpath(MC_dataset(fitReference,fitPHSP),2), "Events", -1));
    }


    //check that the number of pdfs is the same as number of event vectors
    if (pdf_idx.size() != events.size()){
        spdlog::error("Something went very wrong when loading the events and setting the pdfs.");
        spdlog::error("The number of PDFs!= number of event vectors: {0:d} vs {1:d}",pdf_idx.size(), events.size());
        return 5;
    }
    //we are good to go, now; how many individual pdfs are used?
    const UInt_t nPDFs = pdf_idx.size();

    UInt_t N_tot = 0;
    for (UInt_t n = 0; n < nPDFs; n++){
        spdlog::debug("Event vector {0:d}:\t"+(Blind?"Larger 1 ":std::to_string(events.at(n).size())),n);
        if (events.at(n).size()==0){
            spdlog::error("Empty event vector!");
            return 404;
        }
        N_tot += events.at(n).size();
    }
    spdlog::info("Total number of used events:\t{0:d}", N_tot);

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

    //Initialize common parameters //Includes sWave if opts.swave is on
    std::vector<std::string> common_params = param_string(opts, true);
    if(SimultaneousFit) f.set_common_parameters(common_params);
    spdlog::info("Shared parameters: " + convert_vector_to_string(common_params));

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

    //Loop over PDFs and initialize parameters
    for(unsigned n = 0; n < nPDFs; n++){
        opts.run = pdf_idx.at(n); //Set proper run to options
        spdlog::debug("Run {0:d}", opts.run);
        opts.name = get_MCfit_label(fitReference, fitPHSP, nBins, -1, SimultaneousFit, params, opts.run, opts.only_angles, dontFitAngles);
        opts.update_angle_ranges(); //Set angles in options back to defaults
        opts.update_efficiencies = true; //This ensures the acceptance weights to be taken into account.
        //It needs to be set properly as the MC needs to take into account the MC-data weights and hence the weight multiplication is on. However, the weights are recalculated later on (see fitter::fit and fitter::init), so they need to be turned off after loading the efficiencies.

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

        //Loop over bins now
        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 ONLY, init the mass parameters
                leParameters->init_mass_parameters(n,nBins,b,0.01);                
                leParameters->m_b.init(PDGMASS_B_PLUS, B_MASS_LOW, B_MASS_HIGH, fixBMass ? 0.0: 1.0);
                if (!fitReference && !fitPHSP){
                    std::string file = final_result_name_MC(params, 1, true, false, true, false, true);
                    leParameters->fix_param_from_rootfile(file, {"alpha_1","alpha_2","n_1","n_2"}, pdf_idx.at(n), 0);
                }
                leParameters->m_scale.init_fixed(1.0); //Scale of sigma between the signal/ref MC, used in the data mass fit
            }

            if (!dontFitAngles){
                leParameters->init_angular_parameters(nBins,b,angleStepSize,angleRange, Blind);
            }


            //TODO: DELETE: only for testing the folding
            //double FLinit[4] = {0.662,0.689,0.444,0.357};
            //if (params.folding == 4) leParameters->Fl.init_fixed(FLinit[b]);//,0.0,1.0,0.05);
            //if (params.folding == 4) leParameters->S8.init_fixed(0.0);
            if (!fitReference && !fitPHSP){
                std::string file = final_result_name_MC(params, 1, true, false, true, false, true);
                leParameters->fix_param_from_rootfile(file, {"alpha_1","alpha_2","n_1","n_2"}, pdf_idx.at(n), 0);
            }
            //Kpi mass p-wave
            if(opts.fit_mkpi || opts.use_mkpi){
                leParameters->init_mkpi_pWave_parameters(fitReference,0.001);
            }
            //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("[PDF{0:d}]\tSaved PDF and parameters!", pdf_idx.at(n) );

            //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_indexed(auto meas: events.at(n)){
                theOptions[n].update_angle_ranges(); //Update angle ranges based on the folding
                //Reject events in the case prescale is used
                if (prescale>1 && i%prescale!=0) continue;
                //Select either magUp or magDown
                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(meas.m < B_MASS_LOW || meas.m > B_MASS_HIGH) continue;
                if(!filterFldFour(&meas, &theOptions[n])) continue; //remove ctk events
                if(!opts.full_angular) fldr.fold(&meas);
                leEvents->push_back(meas);
            }

            //Update efficiencies (aka take into account angular parametrization weights) ONCE
            lePDF->load_coeffs_eff_phsp_4d();
            lePDF->update_cached_normalization(leParameters);
            lePDF->update_cached_efficiencies(leParameters, leEvents);
            opts.update_angle_ranges();
            spdlog::info("[PDF{0:d}]\tFinished selecting the events: {1:d}",n, leEvents->size());
            //save event vector in vector
            selection[b].push_back(leEvents);
            if(selection[b].back()->size() > 0){
                spdlog::info("[PDF{0:d}]\t[BIN{1:d}]\tDone!", n, b);
            }
            else{
                spdlog::critical("No events found for PDF={0:d} and q2 bin={1:d}. Exit!",n,b);
                assert(0);
            }

        } //End loop over bins
        opts.update_efficiencies = false; //Prevent the weights to be applied several more times in fitter::fit

        theOptions[n] = opts;
        //Allocate the plotter... sigh
        thePlotter[n] = new fcnc::bu2kstarmumu_plotter(&theOptions[n]);

    }//end loop over PDFs

    //--------------------------------
    //          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];

    //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;

        if(!SimultaneousFit){
            for(UInt_t n = 0; n < nPDFs; n++){

                //Delete the texFile
                std::string tag = get_MCfit_label(fitReference, fitPHSP,
                                                  nBins, b,
                                                  SimultaneousFit, params, theOptions[n].run,
                                                  opts.only_angles, dontFitAngles);

                clear_Latex_noteFile(latex_fitterFile(tag));

                spdlog::info("[FIT{0:d}]\tRunning the fitter...", n);
                fitresult = f.fit(theProbs[b].at(n), theParams[b].at(n), selection[b].at(n),tag);
                fit_results[b].push_back(fitresult);
            }
        }
        else{
            spdlog::info("[FIT]\tFitting simultaenously....");
            //Delete the texFile
            std::string tag = get_MCfit_label(fitReference, fitPHSP,
                                              nBins, b,
                                              SimultaneousFit, params, opts.run,
                                              opts.only_angles, dontFitAngles);
            clear_Latex_noteFile(latex_fitterFile(tag));
            for(UInt_t s = 0; s < selection[b].size(); s++)spdlog::info("PDF #{0:d}: {1:d}", pdf_idx.at(s), selection[b].at(s)->size());
            fitresult = f.fit(theProbs[b], theParams[b], selection[b],tag);
            fit_results[b].push_back(fitresult);
            spdlog::info("Q2BIN={0:d}\tLLH={1:f}", b, f.likelihood());
            if(params.likelyhood){
                spdlog::error("TODO.");
                return 5;
            }
            if(params.FeldCous){
                spdlog::error("TODO.");
                return 5;
            }
        }

        //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:
        for(UInt_t n = 0; n < nPDFs; n++){
            f_sigs[b]   .push_back(((fcnc::bu2kstarmumu_parameters *) theParams[b].at(n))->f_sig.get_value());
            f_sigserr[b].push_back(((fcnc::bu2kstarmumu_parameters *) theParams[b].at(n))->f_sig.get_error());
            evts_cntr[b].push_back(selection[b].at(n)->size());
        }

        //Plot everything
        for(UInt_t n = 0; n < nPDFs; n++){
            std::string eps_label = get_MCfit_label(fitReference, fitPHSP,
                                                    nBins, b,
                                                    SimultaneousFit, params, theOptions[n].run,
                                                    opts.only_angles, dontFitAngles);
            theOptions[n].q2_label = q2_label(theOptions[n].TheQ2binsmin.at(b), theOptions[n].TheQ2binsmax.at(b));
            spdlog::info("[PLOT]\t"+eps_label);
            thePlotter[n]->SetPulls(plotPulls);
            thePlotter[n]->plot_data((fcnc::bu2kstarmumu_pdf*)theProbs[b].at(n),
                                     (fcnc::bu2kstarmumu_parameters*)theParams[b].at(n),
                                     selection[b].at(n),
                                     get_MCfitPlot_path(fitReference, fitPHSP), eps_label, false);
        }

        std::string eps_label = get_MCfit_label(fitReference, fitPHSP,
                                                    nBins, b,
                                                    SimultaneousFit, params, params.Run,
                                                    opts.only_angles, dontFitAngles);
        if (!plotPulls) eps_label = eps_label+"_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[0]->SetPulls(plotPulls);
        thePlotter[0]->plot_added_pdfs(prober, paramser, &selection[b],
                                       get_MCfitPlot_path(fitReference, fitPHSP),eps_label, false);

    } //end bin loop


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

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

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


    //Save the fit results
    for(unsigned int b = 0; b < nBins; b++){
        for(UInt_t i = 0; i < pdf_idx.size(); i++){
            std::vector<fcnc::bu2kstarmumu_parameters*> * paramser = (std::vector<fcnc::bu2kstarmumu_parameters*> *) & theParams[b];
                paramser->at(i)->save_param_values(finalResult_MCfit_txt(i,fitReference,fitPHSP, nBins, b, SimultaneousFit, params, params.Run, opts.only_angles, dontFitAngles));

        }
    }

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

    print_sig_yields_tex(get_MCfit_label(fitReference, fitPHSP,
                                         nBins, -1,
                                         SimultaneousFit, params, params.Run,
                                         opts.only_angles, dontFitAngles),
                         nBins, pdf_idx, theOptions, evts_cntr, f_sigs, f_sigserr);

    //Save results to root file
    std::string results_file = final_result_name_MC(params, nBins, fitReference, fitPHSP, SimultaneousFit, opts.only_angles, dontFitAngles);
    save_results(results_file, nBins, pdf_idx, fit_results, theParams, SimultaneousFit, &opts);

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

}
