//Renata Kopecna

#include <TCanvas.h>
#include <TPaveText.h>
#include <TTree.h>
#include <TBranch.h>
#include <TF1.h>
#include <TH1D.h>
#include <TStyle.h>

#include <string>
#include <vector>
#include <event.hh>
#include <parameters.hh>
#include <funcs.hh>
#include <pdf.hh>
#include <options.hh>

#include <multifit.hh>

using namespace std;
using namespace fcnc;

multifit::multifit(options* o):
    opts(o)
{
    output = new TFile(("result" + opts->name + ".root").c_str(), "RECREATE");
};

multifit::~multifit(){
    output->Write();
    output->Close();
    delete output;
};

void multifit::fit_gaussian(vector<double>& values, double& gauss_mean, double& sigma_gauss_mean, double& gauss_width, double& sigma_gauss_width, double& chi_squared) const
{

    double max = *max_element(std::begin(values), std::end(values));
    double min = *min_element(std::begin(values), std::end(values));

    double mean=0.0;
    double rms=0.0;
    TTree* t = new TTree("tree", "pull values");
    double value;
    TBranch *branch = t->Branch("value",&value,"value/D");
    for (unsigned int i = 0; i < values.size(); i++) {
        value = values.at(i);
        branch->Fill();
        t->Fill();
        mean += value/values.size();
    }
    for (unsigned int i = 0; i < values.size(); i++){
        value = values.at(i);
        rms += (value - mean)*(value - mean)/values.size();
    }
    rms = sqrt(rms);
    TF1* f1 = new TF1("f1", "gaus(0)/sqrt(2.0*3.1415926536)/abs([2])", min-5.0, max+5.0);
    f1->SetParameters(1.0,mean,rms);
    //f1->SetParLimits(0, 1.0, 1.0);
    f1->FixParameter(0, 1.0);
    t->UnbinnedFit("f1", "value");
    gauss_mean = f1->GetParameter(1);
    sigma_gauss_mean = f1->GetParError(1);
    gauss_width = f1->GetParameter(2);
    sigma_gauss_width = f1->GetParError(2);
    delete f1;
    delete t;
};

void multifit::update_pull(parameter* p, vector<double>& values, vector<double>& errors, double& gauss_mean, double& sigma_gauss_mean, double& gauss_width, double& sigma_gauss_width, double& chi_squared, string appendix) const
{
    unsigned int runs=values.size();
    string parname(p->get_name());
    string descr(p->get_description());
    gStyle->SetOptStat(0);
    gStyle->SetOptFit(0);
    //finally the pull histo
    TH1D* pull_histo = new TH1D((parname + "_pull_").c_str(),
                                (descr + " pull distribution;" + "(" + descr
                                 + "^{fitted}-" + descr
                                 + "^{generated})/#sigma").c_str(),
                                50, -5.0, 5.0);//sensible for pulls
    //fill histos
    //this is done every time because the time is negligible
    double start_value = p->get_start_value();
    vector<double> pull_values;
    for (unsigned int k=0; k<runs; k++){
        if (errors.at(k) != 0.0){
            double pull = (values.at(k)-start_value)/errors.at(k);
            pull_histo->Fill(pull);
            pull_values.push_back(pull);
        }
    }
    fit_gaussian(pull_values, gauss_mean, sigma_gauss_mean, gauss_width, sigma_gauss_width, chi_squared);

    TCanvas *pull_canvas = new TCanvas((parname + "_pull_" + appendix).c_str(), "Bs Likelihood Analysis: Pull distribution",1600,1200);
    pull_canvas->cd();
    pull_histo->Draw();
    TF1 *g = new TF1("g","[0]*exp(-(x-[1])*(x-[1])/2.0/[2]/[2])",-5.0,5.0);

    //for normalization: multiply with dx*histint = width/Nbins*histint
    g->SetParameters(1.0/sqrt(2.0*TMath::Pi())/gauss_width*10.0/50.0*pull_histo->Integral(), gauss_mean, gauss_width);
    g->Draw("same");
    TPaveText* text = new TPaveText(0.6,0.7,0.88,0.88,"NDC");
    text->SetFillColor(0);
    std::ostringstream stream_mean;
    stream_mean << "Mean: " << fixed << setprecision(3) << gauss_mean << "#pm" << sigma_gauss_mean;
    std::ostringstream stream_width;
    stream_width << "Width: " << fixed << setprecision(3) << gauss_width << "#pm" << sigma_gauss_width;
    text->AddText(stream_mean.str().c_str());
    text->AddText(stream_width.str().c_str());
    text->Draw("same");

    string afilename; //TODO: move to paths
    if (appendix != "") afilename = "plots/pull_" + parname + "_" + appendix + ".eps";
    else                afilename = "plots/pull_" + parname + ".eps";
    if (opts->write_eps)pull_canvas->Print(afilename.c_str(), "eps");
    output->cd();
    pull_canvas->Write();
    delete text;
    delete g;
    delete pull_histo;
    delete pull_canvas;
}

