#include "../GlobalFunctions.hh"
#include "../HeidelbergFitter/LHCbStyle.h"

using namespace std;
using namespace RooFit ;

//define mass regions:
const UInt_t N = 2;
//const Double_t low[N]  = {5170., 5330., 5500.};
//const Double_t high[N] = {5230., 5500., 5700.};
const Double_t low[N]  = {5170., 5350.};
const Double_t high[N] = {5210., 5700.};

std::vector<double>values[3][N];
std::vector<double>errors[3][N];

std::string angles_latex[3] = {"cos(#Theta_{L})", "cos(#Theta_{K})", "#phi"};
std::string angles_names[3] = {"ctl", "ctk", "phi"};

Int_t BkgCheck(Int_t Run = 1, std::string DDLL = "DD") {

    if(Run != 1 && Run != 2 && Run != 12){
        std::cout << "[ERROR]\t\tInvalid Run number given: " << Run  << ". Exit program!" << std::endl;
        return 0;
    }

    const bool ScanChebyMaxOrder = false;
    const bool IncludeChebyFit = false;
    const UInt_t MaxOrderChebyPol = 3;

    gStyle -> SetOptStat(0);
    LHCbStyle();
    gROOT->SetBatch(kTRUE);
    TH1::SetDefaultSumw2(kTRUE);
    //load tree from FCNC file
    TChain * tree = new TChain("Events");
    tree->Add(Form("%s/data/*Run%i%s_FCNC.root",path_to_data.c_str(),Run, KshortChannel && SplitDDandLL ? ("_"+DDLL).c_str() : ""));

    if(tree->GetEntries() == 0){
        std::cout << "No entries found in file!" << std::endl;
        return 0;
    }

    //assign values to angles and mass
    double angles[3], m;
    tree->SetBranchAddress("m" , &m);
    tree->SetBranchAddress("costhetal" , &angles[0]);
    tree->SetBranchAddress("costhetak" , &angles[1]);
    tree->SetBranchAddress("phi" , &angles[2]);

    std::vector<TH1D*> h[3];
    const UInt_t nBins = 10;

    //create histograms
    for(UInt_t a = 0; a < 3; a++){
        for(UInt_t n = 0; n < N; n++){
           values[a][n].clear();
           errors[a][n].clear();
           TH1D* hist = new TH1D((angles_names[a]+"_"+std::to_string(low[n])+"-"+std::to_string(high[n])+"_MeV").c_str(),
                                  Form("%.1f-%.1f MeV;%s;norm. Events", low[n], high[n], angles_latex[a].c_str()),
                                  nBins, a == 2 ? -TMath::Pi() : -1., a == 2 ? +TMath::Pi() : +1.);
           h[a].push_back(hist);
        }
    }

    //fill all histograms with events, according to mass bin
    for(UInt_t e = 0; e < tree->GetEntries(); e++){
        tree->GetEntry(e);
        for(UInt_t n = 0; n < N; n++){
            if(m > low[n] && m <= high[n]){
                for(UInt_t a = 0; a < 3; a++){
                    h[a].at(n)->Fill(angles[a]);
                }
                continue;
            }
        }
    }

    //put event numbers into hist names
    for(UInt_t a = 0; a < 3; a++){
      for(UInt_t n = 0; n < N; n++){
        h[a].at(n)->SetTitle(Form("%s (%.0f Events)", h[a].at(n)->GetTitle(), h[a].at(n)->Integral()));
      }
    }

    //make .eps plots:
    TCanvas * c = new TCanvas();
    c->cd();
    for(UInt_t a = 0; a < 3; a++){
        std::vector<TF1*> ff;
        c->Clear();
        for(UInt_t n = 0; n < N; n++){
            h[a].at(n)->Scale(1./h[a].at(n)->Integral()*nBins);
            h[a].at(n)->GetYaxis()->SetRangeUser(0, 2.5);///nBins);
            h[a].at(n)->SetLineColor(n+1);
            h[a].at(n)->SetMarkerStyle(n+20);
            h[a].at(n)->SetMarkerColor(n+1);
            h[a].at(n)->SetFillStyle(n == 1 ? 3345 : (n == 2 ? 3354 : 1));
            h[a].at(n)->SetFillColor(n+1);
            h[a].at(n)->Draw(n == 0 ? "PE3" : "PE3SAME");

            //Fit Chebyshev polynomials to histograms:
/*
            std::string fitformular = "[0]*cheb0";
            for(UInt_t ch = 1; ch <= MaxOrderChebyPol; ch++){
//              fitformular.append(TString::Format(" + [%d]*cheb%d", ch, ch));
              fitformular = TString::Format("cheb%d", ch, ch);
            }
            TF1 * f = new TF1(("f_"+angles_names[a]+"_"+std::to_string(n)).c_str(),fitformular.c_str(), a == 2 ? -TMath::Pi() : -1., a == 2 ? +TMath::Pi() : +1.);
            f->SetLineColor(n+1);
            for(UInt_t ch = 0; ch <= MaxOrderChebyPol; ch++){
              f->SetParameter(ch, 0.8);
            }
            h[a].at(n)->Fit(f, "R+");
            ff.push_back(f);
*/

          UInt_t cheb = 1;
          if(ScanChebyMaxOrder){
            while(cheb <= MaxOrderChebyPol){
              TF1 * f = new TF1(("f_"+angles_names[a]+"_"+std::to_string(n)).c_str(), TString::Format("cheb%d", cheb), a == 2 ? -TMath::Pi() : -1., a == 2 ? +TMath::Pi() : +1.);
              f->FixParameter(0, 1.);
              h[a].at(n)->Fit(f, "RQ");
              if(TMath::Abs(f->GetParameter(cheb)) < TMath::Abs(f->GetParError(cheb))){
                cheb--;
                break;
              }
              cheb++;
            }
          }
          else{
            cheb = MaxOrderChebyPol;
          }

          //fit Chebyshev with according max order:
          if(IncludeChebyFit){
            TF1 * f = new TF1(("f_"+angles_names[a]+"_"+std::to_string(n)).c_str(), TString::Format("cheb%d", cheb), a == 2 ? -TMath::Pi() : -1., a == 2 ? +TMath::Pi() : +1.);
            f->SetLineColor(n+1);
            f->FixParameter(0, 1.);
            h[a].at(n)->Fit(f, "RQ+");
            for(UInt_t i = 0; i < cheb; i++){
              values[a][n].push_back(f->GetParameter(i));
              errors[a][n].push_back(f->GetParError(i));
            }
            ff.push_back(f);
            h[a].at(n)->GetListOfFunctions()->Clear();
            h[a].at(n)->SetTitle(Form("%s (Cheby %d)", h[a].at(n)->GetTitle(), cheb));
          }
        }
        c->BuildLegend(0.25, 0.75, 0.75, 0.92);
        for(UInt_t i = 0; i < ff.size(); i++)
          ff.at(i)->Draw("SAME");
        c->SaveAs((angles_names[a]+"_bkg_Run"+std::to_string(Run)+(KshortChannel && SplitDDandLL ? "_"+DDLL : "")+".eps").c_str());
    }
    delete c;
    for(UInt_t a = 0; a < 3; a++){
      for(UInt_t n = 0; n < N; n++){
        delete h[a].at(n);
      }
    }
    return 1;
}

