//Renata Kopecna

#include "generatetoys.hh"
#include <bu2kstarmumu_parameters.hh>
#include <helpers.hh>
#include <event.hh>
#include <paths.hh>
#include <folder.hh>

#include <spdlog.h>

//Options to generate toys:
//Runs: 1, 2 or 1+2
//Signal only / Bkg only / Both
//Rare / Jpsi
//Bins
//Mass only / Angles only
//S or P

//Values taken from: MCsig, MCref, Bkg (top/bottom/both), Reference
//Can be also just mass/angle fit

//TODO: move this somewhere else
bool LowMassFit = false;
bool HighMassFit = true;


int initialize_parameters_from_MCfit(fcnc::bu2kstarmumu_parameters *parameters, bool fromRef, int PDF, int bin, std::vector<std::string> names, basic_params params, fixConstr fixConstrain){
    bool SimFit = true;
    bool onlyAngs = false;
    int nBins = fromRef ? 1 : params.nBins; //When looping over bins, this protects from trying to read from refMC
    std::string fileName_MC = final_result_name_MC(params, nBins, fromRef, false, SimFit, onlyAngs, fromRef);
    parameters->get_param_from_rootfile(fileName_MC, names, PDF, fromRef ? 0 : bin, fixConstrain); //TODO: use the struct fixConstr also in the parameters
    return 0;
}

int initialize_parameters_from_BkgFit(fcnc::bu2kstarmumu_parameters *parameters, bool fromRef, bool LowMassFit, bool HighMassFit, bool fitKpi, int PDF,  int bin, std::vector<std::string> names, basic_params params,fixConstr fixConstrain){
    int nBins = fromRef ? 1 : params.nBins;
    std::string fileName_topBkg =  final_result_name_bkg(nBins, fromRef, LowMassFit, HighMassFit, fitKpi, params);
    parameters->get_param_from_rootfile(fileName_topBkg, names, PDF,fromRef ? 0 : bin, fixConstrain);
    return 0;
}

int initialize_parameters_from_MassFit(fcnc::bu2kstarmumu_parameters *parameters, bool fromRef, int PDF,  int bin, std::vector<std::string> names, basic_params params, fixConstr fixConstrain){
    int nBins = fromRef ? 1 : params.nBins; //When looping over bins, this protects from trying to read from refMC
    bool splitRuns = true; //Get it per run or not
    int b = fromRef ? 0 : bin;
    //If not, use the PDF number form the parameter Run
    std::string fileName_dataMass = final_result_name_mass(fromRef, nBins, params.Run, params, params.Run);
    parameters->get_param_from_rootfile(fileName_dataMass, names, splitRuns ? PDF : params.Run, b, fixConstrain);

    return 0;
}

int initialize_B_mass(fcnc::bu2kstarmumu_parameters *parameters, bool fromRef, int PDF,  int bin, basic_params params){
    int nBins = fromRef ? 1 : params.nBins; //When looping over bins, this protects from trying to read from refMC
    bool splitRuns = true; //Get it per run or not
    int b = fromRef ? 0 : bin;
    //Not part of the other functions as the range needs to be full
    std::string fileName_dataMass = final_result_name_mass(fromRef, nBins, params.Run, params, params.Run);
    parameters->m_b.init(get_param_value_from_rootfile(fileName_dataMass,"m_b",  splitRuns ? PDF : params.Run,b),B_MASS_LOW,B_MASS_HIGH,0);
    return 0;
}


//Take the default values, set in constants.hh and bu2kstarmumu_parameters.cc
int initialize_default_parameters(fcnc::bu2kstarmumu_parameters *parameters){
    parameters->use_default();
    return 0;
}

