/**
 * @file likelihoodscan.hh
 * @author Christoph Langenbruch Renata Kopecna
 * @date 2020-01-12
 *
 */

//TODO: remove unused stuff
#include <likelihoodscan.hh>
#include <iostream>
#include <TMinuit.h>
#include <TTree.h>
#include <TRandom3.h>
#include <complex>
#include <assert.h>
#include <event.hh>
#include <parameters.hh>
#include <funcs.hh>
#include <pdf.hh>
#include <plotter.hh>
#include <options.hh>
#include <thread>
#include <mutex>
#include <TMatrix.h>
#include <TH1D.h>
#include <TCanvas.h>
#include <TROOT.h>
#include <TStyle.h>
#include <TVector.h>
#include <TMatrixD.h>
#include <TMatrixDSym.h>
#include <TMatrixTSym.h>
#include <TMatrixT.h>
#include <TFile.h>
#include <TDecompChol.h>

#include <fitter.hh>
#include <spdlog.h>


bool fcnc::likelihoodscan::scan_1d(std::string parname, unsigned int nsteps, double parmin, double parmax,
                             std::vector<pdf*> probs, std::vector<parameters*> params,
                             std::vector<std::vector<event> *> events, std::vector<std::string> common_par,
                             unsigned int bin, int from, int to){
    spdlog::info("Likelihood scan study for parameter "+parname);
    bool cache_minos = opts->minos_errors;
    bool cache_hesse = opts->hesse_postrun;
    bool cache_hesse2 = opts->squared_hesse;
    bool cache_shift = opts->shift_lh;
    //int cache_print = spdlog::default_logger_raw()->level();

    parameter* par = params.at(0)->get_parameter(parname);
    double cache_value = par->get_start_value();
    double cache_step_size = par->get_start_value();
    if (par != 0){
        opts->minos_errors = false;
        opts->hesse_postrun = false;
        opts->squared_hesse = false;
        opts->shift_lh = false;

        fitter f(opts);
        if(common_par.size() > 0) f.set_common_parameters(common_par);

        f.fit(probs, params, events);
        double lhdatafloated = f.likelihood();
        spdlog::info("Data fit finished with -2logLh = {0:f}", lhdatafloated);

        if (from < 0) from = 0;
        if (to > int(nsteps) || to < 0) to = int(nsteps);

        std::vector<double> lhdatafixed(to-from, 0.0);
        std::vector<double> statusdatafixed(to-from, 0.0);

        std::string pardesc = par->get_description();
        for (int j=from; j<to; j++){
            double parcur = parmin+(j+0.5)*(parmax-parmin)/nsteps;
            spdlog::info("Studying point " + parname + "={0:f}", parcur);
            for(unsigned int p=0; p<params.size(); p++){
                params.at(p)->reset_parameters();
                params.at(p)->get_parameter(parname)->init(parcur, parmin, parmax, 0.0);
            }
            statusdatafixed.at(j-from) = f.fit(probs, params, events);
            lhdatafixed.at(j-from) = f.likelihood();
        }
        //reactivate parameter for consecutive llh scans:
        for(unsigned int p=0; p<params.size(); p++){
            params.at(p)->get_parameter(parname)->init(cache_value, parmin, parmax, cache_step_size);
        }

        //save to root file
        std::ostringstream sout;
        sout << "llhscan_results_" << (opts->full_angular ? "full_angular" : "folding"+std::to_string(opts->folding)) << "_bin" << bin << ".root"; //TODO names

        TFile* fout = new TFile(sout.str().c_str(), "UPDATE");
        fout->cd();
        TTree* t = new TTree(parname.c_str(), pardesc.c_str());
        int step=0, statusdfixed=0;
        double value=0.0, lhdfloated=lhdatafloated, lhdfixed=0.0;
        t->Branch("bin",&bin,"bin/I");
        t->Branch("step",&step,"step/I");
        t->Branch("value",&value,"value/D");
        t->Branch("statusdatafixed",&statusdfixed,"statusdatafixed/I");
        t->Branch("lhdatafloated",&lhdfloated,"lhdatafloated/D");
        t->Branch("lhdatafixed",&lhdfixed,"lhdatafixed/D");
        for (int j=from; j<to; j++){
            step = j;
            value = parmin+(j+0.5)*(parmax-parmin)/nsteps;
            statusdfixed = statusdatafixed.at(j-from);
            lhdfixed = lhdatafixed.at(j-from);
            lhdfloated = lhdatafloated;
            t->Fill();
        }
        t->Write("",TObject::kWriteDelete);
        fout->Write();
        fout->Close();
        delete fout;
    }
    else{
        spdlog::critical("Could not find parameter "+parname);
        assert(0);
    }
    spdlog::info("Finished likelihood scan for parameter "+parname);
    opts->hesse_postrun = cache_hesse;
    opts->squared_hesse = cache_hesse2;
    opts->minos_errors = cache_minos;
    opts->shift_lh = cache_shift;
    return true;
}

bool fcnc::likelihoodscan::scan_1d(std::string parname, unsigned int nsteps,
                                   double parmin, double parmax,
                                   pdf* prob, parameters* params,
                                   std::vector<event> * events,
                                   unsigned int bin, int from, int to){
    std::vector<pdf *> the_probs;
    the_probs.push_back(prob);

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

    std::vector<std::vector<event>* > the_events;
    the_events.push_back(events);

    std::vector<std::string> common_par;
    common_par.clear();

    return scan_1d(parname, nsteps, parmin, parmax, the_probs, the_params, the_events, common_par, bin, from, to);
}


