//Renata Kopecna

#include <massfit.hh>

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

#include <fstream>

#include <spdlog.h>

#include <TROOT.h>
#include <TStyle.h>

//Fitter part that fits only the mass of DATA

bool checkForCommonValues(std::vector<std::string> vec_one,std::vector<std::string> vec_two){
    for (auto key: vec_one){
        if (std::find(vec_two.begin(), vec_two.end(), key)!= vec_two.end()){
            spdlog::warn("Parameter " + key + " is shared between two vectors!");
            return true;
        }
    }
    return false;
}

int massfit(fcnc::options opts,  bool fitReference, bool splitRuns, basic_params pars){
    //splitRuns = true means one gets a simFit of Run 1 and Run 2
    //splitRuns = false means Run 1 and Run 2 are fit together

    const bool fitKstarMass = true;

    bool Blind = false;//!fitReference; //If signal channel, blind to be sure
    const unsigned int nBins = opts.get_nQ2bins(); //unsigned because vector.size is unsigned

    bool constrainMassFromMC = false;  //Use mass fit of MC to constrain mass parameters?
    bool constrainMassFromRefMC = true;//Use mass fit of *RefMC* to constrain mass parameters?
    bool constrainMassFromData   = !fitReference; //Use the mass fit of data Reference to constrain the parameters?

    bool fixBMass = !fitReference; //Fix the B mass to PDG?

    //Set the parameters that should be constrained
    std::vector<std::string> toBeConstrainedMC = {};
    std::vector<std::string> toBeFixedMC = {};

    std::vector<std::string> toBeConstrainedRefMC = {};
    std::vector<std::string> toBeFixedRefMC = {"alpha_1","alpha_2","n_1","n_2"};

    std::vector<std::string> toBeConstrainedRefData = {};
    std::vector<std::string> toBeFixedRefData = {"m_b","m_sigma_1"};


    //Sanity check one doesn't constrain one thing to two different values
    if (constrainMassFromMC && constrainMassFromRefMC){
        if (checkForCommonValues(toBeConstrainedMC, toBeConstrainedRefMC)){
            spdlog::warn("Fix your options, aborting now!");
            return 2;
        }
        if (checkForCommonValues(toBeFixedMC, toBeFixedRefMC)){
            spdlog::warn("Fix your options, aborting now!");
            return 2;
        }
    }
    if (constrainMassFromMC && constrainMassFromData){
        if (checkForCommonValues(toBeConstrainedMC, toBeConstrainedRefData)){
            spdlog::warn("Fix your options, aborting now!");
            return 2;
        }
        if (checkForCommonValues(toBeFixedMC, toBeFixedRefData)){
            spdlog::warn("Fix your options, aborting now!");
            return 2;
        }
    }
    if (constrainMassFromRefMC && constrainMassFromData){
        if (checkForCommonValues(toBeConstrainedRefMC, toBeConstrainedRefData) ){
            spdlog::warn("Fix your options, aborting now!");
            return 2;
        }
        if (checkForCommonValues(toBeFixedRefMC, toBeFixedRefData)){
            spdlog::warn("Fix your options, aborting now!");
            return 2;
        }
    }

    if (constrainMassFromMC && checkForCommonValues(toBeConstrainedMC, toBeFixedMC)){
        spdlog::warn("The parameter will be fixed!");
    }
    if (constrainMassFromRefMC && checkForCommonValues(toBeConstrainedRefMC, toBeFixedRefMC)){
        spdlog::warn("The parameter will be fixed!");
    }
    if (constrainMassFromData && checkForCommonValues(toBeConstrainedRefData, toBeFixedRefData)){
        spdlog::warn("The parameter will be fixed!");
    }

    gROOT->SetBatch(kTRUE);
    gROOT->SetStyle("Plain");
    set_gStyle();

    //Open texFile
    std::ofstream myFile;
    open_Latex_noteFile(latex_massFit(), myFile);

    spdlog::info("[FIT]\t\tFitting only mass...");

    // --- Set the options --- //   
    opts.only_Bmass = !fitKstarMass;
    opts.only_mass2DFit = fitKstarMass;

    opts.swave = fitKstarMass;
    opts.shift_lh = false;    //keep it, it was there before
    opts.weighted_fit = fitKstarMass; // weights for the MASS DATA fit

    opts.plot_chi2 = true;
    opts.fit_mkpi = fitKstarMass;
    opts.use_mkpi = false;
    opts.plot_folder = get_MassFitPlot_path(); //Set the plot path
    opts.minos_errors = true; //TODO!
    opts.squared_hesse = false;
    opts.asymptotic = false;

    if (!fitReference) opts.plots_m_bins = 30; //Plot less bins for signal channel as the stats are sad there
    if (!fitReference) opts.plots_mkpi_bins = 20; //Plot less bins for signal channel as the stats are sad there

    // --- Load events from data tuple --- //
    std::vector<std::vector<fcnc::event>>events;
    std::vector<UInt_t> pdf_idx;

    if (opts.run == 1 || opts.run == 12){
        std::vector<fcnc::event> tmp = fcnc::load_events(get_theFCNCpath(0,1), "Events", -1);
        if (!fitReference) events.push_back(fcnc::filterResonances(tmp));
        else               events.push_back(tmp);
        pdf_idx.push_back(1);
    }
    if (opts.run == 2 || opts.run == 12){
        std::vector<fcnc::event> tmp = fcnc::load_events(get_theFCNCpath(0,2), "Events", -1);
        if (!fitReference) events.push_back(fcnc::filterResonances(tmp));
        else               events.push_back(tmp);
        pdf_idx.push_back(2);
    }
    //No need to keep the datasets separated, as there is no need for the weights
    if (!splitRuns && opts.run == 12){
        pdf_idx.clear();
        pdf_idx.push_back(12);
        //Add the vector from Run 2 into Run 1 vector
        events[0].insert(events[0].end(), events[1].begin(), events[1].end());
        //Delete the Run 2 vector
        events.pop_back();
    }

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

    //Check the number of events
    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);

    //Get the path to the file where to get the values/constraints from
    //Always use simFit from MCand mass only, it is easier
    std::string ParamInitFileMC = final_result_name_MC(pars, nBins, fitReference, false, true, false, false);
    std::string ParamInitFileRefMC = final_result_name_MC(pars, 1, true, false, true, false, true);
    std::string ParamInitFileRefData = final_result_name_mass(true, 1, splitRuns, pars, pars.Run);

    //current fitter, plotter, parameter and pdf:
    fcnc::fitter f(&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];

    //Set FS as a common parameter if fitting the sWeight
    if(fitKstarMass) f.set_common_parameters({"FS"});

    //Loop over PDFs and initialize parameters
    for(unsigned n = 0; n < nPDFs; n++){
        opts.run = pdf_idx.at(n); //Set proper run to options
        opts.name = get_MassFit_label(fitReference, -1, nBins, splitRuns, pars, pdf_idx.at(n));
        spdlog::debug("Initializing run {0:d}", opts.run);
        opts.update_efficiencies = true; //This ensures the acceptance weights to be taken into account.
        //loop over number of bins
        for(UInt_t b = 0; b < nBins; b++){
            //Create parameter set
            fcnc::bu2kstarmumu_parameters * leParameters = new fcnc::bu2kstarmumu_parameters(&opts);
            //create PDF
            fcnc::bu2kstarmumu_pdf * lePDF = new fcnc::bu2kstarmumu_pdf(&opts, leParameters);

            //Init the sig/bkg fraction
            leParameters->f_sig.init(fitReference ? 0.8 : 0.3, 0.1, 1.0, 0.1);

            //Init PDG mass, keep the range this large!
            leParameters->m_b.init(PDGMASS_B, B_MASS_LOW, B_MASS_HIGH, fixBMass ? 0.0 : 0.1);

            //Init mass parameters to defaults
            leParameters->init_mass_parameters(n,nBins,b,0.01);

            //Overwrite the defaults with constraints/fixes from MC/Ref, if needed
            if (constrainMassFromMC){
                leParameters->constrain_param_from_rootfile(ParamInitFileMC, toBeConstrainedMC,  splitRuns ? pdf_idx.at(n) : 2, b);
                leParameters->fix_param_from_rootfile(ParamInitFileMC, toBeFixedMC,  splitRuns ? pdf_idx.at(n) : 2, b);
            }
            if (constrainMassFromRefMC){ //Fix to refrence channel, so always read 0th bin
                leParameters->constrain_param_from_rootfile(ParamInitFileRefMC, toBeConstrainedRefMC,  splitRuns ? pdf_idx.at(n) : 2, 0);
                leParameters->fix_param_from_rootfile(ParamInitFileRefMC, toBeFixedRefMC,  splitRuns ? pdf_idx.at(n) : 2, 0);
            }
            if (constrainMassFromData){ //Fix to refrence channel, so always read 0th bin
                leParameters->constrain_param_from_rootfile(ParamInitFileRefData, toBeConstrainedRefData,  splitRuns ? pdf_idx.at(n) : 2, 0);
                leParameters->fix_param_from_rootfile(ParamInitFileRefData, toBeFixedRefData,  splitRuns ? pdf_idx.at(n) : 2, 0);
            }                                    

            //Make sure the B mass range is correct
            leParameters->m_b.set_min(B_MASS_LOW);
            leParameters->m_b.set_max(B_MASS_HIGH);


            //Init the ratio of sigmas in signal MC/reference MC
            if (!fitReference) leParameters->m_sigma_1.init_fixed(leParameters->m_sigma_1.get_value()*get_sigmaRatio_fromMC(pars,nBins,b,pdf_idx.at(n)));
            //I don't trust the m_scale


            //Init background
            leParameters->init_mass_background_parameters(nBins, b, opts.fit_lambda);

            //define center of q2bin as effective q2:
            leParameters->eff_q2.init_fixed(bin_center_q2(opts,b));

            //Init the Kstar mass fit
            if(opts.fit_mkpi || opts.use_mkpi){
                //p-wave;
                leParameters->init_mkpi_pWave_parameters(fitReference,0.001);
                leParameters->init_kpi_background_parameters(fitReference,0.05);

                //s-wave
                leParameters->FS.init(0.055, 0.00, 0.5, 0.001); //David got 0.055
                leParameters->init_mkpi_sWave_parameters(fitReference, 0.0); //Fix it, otherwise it doesn't converge
            }

            //Control print of the parameters
            //leParameters->print_parameters(false);

            //Save the parameters and PDF per bin
            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 (auto meas: events.at(n)){
                if(meas.m < opts.m_low || meas.m > opts.m_high) continue;
                if(pars.polarity==1  && meas.magnet > 0) continue;
                if(pars.polarity==-1 && meas.magnet < 0) continue;
                if(meas.q2 < opts.TheQ2binsmin.at(b) || meas.q2 > opts.TheQ2binsmax.at(b))  continue;
                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);

            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 #ItIsn'tThe90sAnymore
        thePlotter[n] = new fcnc::bu2kstarmumu_plotter(&theOptions[n]);
    }//end loop over PDFs

    //--------------------------------
    //          FIT
    //--------------------------------
    spdlog::info("[MASSFIT]\tMass 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<double> bkg_int_full_range[nBins]; //Not really used
    std::vector<UInt_t>evts_cntr[nBins];

    //fit per 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);

        std::string tag = get_MassFit_label(fitReference, b,nBins, splitRuns, pars, 0);
        clear_Latex_noteFile(latex_fitterFile(tag));
        spdlog::info("[FIT]\tRunning the fitter...");
        fit_results[b].push_back(f.fit(theProbs[b], theParams[b], selection[b], tag));
        spdlog::info("Q2BIN={0:d}\tLLH={1:f}", b, f.likelihood());

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

        //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());
            //integrate the background pdf over the range [opts.m_min; opts.m_max] to estimate the background contribution in each q2bin:
            bkg_int_full_range[b].push_back(((fcnc::bu2kstarmumu_pdf*)theProbs[b].at(n))->integral_m_bkg_prob((fcnc::bu2kstarmumu_parameters*)theParams[b].at(n),B_MASS_LOW, B_MASS_HIGH));
        }


        //Plot everything
        bool sigRegion = true; //Plot only in the signal region
        for(UInt_t n = 0; n < nPDFs; n++){
            tag = get_MassFit_label(fitReference, b,nBins, splitRuns, pars, pdf_idx.at(n));
            theOptions[n].plot_label = "LHCb data";
            theOptions[n].q2_label =q2_label(opts.TheQ2binsmin.at(b), opts.TheQ2binsmax.at(b));

            if (nPDFs>1){ //If only one PDF, just plot the "added pdfs"
                spdlog::info("[PLOT]\t"+theOptions[n].plot_label);
                thePlotter[n]->plot_data((fcnc::bu2kstarmumu_pdf*)theProbs[b].at(n),
                                         (fcnc::bu2kstarmumu_parameters*)theParams[b].at(n),
                                         selection[b].at(n), get_MassFitPlot_path(),tag, sigRegion);
            }
        }
        //Plot joint result
        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];

        tag = get_MassFit_label(fitReference, b,nBins, splitRuns, pars, pars.Run);
        thePlotter[0]->plot_added_pdfs(prober, paramser, &selection[b],
                                       get_MassFitPlot_path(), tag, sigRegion);
    }//end bin loop

     //Plot the significances and yields in the fancy plot
    //TODO: unhardcode the runs
    if (!fitReference && nBins>1){
        std::string tag = get_MassFit_label(fitReference, -1,nBins, splitRuns, pars, 12);
        spdlog::info("Saving yield plots into " + tag);
        thePlotter[0]->plotYieldInQ2(true, theParams, theProbs, nPDFs, evts_cntr, get_MassFitPlot_path(), tag);
        thePlotter[1]->plotYieldInQ2(false, theParams, theProbs, nPDFs, evts_cntr, get_MassFitPlot_path(), tag);
    }

    //--------------------------------
    //        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];
            std::string results_file = get_MassFitResult_path()+ "fitresult_"
                    +get_MassFit_label(fitReference, b, nBins, splitRuns, pars, pars.Run)+ ".txt";
            paramser->at(i)->save_param_values(results_file);
        }
    }

    //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_bkg_yields(nBins, pdf_idx, evts_cntr, f_sigs, f_sigserr);
    }
    print_sig_yields_tex(get_MassFit_label(fitReference, -1,nBins,
                                           splitRuns, pars, pars.Run),
                         nBins, pdf_idx, &opts, evts_cntr, f_sigs, f_sigserr);


    //Save results to root file
    std::string results_file = final_result_name_mass(fitReference, nBins, splitRuns, pars, pars.Run);
    save_results(results_file, nBins, pdf_idx, fit_results, theParams, true, &opts);

    //Close Latex file
    myFile.close();

    spdlog::info("[MASSFIT]\tMass fit finished.");
    return 0;
}