//Generates toys identical to the given parameters
std::vector<fcnc::event> generateToys(basic_params params, int bin, int PDF,
                                      fcnc::bu2kstarmumu_parameters *toyParameters,
                                      fcnc::options opts){
/* These options should not be modified inside here
    //        Set options         //
    //The options will be taken from the fitter itself
    opts.only_angles = false;
    opts.only_Bmass = false;

    opts.weighted_fit = true;
    opts.use_weighted_bkg = false;
    opts.always_generate_full_angular = true;

    opts.swave = false;
    opts.shift_lh = false;
    opts.individual_penalties = false;

    //Stat. uncertainty determination tool:
    opts.squared_hesse = true;
    opts.minos_errors = false;

    opts.generate_mkpi = true;
    opts.simple_mkpi = false;
    opts.isobar = false;
    opts.mkpi_full_range_norm = false;

    opts.flat_bkg = false;
*/
    //Init the q2 options, to be sure
    opts.q2_min = opts.TheQ2binsmin[bin];
    opts.q2_max = opts.TheQ2binsmax[bin];

    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. //TODO
    if (opts.run == 1) opts.year = 2012; //Cause why not :)
    else               opts.year = 2017;

    //Create parameter set
    fcnc::bu2kstarmumu_pdf        prob(&opts, toyParameters);
    fcnc::bu2kstarmumu_generator  gen(&opts);

    prob.load_coeffs_eff_phsp_4d();
    prob.update_cached_normalization(toyParameters);
    //And finally generate the events
    //How much events relative to the requedred total number of events?
    double frac = eventsInBin_fraction(bin,opts.run,params.nBins,params.reference);
    int N = std::lround(params.nEvents*frac);
    spdlog::info("Generate {0:d} toy events", N);
 
    std::vector<fcnc::event> genToys = gen.generate(N, toyParameters, &prob);

    //Update cached efficiencies of the angular acceptance correction weights
    prob.update_cached_efficiencies(toyParameters, &genToys);

    //for folded analysis, fold the angular values of all generated events before returning
    if(!opts.full_angular){
        spdlog::debug("Folding the generated {0:d} events according to fold #{1:d}", genToys.size(), opts.folding);
        fcnc::folder theFolder(&opts);
        for(UInt_t e = 0; e < genToys.size(); e++){
            fcnc::event * evt = &genToys.at(e);
            theFolder.fold(evt);
            spdlog::trace("Folded event: ctk={0:f}\t ctl={1:f}\t phi={2:f}\t m={3:f} \tmkpi={4:f}", evt->costhetak, evt->costhetal, evt->phi, evt->m, evt->mkpi);
        }
    }
    return genToys;
}

