//Renata Kopecna

#include <backgroundfit.hh>

#include <fstream>
#include <sstream>      // std::istringstream

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

#include <spdlog.h>

#include <TStyle.h>
#include <TH2D.h>
#include <TLatex.h>
#include <TCanvas.h>
#include <TROOT.h>

//Do the actuall background fit
int backgroundfit(fcnc::options opts,
                  bool fitReference, bool LowMassFit, bool HighMassFit, bool fitKpi,
                  bool Use2DAngularBins, basic_params params){
    //params now only used to tag nBins in the names

    //Technically, this should be simFit between Run 1 and 2, but since we need it only for the toys atm, no need for it

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

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

    spdlog::info("[FIT]\t\tFit only angular and m(Kpi) background");


    //Used to check whether the background is "correlated" between each other
    //TODO: not needed atm, maybe fix later
    /* 1D angular bins
      1:        ctl < -0.5
      2: -0.5 < ctl <  0.0
      3:  0.0 < ctl <  0.5
      4:  0.5 < ctl
      5:        ctk < -0.5
      6: -0.5 < ctk <  0.0
      7:  0.0 < ctk <  0.5
      8:  0.5 < ctk
      */
    /* 2D angular bins
      1: ctl < 0.0 && ctk < 0.0
      2: ctl < 0.0 && ctk > 0.0
      3: ctl > 0.0 && ctk < 0.0
      4: ctl > 0.0 && ctk > 0.0
      */

    double HighBpeakCut = fitReference ? B_MASS_HIGH_BKG_REF : B_MASS_HIGH_BKG;
    double LowBpeakCut  = fitReference ? B_MASS_LOW_BKG_REF : B_MASS_LOW_BKG;

    assert(!(HighMassFit && LowMassFit));

    opts.fit_full_angular_bkg = true;

    opts.only_angles = !(HighMassFit || LowMassFit); //Fit mass only if one sideband is selected
    opts.only_Bmass = false;
    opts.only_mkpi = false;

    opts.swave = false;

    opts.bkg_order_costhetak = 5;
    opts.shift_lh = false; //keep it, it was there before
    opts.weighted_fit = true;
    opts.squared_hesse = true;
    opts.minos_errors = false;
    opts.asymptotic = false;

    opts.flat_bkg = false;

    opts.fit_mkpi = false; //set to true only later
    opts.use_mkpi = false;
    opts.simple_mkpi = false; //S-wave model
    opts.isobar = false; //S-wave model

    //Plotting options
    opts.plot_chi2 = false;
    opts.plots_m_bins = 20;
    opts.plots_mkpi_bins = 20;

    //load events from data tuple
    std::vector<std::vector<fcnc::event>>events;
    std::vector<int> run_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);
        run_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);
        run_idx.push_back(2);
    }

    std::vector<int> fitresults;

    //determine number of bins:
    const UInt_t nBins = opts.TheQ2binsmin.size();

    //current fitter, plotter, parameteres and pdfs:
    fcnc::fitter f(&opts);
    fcnc::folder fldr(&opts);
    fcnc::bu2kstarmumu_plotter thePlotter(&opts);

    fcnc::bu2kstarmumu_parameters * theParams[nBins];
    fcnc::bu2kstarmumu_pdf        * theProb[nBins];
    std::vector<fcnc::event> * theEvents[nBins];


    for(UInt_t b = 0; b < nBins; b++){

        theParams[b] = new fcnc::bu2kstarmumu_parameters(&opts);
        theProb[b]   = new fcnc::bu2kstarmumu_pdf(&opts, theParams[b]);
        theEvents[b] = new std::vector<fcnc::event>();

        //Init parameters
        theParams[b]->f_sig.init_fixed(0.0);
        theParams[b]->m_b.init(PDGMASS_B, HighMassFit ? HighBpeakCut : B_MASS_LOW, LowMassFit ? LowBpeakCut : B_MASS_HIGH, 0.0);
        //The rest of the mass parameters can be left to default values
        //They are defined in bu2kstamumu_parameters

        //Init the bkg shape
        theParams[b]->init_angular_background_parameters(fitReference,0.1);

        //Init the mass bkg shape
        theParams[b]->init_mass_background_parameters(nBins,b,true);
        //theParams[b]->cbkgctk3.init_fixed(-2.5);

        //Select events        
        for_indexed(auto evts: events){
            //update weights according to sub-set for all events
            opts.run  = run_idx.at(i);
            opts.update_efficiencies = true;

            theProb[b]->load_coeffs_eff_phsp_4d();
            theProb[b]->update_cached_normalization(theParams[b]);
            theProb[b]->update_cached_efficiencies(theParams[b], &evts);
            opts.update_angle_ranges(); //Update angles to get rid of events that cannot be folded
            //select events
            spdlog::debug("Loop over {0:d} events and select the suitable ones", evts.size());
            for(auto meas: evts){

                if(meas.q2 < opts.TheQ2binsmin.at(b) || meas.q2 > opts.TheQ2binsmax.at(b)) continue;
                if(meas.m < opts.m_low || meas.m > opts.m_high) continue;

                if(LowMassFit && meas.m > LowBpeakCut) continue;
                else if(HighMassFit && meas.m < HighBpeakCut) continue;
                else if(meas.m > LowBpeakCut && meas.m < HighBpeakCut) continue;

                if(meas.mkpi < opts.mkpi_min || meas.mkpi > opts.mkpi_max) continue;

                if(params.bin > 1){
                    bool keep_event = true;
                    if(Use2DAngularBins){
                        switch(params.bin){
                        case 2:
                            if(meas.costhetal > 0.0 || meas.costhetak > 0.0) keep_event = false;
                            break;
                        case 3:
                            if(meas.costhetal > 0.0 || meas.costhetak < 0.0) keep_event = false;
                            break;
                        case 4:
                            if(meas.costhetal < 0.0 || meas.costhetak > 0.0) keep_event = false;
                            break;
                        case 5:
                            if(meas.costhetal < 0.0 || meas.costhetak < 0.0) keep_event = false;
                            break;
                        }
                    }
                    else{ //1D angular bins
                        switch(params.bin){
                        case 2:
                            if(meas.costhetal > -0.5) keep_event = false;
                            break;
                        case 3:
                            if(meas.costhetal < -0.5 || meas.costhetal > 0.0) keep_event = false;;
                            break;
                        case 4:
                            if(meas.costhetal < 0.0 || meas.costhetal > 0.5) keep_event = false;;
                            break;
                        case 5:
                            if(meas.costhetal > 0.5) keep_event = false;;
                            break;
                        case 6:
                            if(meas.costhetak > -0.5) keep_event = false;;
                            break;
                        case 7:
                            if(meas.costhetak < -0.5 || meas.costhetak > 0.0) keep_event = false;;
                            break;
                        case 8:
                            if(meas.costhetak < 0.0 || meas.costhetak > 0.5) keep_event = false;;
                            break;
                        case 9:
                            if(meas.costhetak > 0.5) keep_event = false;;
                            break;
                        }
                    }
                    if(!keep_event) continue;
                }

                if(!filterFldFour(&meas, &opts)) continue;
                fldr.fold(&meas);

                //only one combined event vector
                theEvents[b]->push_back(meas);
            }
            spdlog::debug("Selected {0:d} events.", theEvents[b]->size());
        }
        spdlog::info("[BIN {0:d}]\tDone!",b);
    }//end loop over bins

    //do not update cached efficiencies anyomre, to keep individual ones for four sub-sets
    opts.update_efficiencies = false;

    //create 2D correlation plots for all angles (3) and all bins (8): 3! * 8 = 24
    const unsigned int nANGLES = 3;

    double hmin[nANGLES] = {CTL_MIN, CTK_MIN, PHI_MIN};
    double hmax[nANGLES] = {CTL_MAX, CTK_MAX, PHI_MAX};
    double corrfactor[nBins][nANGLES][nANGLES];
    const int nCorrBins = 10;
    for(UInt_t b = 0; b < nBins; b++){
        spdlog::info("[START]\tStart the background 2D correlation plotting for bin #{0:d}", b);
        for(UInt_t a = 0; a < nANGLES - 1; a++){
            for(UInt_t aa = a + 1; aa < nANGLES; aa++){
                std::string mainName = "bckgnd_correl_"+ANGLES[a]+"_"+ANGLES[aa]+"_q2bin"+std::to_string(b);
                TCanvas * cAngCorr = new TCanvas(("c"+mainName).c_str(),
                                                 ("c"+mainName).c_str(), 1400, 1200);
                cAngCorr->cd()->SetMargin(0.1,0.125,0.125,0.125);
                cAngCorr->cd();
                TH2D * h = new TH2D((mainName).c_str(),
                                    (mainName+";"+latex_angles[a]+";"+latex_angles[aa]).c_str(),
                                    nCorrBins, hmin[a], hmax[a], nCorrBins, hmin[aa], hmax[aa]);

                for(UInt_t e = 0; e < theEvents[b]->size(); e++){
                    fcnc::event meas = theEvents[b]->at(e);
                    double x=0.0, y=0.0;
                    switch(a){ //TODO: make a function in helpers
                    case 0:
                        x = meas.costhetal;
                        break;
                    case 1:
                        x = meas.costhetak;
                        break;
                    }
                    switch(aa){
                    case 1:
                        y = meas.costhetak;
                        break;
                    case 2:
                        y = meas.phi;
                        break;
                    }
                    h->Fill(x, y);
                }
                h->GetZaxis()->SetRangeUser(0, fitReference ? 25 : 5);
                h->GetZaxis()->SetTitle("Number of entries");

                h->GetXaxis()->SetTitleSize(0.05);
                h->GetXaxis()->SetLabelSize(0.05);
                h->GetXaxis()->SetTitleOffset(0.95);

                h->GetYaxis()->SetTitleSize(0.05);
                h->GetYaxis()->SetLabelSize(0.05);
                h->GetYaxis()->SetTitleOffset(0.75);

                h->GetZaxis()->SetTitleSize(0.05);
                h->GetZaxis()->SetLabelSize(0.05);
                h->GetZaxis()->SetTitleOffset(0.75);

                h->SetTitle("");
                h->Draw("COLZ");

                corrfactor[b][a][aa] = h->GetCorrelationFactor();

                TLatex * leg = getPrettyTex(0.06,13);
                leg->SetTextColor(kBlack);
                std::ostringstream sCorrout; //TODO: move the path
                sCorrout << std::fixed << std::setprecision(2) << "corr(" << latex_angles[a] << ", " << latex_angles[aa] <<") = " <<  corrfactor[b][a][aa]*100. << "\%";
                leg->DrawLatex(0.22,0.98, sCorrout.str().c_str());
                leg->SetTextSize(0.04);
                leg->DrawLatex(0.25,0.92, ("Events: "+std::to_string((int) h->GetEntries())).c_str());
                cAngCorr->Print(get_bkgCorrPlot_path(ANGLES[a],ANGLES[aa],b,nBins,fitReference, LowMassFit, HighMassFit,  fitKpi, params).c_str(), "eps");

            }
        }

        //Save the correlations s
        myFile << "\\begin{tabular}{r|ccc}\\hline" << std::endl;
        myFile << std::fixed << std::setprecision(2) << "[" << opts.TheQ2binsmin.at(b) << ", " << opts.TheQ2binsmax.at(b) << "]";
        for(UInt_t a = 0; a < nANGLES; a++)  myFile << "\t&" << latex_2angles[a];
        myFile << "\\\\" << std::endl;
        myFile << "\\hline\\hline" << std::endl;
        for(UInt_t a = 0; a < nANGLES; a++){
            myFile << latex_2angles[a];
            for(UInt_t aa = 0; aa < nANGLES; aa++){
                myFile << "\t&";
                if(aa > a)       myFile << std::setprecision(3) << std::fixed << corrfactor[b][a][aa];
                else if(aa == a) myFile << "1.00";
                else if(aa < a)  myFile << "    ";
            }
            myFile << "\\\\" << std::endl;
        }
        myFile << "\\hline" << std::endl;
        myFile << "\\end{tabular}"<< std::endl;
        myFile << std::endl;

    }

    //FIT ANGULAR BACKGROUND + B mass
    opts.update_angle_ranges(); //Just in case
    //fit all bins:
    for(UInt_t b = 0; b < nBins; b++){

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

        //fit the events:
        int fitresult = f.fit(theProb[b], theParams[b], theEvents[b]);
        fitresults.push_back(fitresult);

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

        //plot pdf with the data points:
        opts.plot_label = "LHCb data";
        opts.q2_label = q2_label(opts.TheQ2binsmin.at(b), opts.TheQ2binsmax.at(b));

        if(opts.write_eps){
            std::string label = get_bkg_tag(b,nBins,fitReference, LowMassFit, HighMassFit,  false, params);
            spdlog::info("[PLOT]\t" + label);
            thePlotter.plot_data(theProb[b], theParams[b], theEvents[b], get_bkgFitPlot_path(), label, false);
        }
        //print all parameters
        std::string fitResultTxt = get_bkgFitResult_path()+ "fitresult_" + get_bkg_tag(b,nBins,fitReference, LowMassFit, HighMassFit,  false, params);
        spdlog::debug("Saving into "+fitResultTxt);
        print_all_parameters(nBins, theParams, 2, fitResultTxt);
    }

    //print all fitresults
    print_fit_results(nBins, fitresults);

    //Save results to root file
    std::string results_file = final_result_name_bkg(nBins, fitReference, LowMassFit, HighMassFit, false, params);
    save_results(results_file, nBins, params.Run, fitresults, theParams, &opts);

    //FIT M(KPI) BACKGROUND
    if (!fitKpi){ //first check if you really want to fit M(KPI) :)
        //Close Latex file
        myFile.close();
        return 0;
    }

    //Reset the options
    opts.only_angles = false;
    opts.only_mkpi = true;
    opts.fit_mkpi = true;


    //fit all bins:
    spdlog::info("[BKGFIT]\tFit the Kpi mass");
    for(UInt_t b = 0; b < nBins; b++){

        theParams[b]->use_default_bkg();
        //null the lambda(s) and taus
        theParams[b]->m_lambda.init_fixed(0.0);
        theParams[b]->m_lambda_2.init_fixed(0.0);
        theParams[b]->m_tau.init_fixed(0.0);
        theParams[b]->m_tau_2.init_fixed(0.0);
        //Init the floating parameters
        theParams[b]->init_kpi_background_parameters(fitReference,0.01);

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

        //fit the events:
        fitresults.push_back(f.fit(theProb[b], theParams[b], theEvents[b]));
        spdlog::info("[BIN{0:d} ]:\tFitresult: {1:d}",b, fitresults.at(b));

        //plot pdf with the data points:
        opts.q2_label = q2_label(opts.TheQ2binsmin.at(b), opts.TheQ2binsmax.at(b));

        if(opts.write_eps){
            std::string label = get_bkg_tag(b,nBins,fitReference, LowMassFit, HighMassFit, false, params);
            spdlog::info("[PLOT]\t" + label);
            thePlotter.plot_data(theProb[b], theParams[b], theEvents[b], get_bkgFitPlot_path(), label, false);
        }               
        //Print all parameters, but don't save
        print_all_parameters(nBins, theParams, 2, "");

        //print all parameters to txt file
        std::string fitResultTxt = get_bkgFitResult_path()+ "fitresult_" + get_bkg_tag(b,nBins,fitReference, LowMassFit, HighMassFit,  true, params);
        spdlog::debug("Saving into "+fitResultTxt);
        print_all_parameters(nBins, theParams, 2, fitResultTxt);
    }
    spdlog::info("[BKGFIT]\tDone with Kpi fit.");
    //Save results to root file
    results_file = final_result_name_bkg(nBins, fitReference, LowMassFit, HighMassFit, true, params);
    save_results(results_file, nBins, params.Run, fitresults, theParams, &opts);

    //Close Latex file
    myFile.close();

    spdlog::info("[BKGFIT]\tFinished.");
    return 0;
}