bool fcnc::likelihoodscan::scan_2d(std::string parxname, unsigned int nstepsx, double parxmin, double parxmax,
                             std::string paryname, unsigned int nstepsy, double parymin, double parymax,
                             std::vector<pdf*> probs, std::vector<parameters*> params,
                             std::vector<std::vector<event> *> events, std::vector<std::string> common_par,
                             unsigned int bin, int from, int to){
    spdlog::info("Likelihood scan study for parameters "+parxname+", "+paryname);
    bool cache_minos = opts->minos_errors;
    bool cache_hesse = opts->hesse_postrun;
    bool cache_hesse2 = opts->squared_hesse;
    bool cache_shift = opts->shift_lh;
    //int cache_print = spdlog::default_logger_raw()->level();

    parameter* parx = params.at(0)->get_parameter(parxname);
    parameter* pary = params.at(0)->get_parameter(paryname);
    if (parx != 0 && pary != 0){
        opts->minos_errors = false;
        opts->hesse_postrun = false;
        opts->squared_hesse = false;
        opts->shift_lh = false;

        fitter f(opts);
        if(common_par.size() > 0) f.set_common_parameters(common_par);

        f.fit(probs, params, events);
        double lhdatafloated = f.likelihood();
        spdlog::info("Data fit finished with -2logLh = {0:f}", lhdatafloated);

        if (from < 0) from = 0;
        if (to > int(nstepsx*nstepsy) || to < 0) to = nstepsx*nstepsy;

        std::vector<double> lhdatafixed(to-from, 0.0);
        std::vector<int > statusdatafixed(to-from, 0.0);

        std::string parxdesc = parx->get_description();
        std::string parydesc = pary->get_description();
        for (int j=from; j<to; j++){
            double parxcur = parxmin+((j%nstepsx)+0.5)*(parxmax-parxmin)/nstepsx;
            double parycur = parymin+((j/nstepsx)+0.5)*(parymax-parymin)/nstepsy;

            spdlog::info("Studying point {0:s}={1:f}, {2:s}={3:f}  idx={4:d}", parxname, parxcur, paryname,parycur, j);
            for(unsigned int p = 0; p < params.size(); p++){
                params.at(p)->reset_parameters();
                params.at(p)->get_parameter(parxname)->init(parxcur, parxmin, parxmax, 0.0);
                params.at(p)->get_parameter(paryname)->init(parycur, parymin, parymax, 0.0);
            }

            statusdatafixed.at(j-from) = f.fit(probs, params, events);
            lhdatafixed.at(j-from) = f.likelihood();
        }

        //save to root file
        std::ostringstream sout;
        sout << "llh2Dscan_results_" << (opts->full_angular ? "full_angular" : "folding"+std::to_string(opts->folding)) << "_bin" << bin << ".root"; //TODO: name
        TFile* fout = new TFile(sout.str().c_str(), "RECREATE");
        fout->cd();
        TTree* t = new TTree((parxname+std::string("_")+paryname).c_str(), (std::string(";")+parxdesc+std::string(";")+parydesc).c_str());
        int step=0, stepx=0, stepy=0, statusdfixed=0;
        double valuex=0.0, valuey, lhdfloated=lhdatafloated, lhdfixed=0.0;
        t->Branch("bin",&bin,"bin/I");
        t->Branch("step",&step,"step/I");
        t->Branch("stepx",&stepx,"stepx/I");
        t->Branch("stepy",&stepy,"stepy/I");
        t->Branch("valuex",&valuex,"valuex/D");
        t->Branch("valuey",&valuey,"valuey/D");
        t->Branch("statusdatafixed",&statusdfixed,"statusdatafixed/I");
        t->Branch("lhdatafloated",&lhdfloated,"lhdatafloated/D");
        t->Branch("lhdatafixed",&lhdfixed,"lhdatafixed/D");
        for (int j=from; j<to; j++){
            step = j;
            stepx = j%nstepsx;
            stepy = j/nstepsx;
            valuex = parxmin+((j%nstepsx)+0.5)*(parxmax-parxmin)/nstepsx;
            valuey = parymin+((j/nstepsx)+0.5)*(parymax-parymin)/nstepsy;
            lhdfixed = lhdatafixed.at(j-from);
            statusdfixed = statusdatafixed.at(j-from);
            lhdfloated = lhdatafloated;
            t->Fill();
        }
        t->Write();
        fout->Write();
        fout->Close();
        delete fout;
    }
    else{
        if (parx == 0) spdlog::info("Could not find parameter "+parxname);
        if (pary == 0) spdlog::info("Could not find parameter "+paryname);
        assert(0);
    }
    spdlog::info("Finished likelihood scan for parameters "+parxname+", "+paryname);
    opts->hesse_postrun = cache_hesse;
    opts->squared_hesse = cache_hesse2;
    opts->minos_errors = cache_minos;
    opts->shift_lh = cache_shift;
    return true;
}

bool fcnc::likelihoodscan::scan_2d(std::string parxname, unsigned int nstepsx, double parxmin, double parxmax,
                             std::string paryname, unsigned int nstepsy, double parymin, double parymax,
                             pdf* prob, parameters* params,
                             std::vector<event> * events, unsigned int bin, int from, int to){

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

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

    std::vector<std::vector<event>* > the_events;
    the_events.push_back(events);

    std::vector<std::string> common_par;
    common_par.clear();

    return scan_2d(parxname, nstepsx, parxmin, parxmax,
                   paryname, nstepsy, parymin, parymax,
                   the_probs, the_params, the_events, common_par, bin, from, to);
}