std::vector<fcnc::event> generateIndividualToys(basic_params params, int bin, int PDF,
                                                fcnc::bu2kstarmumu_parameters *toyParameters,
                                                fcnc::options opts, bool genOnlySig, bool genOnlyBkg){

    if (genOnlySig && genOnlyBkg){
        spdlog::error("Cannot generate only bkg and only sig in one go!");
        spdlog::error("Check your options.");
        return {{}};
    }

    bool genReference = params.reference;

    //        Set options         //
    //The options will be taken from the fitter itself
    opts.only_angles = false;
    opts.only_Bmass = false;

    opts.weighted_fit = true;
    opts.use_weighted_bkg = false;
    opts.always_generate_full_angular = true;

    opts.swave = false;
    opts.shift_lh = false;
    opts.individual_penalties = false;

    //Stat. uncertainty determination tool:
    opts.squared_hesse = true;
    opts.minos_errors = false;

    opts.generate_mkpi = true;
    opts.generate_only_bkg = genOnlyBkg;
    opts.simple_mkpi = false;
    opts.isobar = false;
    opts.mkpi_full_range_norm = false;

    opts.flat_bkg = false;

    //Init the q2 options, to be sure
    opts.q2_min = opts.TheQ2binsmin[bin];
    opts.q2_max = opts.TheQ2binsmax[bin];

    opts.run = PDF; //Set proper run to options
    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.
    if (opts.run == 1) opts.year = 2012; //Cause why not :)
    else               opts.year = 2017;

    //Create parameter set
    fcnc::bu2kstarmumu_pdf        prob(&opts, toyParameters);
    fcnc::bu2kstarmumu_generator  gen(&opts);

    prob.load_coeffs_eff_phsp_4d();
    prob.update_cached_normalization(toyParameters);

    //    Initialize everything   //
    //Note that hte code will crash if the file is not there, so test before submititng pls

    //Set the f_sig
    if (genOnlySig)      toyParameters->f_sig.init_fixed(1.0);
    else if (genOnlyBkg) toyParameters->f_sig.init_fixed(0.0);
    else{ //Or both, in that case load the f_sig from mass only fit
         initialize_parameters_from_MassFit(toyParameters,genReference,PDF,bin,{"f_sig"},params,fixConstr(false,false));
    }

    //Init angular parameters from MC
    if (genReference) toyParameters->init_ang_parameters_fromRefDavid(bin,1.0,0.01);
    else initialize_parameters_from_MCfit(toyParameters,false,PDF,bin,param_string(opts, true),params,fixConstr(false,false));

    //Init the ratio of sigmas in signal MC/reference MC
    //As everything is fixed everywhere and not contrained in mass, this is fine, but this should be done properly //TODO
    initialize_parameters_from_MCfit(toyParameters,genReference,PDF,bin,params_string_mass(opts),params,fixConstr(true,false));

    if (!genReference) toyParameters->m_scale.init_fixed(get_sigmaRatio_fromMC(params,params.nBins,bin, PDF));
    else               toyParameters->m_scale.init_fixed(1.0);


    //Init the bkg parameters from the BKG fit file
         //First init the background for angles
    initialize_parameters_from_BkgFit(toyParameters, true, LowMassFit, HighMassFit, false, 12, bin,param_string_bkg(),params,fixConstr(false,false));
        //Then init the background for m_Kpi
    initialize_parameters_from_BkgFit(toyParameters, true, LowMassFit, HighMassFit, true, 12, bin,param_string_bkg_mkpi(),params,fixConstr(false,false));

    //Init the mass background from mass fit to reference for now
    std::vector<std::string> fixFromMassFit; //TODO: set as a proper vector depending on using 2 lambdas or anything
    if (opts.fit_lambda) fixFromMassFit.push_back("m_lambda");
    else                 fixFromMassFit.push_back("m_tau");
    initialize_parameters_from_MassFit(toyParameters,true,PDF,bin,fixFromMassFit,params,fixConstr(false,false));

    //Initialize the q2 by hand, as it is the easiest
    toyParameters->eff_q2.init(bin_center_q2(opts,bin), opts.q2_min, opts.q2_max, 0.0);

    //Set the Bmass separatelly as it's range is important
    initialize_B_mass(toyParameters,true,PDF,bin,params);

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

    //Update the normalisation with the new parameter settings
    prob.update_cached_normalization(toyParameters);

    //And finally generate the events
    //How much events relative to the requedred total number of events?
    double frac = eventsInBin_fraction(bin,opts.run,params.nBins,params.reference);
    int N = std::lround(params.nEvents*frac);
    spdlog::info("Generate {0:d} toy events", N);
    std::vector<fcnc::event> genToys = gen.generate(N, toyParameters, &prob);

    //for folded analysis, fold the angular values of all generated events before returning
    if(!opts.full_angular){
        spdlog::debug("Folding the generated {0:d} events according to fold #{1:d}", genToys.size(), opts.folding);
        fcnc::folder theFolder(&opts);
        for(auto evt: genToys){
            theFolder.fold(&evt);
            spdlog::trace("Folded event: ctk={0:f}, ctl={1:f}, phi={2:f}", evt.costhetak, evt.costhetal, evt.phi);
        }
    }
    return genToys;
}

int saveToys(basic_params params, fcnc::options opts,  bool genOnlySig, bool genOnlyBkg){

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


    const int nBins = params.nBins;

    std::vector<fcnc::parameters*> toyParams [nBins];
    std::vector<int>tmpRes[nBins];  //instead of fit results

    //Now just for saving do this
    for (auto idx: pdf_idx){
        std::vector<fcnc::event> events;
        for (int bin = 0; bin < params.nBins; bin++){
            fcnc::bu2kstarmumu_parameters * tmpParams = new fcnc::bu2kstarmumu_parameters (&opts);
            std::vector<fcnc::event> tmp = generateIndividualToys(params, bin, idx, tmpParams, opts,genOnlySig, genOnlyBkg);
            events.insert(events.end(),tmp.begin(),tmp.end());
            toyParams[bin].push_back(tmpParams);
            tmpRes[bin].push_back(0);
        }
        std::string saveFile = get_finalToys_file(params.reference, nBins, true, params, idx);
        fcnc::save_events(saveFile,events);
    }
    std::string paramFile = init_params_name_toys(-1,params.reference, nBins, true, params, params.Run, false, false, genOnlySig, genOnlyBkg, false, false); //TODO: fix the naming, as you actually do take stuff from upper-mass sideband

    save_results(paramFile,nBins,pdf_idx,tmpRes,toyParams,true,&opts);
    return 0;

}