#include <THStack.h>
#include <TH3.h>
#include <TH2.h>
#include <TH1.h>
#include <TF1.h>
#include <TProfile.h>
#include <TObjArray.h>
#include <TObject.h>
#include <TString.h>
#include <TFile.h>
#include <THStack.h>
#include <TLegend.h>
#include <TCanvas.h>
#include <TGraph.h>
#include <TMultiGraph.h>
#include <TList.h>
#include <TStyle.h>
#include <TColor.h>
#include <TROOT.h>
#include <iostream>
#include <vector>
#include <set>
#include <cmath>
#include <iostream>
#include "../scripts/palettes.C"

using namespace std;

void HandleTH1(TH1* hist, double normalisation, int rebin_x, int ){
    if( normalisation!=1 ){
        hist->Scale(1./normalisation);
    }
    if(rebin_x >0 ){
        hist->RebinX(rebin_x);
    }
}

void HandleTH2(TH2* hist, double normalisation, int rebin_x, int rebin_y){
    if( normalisation!=1 ){
        hist->Scale(1./normalisation);
    }
    if(rebin_x >0 ){
        hist->RebinX(rebin_x);
    }
    if(rebin_y >0 ){
        hist->RebinY(rebin_y);
    }
}

void HandleTH3(TH3* hist, double normalisation, int rebin_x, int rebin_y){
    HandleTH2((TH2*)hist,normalisation,rebin_x,rebin_y);
}

TLegend* HandleStack(THStack* stack, double normalisation, int rebin_x, int rebin_y){
  TLegend* leg=new TLegend(0.75,.75,0.9,0.9);
  leg->SetFillColor(0);
    TList* all_hists=stack->GetHists();
    TIterator* iterator=all_hists->MakeIterator();
    TObject* object=NULL;
    bool first=true;
    while((object=iterator->Next())){
        if( object->InheritsFrom("TH2")){
            HandleTH2((TH2*)object,normalisation,rebin_x,rebin_y);
        } else if( object->InheritsFrom("TH1")){
            HandleTH1((TH1*)object,normalisation,rebin_x,rebin_y);
        }else continue;

        TH1* hist=(TH1*) object;
        TIter next_funct(hist->GetListOfFunctions());
        while(TObject* funct=next_funct()){
            funct->ls();
            TNamed* named=dynamic_cast<TNamed*>(funct);
            if(named){
                if(TString(named->GetName())=="legend"){
                    TString title=named->GetTitle();
                    leg->AddEntry(hist,title.Data(),"l");
                    cout<<title<<endl;
                }
            }
        }
        if(first){
            first=false;
            hist->GetListOfFunctions()->Add(leg);
        }
    }
    leg->SetX1(0.1); leg->SetX2(0.4);
    leg->SetY1(0.1); leg->SetY2(0.4);
    leg->Draw();
    cout<<leg->GetX1()<<", "<<leg->GetX2()<<endl;
    cout<<leg->GetX1NDC()<<", "<<leg->GetX2NDC()<<endl;
    return leg;
}

int PlotOneSelected(
        const char* histogram_file,
        const char* hist_name,
        const char* draw,
        const char* outFileName,
        const double  normalisation=1,
        const int rebin_x=1,
        const int rebin_y=1
        ){

    // Open the file
    TFile* infile=TFile::Open(histogram_file);
    if(!infile) {
        cout<<"Error: Cannot open file: "<<histogram_file<<endl;
        return 1;
    }

    // Get the source histogram
    TNamed* hist_obj=(TNamed*)infile->Get(hist_name);
    if(!hist_obj) {
        cout<<"Error: Cannot find histogram ("<<hist_name<<") in file: "<<histogram_file<<endl;
        return 2;
    }

    hist_obj->Draw(draw);

    TLegend* leg=NULL;
    if(hist_obj->InheritsFrom("THStack")){
        leg=HandleStack((THStack*)hist_obj,normalisation,rebin_x,rebin_y);
    }else if(hist_obj->InheritsFrom("TH3")){
        HandleTH3((TH3*)hist_obj,normalisation,rebin_x,rebin_y);
    }else if(hist_obj->InheritsFrom("TH2")){
        HandleTH2((TH2*)hist_obj,normalisation,rebin_x,rebin_y);
    }else if(hist_obj->InheritsFrom("TH1")){
        HandleTH1((TH1*)hist_obj,normalisation,rebin_x,rebin_y);
    }

    hist_obj->Draw(draw);
    if(leg) leg->Draw();

    gPad->SetGridy();
    gPad->SetGridx();

    gPad->Print(Form("%s.png" ,outFileName));
    gPad->Print(Form("%s.root",outFileName));

    return 0;
}

int PlotSelected(
        const TString histogram_files,
        const char* hist_name,
        const char* draw,
        const TString outFileName,
        const TString  normalisation_str="1",
        const int rebin_x=1,
        const int rebin_y=1,
        const TString value_prefix="analysis-Tor1_",
        const TString value_postfix="-Tor2_1.50_profile"
        ){
    SetUserPalette(1,255,0.5);

    TObjArray* names=histogram_files.Tokenize(" ");

    TObjArray* normalisation_tkns=normalisation_str.Tokenize(" ");
    double normalisation=1;
    bool normPerFile=false;
    if(normalisation_tkns->GetEntries()==1){
        normalisation=((TObjString*)normalisation_tkns->At(0))->GetString().Atof();
    }else if(normalisation_tkns->GetEntries()!=0){
        if (normalisation_tkns->GetEntries()!=names->GetEntries()){
            cout<<"Error: Number of normalisation tokens must be either 0, 1 (a single value for all files), or the same as the number of files"<<endl;
            return 1;
        }
        normPerFile=true;
    }

    int count=0;
    TObject* object=NULL;
    TIterator* iterator=names->MakeIterator();
    //TIterator* it_out=out_names->MakeIterator();
    while((object=iterator->Next())){
        //TObject* outputName=it_out->Next();
        TString outName=object->GetName();
        outName.Remove(0,value_prefix.Length());
        outName.Remove(outName.Length()-value_postfix.Length());
        if(outName==object->GetName()) outName=outFileName;
        else outName=outFileName+outName;

        if(normPerFile) {
            normalisation=((TObjString*)normalisation_tkns->At(count))->GetString().Atof();
            cout<<"Set normalisation to: "<<normalisation<<endl;
        }

        PlotOneSelected(object->GetName(),hist_name,draw,outName.Data(),normalisation,rebin_x,rebin_y);
        ++count;
    }
    return 0;
}