#include <ostream>
#include <iostream>
#include <fstream>
#include <string>
#include <string.h>
#include <vector>
#include "TTree.h"
#include "TFile.h"
#include "TLorentzVector.h" 
#include "TH1D.h"
#include "TH2.h"
#include "TProfile.h"
#include "TDirectoryFile.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TLegend.h"
#include "TVectorD.h"
#include "TColor.h"
#include "TString.h"
#include "RooDataSet.h"
#include "RooArgSet.h"
#include "RooAbsPdf.h"
#include "RooDoubleCBFast.h"
#include "RooGaussian.h"
#include "RooPlot.h" 
#include "RooBinning.h"
#include "RooDataHist.h"
#include "RooMomentMorph.h"
#include "RooRealVar.h"
#include "RooAddPdf.h"
#include "RooPolynomial.h"
#include "RooExtendPdf.h"
#include "RooWorkspace.h"
#include "RooAbsReal.h"
#include "AtlasLabels.h"
#include "AtlasUtils.h"
#include "RooChi2Var.h"


void polFit(RooDataSet& datasetSR, RooDataSet& datasetAR1, RooDataSet& datasetAR2, RooDataSet& datasetAR3,
RooRealVar& massSR, RooRealVar& massAR1, RooRealVar& massAR2, RooRealVar& massAR3,
RooPolynomial& pol2SR, RooPolynomial& pol2AR1, RooPolynomial& pol2AR2, RooPolynomial& pol2AR3,
RooRealVar& normSR, RooRealVar& normAR1, RooRealVar& normAR2, RooRealVar& normAR3, std::string sampleName) {  
  std::cout<<"beginning pol2fit()"<<std::endl;
  //setting up vectors to store values and easily loop over, as well as variables that will be used in the function
  std::vector<RooDataSet> datasets;
  std::vector<RooRealVar> masses;
  std::vector<RooPolynomial> polynomials;
  std::vector<RooRealVar> norms;
  std::vector<std::string> channels;
  datasets.push_back(datasetSR);datasets.push_back(datasetAR1);datasets.push_back(datasetAR2);datasets.push_back(datasetAR3);
  masses.push_back(massSR);masses.push_back(massAR1);masses.push_back(massAR2);masses.push_back(massAR3);
  polynomials.push_back(pol2SR);polynomials.push_back(pol2AR1);polynomials.push_back(pol2AR2);polynomials.push_back(pol2AR3);
  norms.push_back(normSR);norms.push_back(normAR1);norms.push_back(normAR2);norms.push_back(normAR3);
  //std::string channels[4]={"SR","AR1","AR2","AR3"};
  channels.push_back("SR");channels.push_back("AR1");channels.push_back("AR2");channels.push_back("AR3");
  int numberOfChannels = channels.size(); 
  std::string model_str="euu";
  std::cout<<"local variables setup"<<std::endl;
  //create TFile to save everything in
  TFile* f_out = new TFile("ttZAR_euu.root","recreate");
  f_out->cd();

  for(int i=0;i<numberOfChannels;i++){
    std::cout<<i<<","<<std::endl;
    //std::cout<<"Crash on polynomial fit?"<<std::endl;
    //std::cout<<datasets.at(i)<<std::endl;
    std::cout<<polynomials.at(i).getTitle()<<std::endl;
    polynomials.at(i).fitTo(datasets.at(i),RooFit::Range(12.0, 77.0),RooFit::SumW2Error(true));
    std::cout<<"polynomial fitted"<<std::endl;
    RooPlot* frame = masses.at(i).frame(RooFit::Range(12, 77));
    std::cout<<"RooPlot setup"<<std::endl;
    datasets.at(i).plotOn(frame);
    polynomials.at(i).plotOn(frame);
    //create TDirectoryFile to then save in TFile
    std::string TDirectoryName_str=channels.at(i)+"_"+model_str;
    const char* TDirectoryName_char=TDirectoryName_str.c_str(); 
    TDirectoryFile* TDirectory = new TDirectoryFile(TDirectoryName_char,TDirectoryName_char);
    std::cout<<"TDirectory setup"<<std::endl;
    polynomials.at(i).Print();
    TDirectory->cd();
    gROOT->SetStyle("ATLAS");
    TCanvas* canvas = new TCanvas();
    std::cout<<"TCanvas setup"<<std::endl;
    canvas->cd();
    frame->Draw(); 
    ATLASLabel(0.2, 0.88, "Work in Progress");
    std::string parsed_str=model_str+", "+sampleName+", "+channels.at(i)+", entries = %.3f";
    const char* parsed_char = parsed_str.c_str();
    myText(0.2, 0.82, kBlack, TString::Format(parsed_char, datasets.at(i).sumEntries()));
    myText(0.2, 0.76, kBlack, TString::Format("n = %.3f", norms.at(i).getVal())); 
    std::string savedName_str=sampleName+channels.at(i)+"_"+model_str;
    const char* savedName_char=savedName_str.c_str();
    canvas->SaveAs(savedName_char);     
    std::cout<<"TCanvas saved"<<std::endl;

/*    std::string workspaceName_str = "wpol2"+channels.at(i);
    const char* workspaceName_char=workspaceName_str.c_str();
    RooWorkspace* workspace = new RooWorkspace("wpol2SR", "workspace");
    workspace->import(polynomials.at(i));
    workspace->import(norms.at(i));
    workspace->Print();
    workspace->Write(); */ 
  }

  f_out->Close();

}


