#include <parameterscan.hh>


#include <iostream>
#include <fstream>

#include <paths.hh>
#include <options.hh>
#include <funcs.hh>
#include <generator.hh>
#include <parameters.hh>
#include <event.hh>
#include <funcs.hh>
#include <pdf.hh>
#include <fitter.hh>
#include <helpers.hh>

#include <TFile.h>
#include <TTree.h>
#include <TCanvas.h>

#include <spdlog.h>

/**
 * @file parameterscan.cc
 * @author Renata Kopecna
 * @date 2021-03-02
 *
 */

fcnc::parameterscan::parameterscan(options* o):
  multifit(o)
{
};

void fcnc::parameterscan::scan(std::string param_name, double min, double max, unsigned int nsteps, unsigned int nevents, unsigned int reruns, pdf* prob, parameters* params, generator* gen, bool only_float_in_generation){
    std::vector<unsigned int > events;
    events.push_back(nevents);

    std::vector<pdf* > probs;
    probs.push_back(prob);

    std::vector<parameters* > the_params;
    the_params.push_back(params);

    std::vector<generator* > gens;
    gens.push_back(gen);

    scan(param_name, min, max, nsteps, events, reruns, probs, the_params, gens, only_float_in_generation);
}

//void parameterscan::scan(std::string param_name, double min, double max, unsigned int nsteps, unsigned int nevents, unsigned int reruns, pdf_type* prob, params_type* params, generator_type* gen, bool only_float_in_generation) //fixed in fit on fixed value
void fcnc::parameterscan::scan(std::string param_name, double min, double max, unsigned int nsteps, std::vector<unsigned int> nevents, unsigned int reruns, std::vector<pdf*> pdfs, std::vector<parameters*> params, std::vector<generator*> gens, bool only_float_in_generation){
    parameter* param_varied = 0;
    for (unsigned int i=0; i<params.size(); i++){
        if (params.at(i)->get_parameter(param_name)) {
            param_varied = params.at(i)->get_parameter(param_name);
            break;
        }
    }

    //open TeX file
    std::ofstream myFile;
    open_Latex_noteFile(latex_paramScan(), myFile);

    //parameter* param_varied = params->get_parameter(param_name);

    double fit_start_value = param_varied->get_start_value();
    if (only_float_in_generation) spdlog::info( "The fit start value will be fixed at: {0:f}",fit_start_value);

    int param_index = param_varied->get_index();
    if (param_varied != 0){
        std::string param_description = param_varied->get_description();
        std::string param_root = param_varied->get_root_name();
        spdlog::info( "Doing parameter scan of parameter {0:d}: {1:s}", param_index, param_name);
        spdlog::info( "Varying parameter from {0:f} to {1:f}", min, max);

        unsigned int nparams = 0;
        for (unsigned int i=0; i<params.size(); i++){
            nparams += params.at(i)->nparameters();
        }
        //params->nparameters();

        TH1D* error_hists[nparams];
        TH1D* value_hists[nparams];
        double mean_values[nparams][nsteps+1];
        double sigma_mean_values[nparams][nsteps+1];
        double mean_errors[nparams][nsteps+1];
        double sigma_mean_errors[nparams][nsteps+1];

        double min_hist = min - 0.5*(max-min)/nsteps;
        double max_hist = max + 0.5*(max-min)/nsteps;
        //create histograms
        unsigned int start_param = 0;
        for (unsigned int i=0; i<params.size(); i++){
            for (unsigned int j = 0; j < nparams; j++){
                unsigned int idx = start_param+j;
                parameter* param = params.at(i)->get_parameter(j);
                error_hists[idx] = new TH1D(("error_" + param->get_name()).c_str(),
                                            ("#sigma" + param->get_root_name() + ";" + param_root).c_str(),
                                            nsteps+1, min_hist, max_hist);
                error_hists[idx]->SetMarkerStyle(8);
                error_hists[idx]->SetMarkerColor(4);
                value_hists[idx] = new TH1D(("value_" + param->get_name()).c_str(),
                                            (param->get_root_name() + ";" + param_root).c_str(),
                                            nsteps+1, min_hist, max_hist);
                value_hists[idx]->SetMarkerStyle(8);
                value_hists[idx]->SetMarkerColor(4);
            }
            start_param += params.at(i)->nparameters();
        }
        //do precalculations
        for (unsigned int i=0; i<params.size(); i++){
            pdfs.at(i)->init(params.at(i));
        }
        //create output trees
        //current values to fill into the tree
        double value, error, error_up, error_down, start_value, param_value, nominal_error, nominal_value;
        int step, run, migrad, status_cov;
        std::vector<TTree*> trees;//one tree for every varied parameter
        //start_param = 0;
        for (unsigned int i=0; i<params.size(); i++){
            for (unsigned int j = 0; j < nparams; j++){
                //unsigned int idx = start_param+j;
                parameter* param = params.at(i)->get_parameter(j);
                if (param->get_step_size() != 0.0){
                    std::string parname(param->get_name());
                    std::string pardesc(param->get_description());// + " results");
                    TTree* t = new TTree(parname.c_str(), pardesc.c_str());
                    t->Branch("step",&step,"step/I");
                    t->Branch("run",&run,"run/I");
                    t->Branch("param_value",&param_value,"param_value/D");
                    t->Branch("value",&value,"value/D");
                    t->Branch("error",&error,"error/D");
                    t->Branch("nominal_value",&nominal_value,"nominal_value/D");
                    t->Branch("nominal_error",&nominal_error,"nominal_error/D");
                    t->Branch("error_up",&error_up,"error_up/D");
                    t->Branch("error_down",&error_down,"error_down/D");
                    t->Branch("start_value",&start_value,"start_value/D");
                    t->Branch("migrad",&migrad,"migrad/I");
                    t->Branch("status_cov",&status_cov,"status_cov/I");
                    trees.push_back(t);
                }
            }
            //start_param += params.at(i)->nparameters();
        }

        //enter main loop
        for (unsigned int i = 0; i <= nsteps; i++){
            //reset all results
            std::vector<std::vector<double> > values(nparams, std::vector<double>(reruns, 0.0));
            std::vector<std::vector<double> > errors(nparams, std::vector<double>(reruns, 0.0));
            std::vector<std::vector<double> > errors_up(nparams, std::vector<double>(reruns, 0.0));
            std::vector<std::vector<double> > errors_down(nparams, std::vector<double>(reruns, 0.0));
            std::vector<std::vector<double> > nominal_values(nparams, std::vector<double>(reruns, 0.0));
            std::vector<std::vector<double> > nominal_errors(nparams, std::vector<double>(reruns, 0.0));
            std::vector<int> return_values(reruns);

            //only_float_in_generation means that you use a varied parameter
            //for generation but in the fit use the start_value
            //this can be used for systematic studies
            //we would also need to do a nominal fit
            //then we can see the actual difference

            //choose new value for the scan parameter
            param_value = double(i)*(max - min)/(nsteps) + min;
            if (!only_float_in_generation){ //float in generation and fit
                param_varied->set_value(param_value);
                param_varied->set_start_value(param_value);
            }
            // param_varied->init(param_name, param_description,
            // 		       param_value,
            // 		       param_varied->get_min(), param_varied->get_max(),
            // 		       param_varied->get_step_size());

            //this is the scanned parameter so usually we change it before generation and fit
            //this is the point: we need to do a couple reruns
            fitter fit(opts);
            fit.set_common_parameters(common_params);

            for (unsigned int k = 0; k < reruns; k++){
                spdlog::info( "Parameter "+param_name+" {0:f}", param_value);
                spdlog::info( "Run no. {0:d}", k);
                std::string num;
                std::stringstream out;
                out << opts->name << "_" << (i+1) << "_" << (k+1); //TOOD: move to paths
                num = out.str();

                if (only_float_in_generation){//this generates with updated parameter values
                    param_varied->set_value(param_value);
                    param_varied->set_start_value(param_value);
                }
                // param_varied->init(param_name, param_description,
                // 		   param_value,
                // 		   param_varied->get_min(), param_varied->get_max(),
                // 		   param_varied->get_step_size());

                for (unsigned int j=0; j<params.size(); j++){
                    params.at(j)->reset_parameters();//this line was missing before? did we generate with fitted params?
                }

                std::vector< std::vector <event> > temp_events;
                for (unsigned int j=0; j<gens.size(); j++){
                    std::vector<event> ev = gens.at(j)->generate(nevents.at(j), params.at(j), pdfs.at(j));
                    temp_events.push_back(ev);
                }
                std::vector< std::vector <event> * > events;
                for (unsigned int j=0; j<gens.size(); j++){
                    events.push_back(&temp_events.at(j));
                }

                if (only_float_in_generation){//this tries to fit with original parameter values
                    param_varied->set_value(fit_start_value);
                    param_varied->set_start_value(fit_start_value);
                }
                // param_varied->init(param_name, param_description,
                // 		   fit_start_value, //use start value here!
                // 		   param_varied->get_min(), param_varied->get_max(),
                // 		   param_varied->get_step_size());


                for (unsigned int j=0; j<params.size(); j++){
                    params.at(j)->reset_parameters();
                }
                int result = 0;
                result = fit.fit(pdfs, params, events, num);

                return_values.at(k) = result;
                if (result != 0 && opts->repeat_on_fail){
                    spdlog::warn( "Fit failed, repeating run");
                    k--;
                    continue;
                }
                unsigned int start_param = 0;
                for (unsigned int j = 0; j < params.size(); j++){
                    for (unsigned int l = 0; l < params.at(j)->nparameters(); l++){
                        unsigned int idx = start_param+l;
                        parameter* param = params.at(j)->get_parameter(l);
                        values.at(idx).at(k) = param->get_value();
                        errors.at(idx).at(k) = param->get_error();
                        errors_up.at(idx).at(k) = param->get_error_up();
                        errors_down.at(idx).at(k) = param->get_error_down();
                    }
                    start_param += params.at(j)->nparameters();
                }
                //now we want to do a nominal fit (if the respective options are set)
                if (opts->refitting_nominal && only_float_in_generation){//now we do a nominal fit
                    spdlog::info( "Nominal fit");

                    param_varied->set_value(param_value);
                    param_varied->set_start_value(param_value);
                    // param_varied->init(param_name, param_description,
                    // 		     param_value,
                    // 		     param_varied->get_min(), param_varied->get_max(),
                    // 		     param_varied->get_step_size());

                    for (unsigned int j = 0; j < params.size(); j++){
                        params.at(j)->reset_parameters();
                    }
                    //and lets fit
                    fit.fit(pdfs, params, events, num);

                    //extract the parameters from the fit
                    start_param = 0;
                    for (unsigned int j = 0; j < params.size(); j++){
                        for (unsigned int l = 0; l < params.at(j)->nparameters(); l++){
                            unsigned int idx = start_param+l;
                            parameter* param = params.at(j)->get_parameter(l);
                            nominal_values.at(idx).at(k) = (param->get_value());
                            nominal_errors.at(idx).at(k) = (param->get_error());
                        }
                        start_param += params.at(j)->nparameters();
                    }
                }
                else{
                    for (unsigned int j = 0; j < nparams; j++){
                        nominal_values.at(j).at(k) = (0.0);
                        nominal_errors.at(j).at(k) = (0.0);
                    }
                }

            }//end reruns

            //determine means/widths of the distributions via gauss-fits
            unsigned int start_param = 0;
            for (unsigned int j = 0; j < params.size(); j++){
                for (unsigned int k = 0; k < params.at(j)->nparameters(); k++){
                    unsigned int idx = start_param + k;
                    parameter* param = params.at(j)->get_parameter(k);
                    if (param->get_step_size() != 0.0){
                        std::string step;
                        std::stringstream out;
                        out << (i+1);
                        step = out.str();

                        double dummy1, dummy2, dummy3, dummy4, dummy5;
                        double mean_value, mean_error, sigma_mean_value, sigma_mean_error;
                        update_pull(param, values.at(idx), errors.at(idx),
                                    dummy1, dummy2, dummy3, dummy4, dummy5, step);
                        update_value(param, values.at(idx), errors.at(idx),
                                     mean_value, sigma_mean_value, dummy3, dummy4, dummy5, step);
                        update_error(param, values.at(idx), errors.at(idx),
                                     mean_error, sigma_mean_error, dummy3, dummy4, dummy5, step);

                        value_hists[idx]->SetBinContent(i+1,mean_value);
                        value_hists[idx]->SetBinError(i+1,sigma_mean_value);
                        error_hists[idx]->SetBinContent(i+1,mean_error);
                        error_hists[idx]->SetBinError(i+1,sigma_mean_error);

                        mean_values[idx][i] = mean_value;//for the output tables
                        sigma_mean_values[idx][i] = sigma_mean_value;//for the output tables
                        mean_errors[idx][i] = mean_error;
                        sigma_mean_errors[idx][i] = sigma_mean_error;
                    }
                }
                start_param += params.at(j)->nparameters();
            }

            //save tree results
            int param_idx = 0;
            output->cd();
            start_param = 0;
            for (unsigned int j = 0; j < params.size(); j++){//this gives nice output
                for (unsigned int l = 0; l < params.at(j)->nparameters(); l++){
                    unsigned int idx = start_param+l;
                    parameter* param = params.at(j)->get_parameter(l);
                    if (param->get_step_size() != 0.0){
                        TTree* t = trees.at(param_idx);
                        for (unsigned int k=0; k<reruns; k++){
                            step = i;
                            run = k;
                            //param_value is already set;
                            value = values.at(idx).at(k);
                            error = errors.at(idx).at(k);
                            nominal_value = nominal_values.at(idx).at(k);
                            nominal_error = nominal_errors.at(idx).at(k);
                            error_up = errors_up.at(idx).at(k);
                            error_down = errors_down.at(idx).at(k);
                            start_value = param->get_start_value();
                            migrad = return_values.at(k) % 100;
                            status_cov = return_values.at(k) / 100;
                            t->Fill();
                        }
                        param_idx++;
                    }
                }
                start_param += params.at(j)->nparameters();
            }
        }
        //save output trees
        for (unsigned int j = 0; j < trees.size(); j++){
            trees.at(j)->Write();
        }

        //summary output
        myFile << "Parameter scan results: values\\\\" << std::endl;
        myFile << "\\begin{tabular}{|c|";
        for (unsigned int i=0; i<=nsteps; i++){
            myFile << "c";
        }
        myFile << "|}\\hline" << std::endl;
        myFile << "$" << param_varied->get_description() << "$";
        for (unsigned int i=0; i<=nsteps; i++){
            double param_value = double(i)*(max - min)/(nsteps) + min;
            myFile << " & " << std::setprecision(2) << param_value;
        }
        myFile << "\\\\ \\hline" << std::endl;

        start_param = 0;
        for (unsigned int i=0; i < params.size(); i++){
            for (unsigned int j=0; j < params.at(i)->nparameters(); j++){
                unsigned int idx = start_param+j;
                if (params.at(i)->get_parameter(j)->get_step_size() != 0.0){
                    myFile << "$" << params.at(i)->get_parameter(j)->get_description() << "$";
                    for (unsigned int k=0; k <= nsteps; k++){
                        myFile << " & $" << std::setprecision(2)
                               << mean_values[idx][k] << " \\pm "
                                  //<< mean_errors[i][j]
                               << sigma_mean_values[idx][k]
                                  << "$";
                    }
                    myFile << "\\\\" << std::endl;
                }
            }
            start_param += params.at(i)->nparameters();
        }
        myFile << "\\hline \\end{tabular}" << std::endl;

        myFile << "Parameter scan results: errors\\\\" << std::endl;
        myFile << "\\begin{tabular}{|c|";
        for (unsigned int i=0; i<=nsteps; i++){
            myFile << "c";
        }
        myFile << "|}\\hline" << std::endl;
        myFile << "$" << param_varied->get_description() << "$";
        for (unsigned int i=0; i<=nsteps; i++){
            double param_value = double(i)*(max - min)/(nsteps) + min;
            myFile << " & " << std::setprecision(2) << param_value;
        }
        myFile << "\\\\ \\hline" << std::endl;

        start_param = 0;
        for (unsigned int i=0; i < params.size(); i++){
            for (unsigned int j=0; j < params.at(i)->nparameters(); j++){
                unsigned int idx = start_param+j;
                if (params.at(i)->get_parameter(j)->get_step_size() != 0.0){
                    myFile << "$" << params.at(i)->get_parameter(j)->get_description() << "$";
                    for (unsigned int k=0; k <= nsteps; k++){
                        myFile << " & $" << std::setprecision(2)
                               << mean_errors[idx][k] << " \\pm "
                               << sigma_mean_errors[idx][k] << "$";
                    }
                    myFile << "\\\\" << std::endl;
                }
            }
            start_param += params.at(i)->nparameters();
        }
        myFile << "\\hline \\end{tabular}" << std::endl;

        start_param = 0;
        for (unsigned int i=0; i < params.size(); i++){
            for (unsigned int j=0; j < params.at(i)->nparameters(); j++){
                unsigned int idx = start_param+j;
                //print result histos
                if (params.at(i)->get_parameter(j)->get_step_size() != 0.0){
                    output->cd();
                    parameter* param = params.at(i)->get_parameter(j);

                    std::string error_name = "parameter scan sigma(" + param->get_name() + ")";
                    TCanvas *c1 = new TCanvas(error_name.c_str(), error_name.c_str(), 1600, 1200);
                    c1->cd();
                    error_hists[idx]->Draw("e");
                    if (opts->write_eps) c1->Print(get_param_scan_path(param->get_name(),true).c_str(), "eps");
                    c1->Write();
                    delete error_hists[idx];
                    delete c1;

                    std::string value_name = "Parameter scan "+param->get_name();
                    TCanvas *c2 = new TCanvas(value_name.c_str(), value_name.c_str(), 1600, 1200);
                    c2->cd();
                    value_hists[idx]->Draw("e");
                    if (opts->write_eps) c2->Print(get_param_scan_path(param->get_name(),false).c_str(), "eps");
                    c2->Write();
                    delete value_hists[idx];
                    delete c2;

                    //TODO add pulls
                }
            }
            start_param += params.at(i)->nparameters();
        }
    }
    else{
        spdlog::error("Parameter "+param_name+" not found.");
        assert(0);
    }
}