Int_t BackgroundCheck(){

  if(BkgCheck(1, "DD") == 0){
    return 0;
  }
  std::cout << "/=====================\\" << std::endl;
  std::cout << "|  Results Run 1, DD: |" << std::endl;
  std::cout << "\\=====================/" << std::endl;
  for(UInt_t a = 0; a < 3; a++){
    std::cout << ">> " << angles_names[a] << " <<" << std::endl;
    for(UInt_t v = 0; v < values[a][0].size(); v++){
      double ave_value = 0.0, ave_error = 0.0;
      for(UInt_t n = 0; n < N; n++){
        assert(values[a][n].size() == values[a][0].size());
        ave_value += values[a][n].at(v) / N;
        ave_error += TMath::Power(errors[a][n].at(v), 2);
      }
      ave_error = TMath::Sqrt(ave_error)/N;
      std::cout << "cheb" << v << ":\t" << ave_value << " +/- " << ave_error << std::endl;
    }
  }
  if(BkgCheck(2, "DD") == 0){
    return 0;
  }
  std::cout << "/=====================\\" << std::endl;
  std::cout << "|  Results Run 2, DD: |" << std::endl;
  std::cout << "\\=====================/" << std::endl;
  for(UInt_t a = 0; a < 3; a++){
    std::cout << ">> " << angles_names[a] << " <<" << std::endl;
    for(UInt_t v = 0; v < values[a][0].size(); v++){
      double ave_value = 0.0, ave_error = 0.0;
      for(UInt_t n = 0; n < N; n++){
        assert(values[a][n].size() == values[a][0].size());
        ave_value += values[a][n].at(v) / N;
        ave_error += TMath::Power(errors[a][n].at(v), 2);
      }
      ave_error = TMath::Sqrt(ave_error)/N;
      std::cout << "cheb" << v << ":\t" << ave_value << " +/- " << ave_error << std::endl;
    }
  }

  if(KshortChannel && SplitDDandLL){
    if(BkgCheck(1, "LL") == 0){
      return 0;
    }
    std::cout << "/=====================\\" << std::endl;
    std::cout << "|  Results Run 1, LL: |" << std::endl;
    std::cout << "\\=====================/" << std::endl;
    for(UInt_t a = 0; a < 3; a++){
      std::cout << ">> " << angles_names[a] << " <<" << std::endl;
      for(UInt_t v = 0; v < values[a][0].size(); v++){
        double ave_value = 0.0, ave_error = 0.0;
        for(UInt_t n = 0; n < N; n++){
          assert(values[a][n].size() == values[a][0].size());
          ave_value += values[a][n].at(v) / N;
          ave_error += TMath::Power(errors[a][n].at(v), 2);
        }
        ave_error = TMath::Sqrt(ave_error)/N;
        std::cout << "cheb" << v << ":\t" << ave_value << " +/- " << ave_error << std::endl;
      }
    }
    if(BkgCheck(2, "LL") == 0){
      return 0;
    }
    std::cout << "/=====================\\" << std::endl;
    std::cout << "|  Results Run 2, LL: |" << std::endl;
    std::cout << "\\=====================/" << std::endl;
    for(UInt_t a = 0; a < 3; a++){
      std::cout << ">> " << angles_names[a] << " <<" << std::endl;
      for(UInt_t v = 0; v < values[a][0].size(); v++){
        double ave_value = 0.0, ave_error = 0.0;
        for(UInt_t n = 0; n < N; n++){
          assert(values[a][n].size() == values[a][0].size());
          ave_value += values[a][n].at(v) / N;
          ave_error += TMath::Power(errors[a][n].at(v), 2);
        }
        ave_error = TMath::Sqrt(ave_error)/N;
        std::cout << "cheb" << v << ":\t" << ave_value << " +/- " << ave_error << std::endl;
      }
    }
  }
  return 1;
}