int test() {

      float weight = 0;
      float weight_fake = 0;
      float mass = 0;
      std::vector<float> *base_mu_Pt = nullptr;
      std::vector<float> *mu_corr_PflowTight_VarRad = nullptr;
      std::vector<float> *base_mu_d0 = nullptr;
      std::vector<float> *base_mu_d0Err = nullptr;
      bool isSR_OS = 0;
      bool isFakeSR_OS = 0;      

      std::string channelStrings[5]={"SR","AR1","AR2","AR3","extraStringToAvoidSegFault"};
      int numberOfChannels = sizeof(channelStrings)/sizeof(channelStrings[0]);  
      std::vector<RooRealVar>masses;
      std::vector<RooRealVar>weights;
      std::vector<RooArgSet>hs;
      std::vector<RooDataSet>datasets;
      masses.reserve(numberOfChannels);weights.reserve(numberOfChannels);hs.reserve(numberOfChannels);datasets.reserve(numberOfChannels);  
      for(int i=0;i< numberOfChannels;i++){
            std::string mass_str = "mass"+channelStrings[i];
            std::string weight_str = "weight"+channelStrings[i];
            std::string hs_str = "hs"+channelStrings[i];
            std::string dataset_str = "dataset"+channelStrings[i];
            const char* mass_char = mass_str.c_str();
            const char* weight_char = weight_str.c_str();
            const char* hs_char = hs_str.c_str();
            const char* dataset_char = dataset_str.c_str();
            RooRealVar mass(mass_char, "m_{#mu#mu}", 12,77);
            masses.push_back(mass);
            RooRealVar weight(weight_char, weight_char, 0);
            weights.push_back(weight);
            RooArgSet hs_var(hs_char);
            hs_var.add(mass);
            hs_var.add(weight);
            hs.push_back(hs_var);
            RooDataSet dataset(dataset_char, "", hs_var, RooFit::WeightVar(weight));
            datasets.push_back(dataset);
      }
      //get index for channels (doing this the long way seems uneccessary for such a short list, but it's useful in case other channels want to be added in the future)
      int indexSR=std::distance(channelStrings, std::find(channelStrings, channelStrings + sizeof(channelStrings)/sizeof(channelStrings[0]), "SR"));  
      int indexAR1=std::distance(channelStrings, std::find(channelStrings, channelStrings + sizeof(channelStrings)/sizeof(channelStrings[0]), "AR1"));  
      int indexAR2=std::distance(channelStrings, std::find(channelStrings, channelStrings + sizeof(channelStrings)/sizeof(channelStrings[0]), "AR2"));  
      int indexAR3=std::distance(channelStrings, std::find(channelStrings, channelStrings + sizeof(channelStrings)/sizeof(channelStrings[0]), "AR3"));        
      std::vector<RooRealVar> norms;
      std::vector<RooPolynomial> polynomials;
      norms.reserve(4);polynomials.reserve(4);

      std::cout<<"initializing local variables for data"<<std::endl;
      RooRealVar p0("p0", "p0", 0.1, -5., 5.);
      RooRealVar p1("p1", "p1", 0.1, -5., 5);
      RooRealVar p2("p2", "p2", 0.1, -5., 5);
      RooRealVar p0AR1("p0AR1", "p0AR1", 0.1, -5., 5.);
      RooRealVar p1AR1("p1AR1", "p1AR1", 0.1, -5., 5);
      RooRealVar p2AR1("p2AR1", "p2AR1", 0.1, -5., 5);
      RooRealVar p0AR2("p0AR2", "p0AR2", 0.1, -5., 5.);
      RooRealVar p1AR2("p1AR2", "p1AR2", 0.1, -5., 5.);
      RooRealVar p2AR2("p2AR2", "p2AR2", 0.1, -5., 5.);
      RooRealVar p0AR3("p0AR3", "p0AR3", 0.1, -5., 5.);
      RooRealVar p1AR3("p1AR3", "p1AR3", 0.1, -5., 5.);
      RooRealVar p2AR3("p2AR3", "p2AR3", 0.1, -5., 5.);
      //create variables and push in vectors to be used later            
      std::cout<<"initializing variables to be pushed in vectors for data"<<std::endl;
      RooRealVar normSR("normSR", "normSR", datasets.at(indexSR).sumEntries());
      RooRealVar normAR1("normAR1", "normAR1", datasets.at(indexAR1).sumEntries());
      RooRealVar normAR2("normAR2", "normAR2", datasets.at(indexAR2).sumEntries());
      RooRealVar normAR3("normAR3", "normAR3", datasets.at(indexAR3).sumEntries());
      //std::cout<<normSR<<std::endl;
      norms.push_back(normSR);norms.push_back(normAR1);norms.push_back(normAR2);norms.push_back(normAR3);
      RooPolynomial pol2SR("pol2SR", "pol2SR", masses.at(indexSR), RooArgList(p0,p1));
      RooPolynomial pol2AR1("pol2AR1", "pol2AR1", masses.at(indexAR1), RooArgList(p0AR1)); 
      RooPolynomial pol2AR2("pol2AR2", "pol2AR2", masses.at(indexAR2), RooArgList());
      RooPolynomial pol2AR3("pol2AR3", "pol2AR3", masses.at(indexAR3), RooArgList(p0AR3,p1AR3));
      //std::cout<<pol2SR.evaluate()<<std::endl;  
      polynomials.push_back(pol2SR);polynomials.push_back(pol2AR1);polynomials.push_back(pol2AR2);polynomials.push_back(pol2AR3);
      std::cout<<"everything setup for data"<<std::endl;

      polFit(datasets.at(indexSR), datasets.at(indexAR1), datasets.at(indexAR2), datasets.at(indexAR2), masses.at(indexSR), masses.at(indexAR1), masses.at(indexAR2), masses.at(indexAR3), polynomials.at(indexSR), polynomials.at(indexAR1), polynomials.at(indexAR2), polynomials.at(indexAR3), norms.at(indexSR), norms.at(indexAR1), norms.at(indexAR2), norms.at(indexAR3), "data");

      return 0;
}