void multifit::update_value(parameter * p, vector<double>& values, vector<double>& errors, double& gauss_mean, double& sigma_gauss_mean, double& gauss_width, double& sigma_gauss_width, double& chi_squared, string appendix) const
{
    unsigned int runs=values.size();
    string parname(p->get_name());
    string descr(p->get_description());
    //first find min and max values
    double max = *max_element(std::begin(values), std::end(values));
    double min = *min_element(std::begin(values), std::end(values));
    double width = max-min;
    double hist_min = min - 0.5*width;
    double hist_max = max + 0.5*width;
    gStyle->SetOptStat(0);
    gStyle->SetOptFit(0);
    TH1D* value_histo = new TH1D((parname+"_values").c_str(),
                                 (descr + ";" + descr).c_str(),
                                 50, hist_min, hist_max);
    for (unsigned int k=0; k<runs; k++) {
        value_histo->Fill(values.at(k));
    }

    fit_gaussian(values, gauss_mean, sigma_gauss_mean, gauss_width, sigma_gauss_width, chi_squared);
    TCanvas *value_canvas = new TCanvas((parname + "_value_" + appendix).c_str(), "Bs Likelihood Analysis: Value distribution",1600,1200);
    value_canvas->cd();
    value_histo->Draw();
    TF1 *g = new TF1("g","[0]*exp(-(x-[1])*(x-[1])/2.0/[2]/[2])",-5.0,5.0);
    g->SetParameters(1.0/sqrt(2.0*TMath::Pi())/gauss_width*(hist_max-hist_min)/50.0*value_histo->Integral(), gauss_mean, gauss_width);
    g->Draw("same");

    TPaveText* text = new TPaveText(0.6,0.7,0.88,0.88,"NDC");
    text->SetFillColor(0);
    std::ostringstream stream_mean;
    stream_mean << "Mean: " << fixed << setprecision(3) << gauss_mean << "#pm" << sigma_gauss_mean;
    std::ostringstream stream_width;
    stream_width << "Width: " << fixed << setprecision(3) << gauss_width << "#pm" << sigma_gauss_width;
    text->AddText(stream_mean.str().c_str());
    text->AddText(stream_width.str().c_str());
    text->Draw("same");

    //TODO: move
    string afilename;
    if (appendix != "") afilename = "plots/value_" + parname + "_" + appendix + ".eps";
    else                afilename = "plots/value_" + parname + ".eps";
    if (opts->write_eps) value_canvas->Print(afilename.c_str(), "eps");
    output->cd();
    value_canvas->Write();
    delete text;
    delete value_histo;
    delete value_canvas;
    delete g;
}

void multifit::update_error(parameter* p, vector<double>& values, vector<double>& errors, double& gauss_mean, double& sigma_gauss_mean, double& gauss_width, double& sigma_gauss_width, double& chi_squared, string appendix) const
{
    unsigned int runs=values.size();
    string parname(p->get_name());
    string descr(p->get_description());
    //then find min and max errors
    double max = *max_element(std::begin(errors), std::end(errors));
    double min = *min_element(std::begin(errors), std::end(errors));

    double width = max - min;
    double hist_min = min - 0.5*width;
    double hist_max = max + 0.5*width;
    gStyle->SetOptStat(0);
    gStyle->SetOptFit(0);
    TH1D* error_histo = new TH1D((parname+"_errors").c_str(),
                                 ("#sigma(" + descr + ");#sigma(" + descr + ")").c_str(),
                                 100, hist_min, hist_max);
    for (unsigned int k=0; k<runs; k++){
        error_histo->Fill(errors.at(k));
    }
    fit_gaussian(errors, gauss_mean, sigma_gauss_mean, gauss_width, sigma_gauss_width, chi_squared);
    TCanvas *error_canvas = new TCanvas((parname + "_error_" + appendix).c_str(), "Bs Likelihood Analysis: Error distribution",1600,1200);

    error_canvas->cd();
    error_histo->Draw();
    TF1 *g = new TF1("g","[0]*exp(-(x-[1])*(x-[1])/2.0/[2]/[2])",-5.0,5.0);
    g->SetParameters(1.0/sqrt(2.0*TMath::Pi())/gauss_width*(hist_max-hist_min)/100.0*error_histo->Integral(), gauss_mean, gauss_width);
    g->Draw("same");
    TPaveText* text = new TPaveText(0.6,0.7,0.88,0.88,"NDC");
    text->SetFillColor(0);
    std::ostringstream stream_mean;
    stream_mean << "Mean: " << fixed << setprecision(3) << gauss_mean << "#pm" << sigma_gauss_mean;
    std::ostringstream stream_width;
    stream_width << "Width: " << fixed << setprecision(3) << gauss_width << "#pm" << sigma_gauss_width;
    text->AddText(stream_mean.str().c_str());
    text->AddText(stream_width.str().c_str());
    text->Draw("same");

    //TODO: move
    string afilename;

    if (appendix != "") afilename = "plots/error_" + parname + "_" + appendix + ".eps";
    else                afilename = "plots/error_" + parname + ".eps";

    if (opts->write_eps) error_canvas->Print(afilename.c_str(), "eps");

    output->cd();
    error_canvas->Write();
    delete text;
    delete error_histo;
    delete error_canvas;
    delete g;
}