#include <TH2.h>
#include <TF1.h>
#include <TProfile.h>
#include <TObjArray.h>
#include <TObjString.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 "../scripts/palettes.C"
#include "../scripts/PlotHeightVsBeamline.C"
using namespace std;

struct Histogram_t{
    Histogram_t():norm_hist(NULL),norm_value(1){}
    TH1* hist;
    TH1* norm_hist;
    double x;
    double norm_value;
    double Integrate(TH1* histogram,double low,double high){
        if(!histogram) return -1;
        int bin_low=histogram->FindBin(low);
        int bin_high=histogram->FindBin(high);
        return histogram->Integral(bin_low,bin_high);
    }
    void Integrate(double start, double low_edge, double high_edge,double maximum,double& lower_int,double& middle_int,double& upper_int, TH1* histogram=NULL){
        if(!histogram) histogram=hist;
        lower_int =Integrate(histogram,start,low_edge)/norm_value;
        middle_int=Integrate(histogram,low_edge,high_edge)/norm_value;
        upper_int =Integrate(histogram,high_edge,maximum)/norm_value;
    }

    void IntegrateNormalisation(double start, double low_edge, double high_edge,double maximum,double& lower_int,double& middle_int,double& upper_int){
        if(norm_hist){
            Integrate(start,low_edge,high_edge,maximum,lower_int,middle_int,upper_int,norm_hist);
        }
        else {
            lower_int =1;
            middle_int=1;
            upper_int =1; 
        }
    }
};

void DrawPrint(TMultiGraph* graph, TLegend* legend, const char* outFileName,const char* file_suffix, TString save_method){
    if(save_method=="draw"){
        graph->Draw("AL");
        legend->Draw();
        gPad->SetGridy();
        gPad->SetGridx();

        gPad->Print(Form("%s_%s.png" ,outFileName,file_suffix));
        gPad->Print(Form("%s_%s.root",outFileName,file_suffix));
    }else if(save_method=="save"){
        graph->SaveAs(Form("%s_%s.root" ,outFileName,file_suffix));
    }
}

void PlotIntegrals(
        TString histogram_files,
        TString integral_limits,
        const char* title="Signal Particle Height as a Function of Dipole Field",
        const char* hist_name="truth_IPhaseIISignalAcceptanceAnalyzer/hBeamHeightPrimary",
        TString  normalisation_str="",
        const char* outFileName="overlaid",
        const char* value_format="analysis-Tor1_",
        const char* value_function="",
        double max_field=2.5,
        const char* save_method="draw",
        double min_field=2.5
        ){

    std::set<double> integral_bounds;
    TObjArray* limits=integral_limits.Tokenize(" ");
    if(limits->GetEntries()<2 ){
        cout<<"Error: I need at least 2 integration limits"<<endl;
        return;
    }
    TIterator* iterator=limits->MakeIterator();
    TObject* object=NULL;
    while((object=iterator->Next())){
        TString lim=object->GetName();
        cout<<"Integral limit: "<<lim <<" ("<< lim.Atof()<<")"<<endl;
        integral_bounds.insert(lim.Atof());
        max_field=std::max(lim.Atof(),max_field);
        min_field=std::min(lim.Atof(),min_field);
    }

    PlotLooper plotter(histogram_files,hist_name,normalisation_str,value_format,max_field,min_field,"projectionX");
    plotter.SetValueFunction(value_function);
    bool ok=plotter.Initiailise();
    if(!ok) return;

    std::vector<Histogram_t> histograms;
    for(int i_hist=0; i_hist< plotter.GetNumPlots(); ++i_hist){
        PlotLooper::Plot plot=plotter.GetPlotWithoutNorm(i_hist);
        if(!plot.histogram) {
            cout<<"no histogram found"<<endl;
            continue;
        }

        Histogram_t histogram;
        histogram.hist=plot.histogram;
        histogram.x=plotter.ConvertValue(plot.value);
        plotter.GetNormalisation(i_hist,histogram.norm_value,histogram.norm_hist);

        histograms.push_back(histogram);
    }

    cout<<"Loaded all plots"<<endl;

    std::set<double>::const_iterator i_high=integral_bounds.begin();
    std::set<double>::const_iterator i_low=i_high++;
    TLegend* legend=new TLegend(0.1,0.7,0.3,0.9);
    legend->SetFillColor(kWhite);
    TMultiGraph *adjacentStep=new TMultiGraph("adjacentStep",title);
    TMultiGraph *toZero=new TMultiGraph("toZero",title);
    TMultiGraph *toMax=new TMultiGraph("toMax",title);
    TMultiGraph *ratios=new TMultiGraph("ratios",title);
    double max_colour_field=1.05*max_field;
    for(; i_high != integral_bounds.end() && i_low !=integral_bounds.end();
            ++i_high){
        TGraph* stepGraph=new TGraph();
        TGraph* toZeroGraph=new TGraph();
        TGraph* toMaxGraph=new TGraph();
        TGraph* ratiosGraph=new TGraph();
        int count=0;
        for(std::vector<Histogram_t>::iterator i_hist=histograms.begin();
                i_hist!=histograms.end(); ++i_hist){
            double lower,middle,upper;
            i_hist->Integrate(0,*i_low,*i_high,max_field,lower,middle,upper);
            const double integral=lower+middle;
            const double total=integral+upper;

            double lower_norm,middle_norm,upper_norm;
            i_hist->IntegrateNormalisation(0,*i_low,*i_high,max_field,lower_norm,middle_norm,upper_norm);
            const double integral_norm=lower_norm+middle_norm;
            //const double total_norm=integral_norm+upper_norm;

            stepGraph->SetPoint(count,i_hist->x,middle/middle_norm );
            toZeroGraph->SetPoint(count,i_hist->x,integral/integral_norm );
            toMaxGraph->SetPoint(count,i_hist->x,upper/upper_norm );
            ratiosGraph->SetPoint(count,i_hist->x,integral/total);
            ++count;
        }
        // Set the colour
        double colorIndex=(*i_high-min_field)/max_colour_field*gStyle->GetNumberOfColors();
        int colour=gStyle->GetColorPalette(colorIndex);
        toZeroGraph->SetLineColor(colour);
        toMaxGraph->SetLineColor(colour);
        stepGraph->SetLineColor(colour);
        ratiosGraph->SetLineColor(colour);

        toZeroGraph->Sort();
        toMaxGraph->Sort();
        stepGraph  ->Sort();
        ratiosGraph->Sort();

        adjacentStep->Add(stepGraph);
        toZero->Add(toZeroGraph);
        toMax->Add(toMaxGraph);
        ratios->Add(ratiosGraph);
        TString title;
        title+=*i_high;
        legend->AddEntry(stepGraph,title.Data(),"l");
        i_low=i_high;
    }

    //TLegend* legend2=(TLegend*)legend->Clone();

    legend->SetMargin(0.75);
    legend->SetNColumns(2);

    TCanvas* c1=new TCanvas("c1","c1",899,972);
    toZero->Draw("A");
    toZero->GetYaxis()->SetTitle("Total Yield Below a Given Momentum");
    ratios->Draw("A");
    ratios->GetYaxis()->SetTitle("Fraction of Spectrum Below a Given Momentum");
    DrawPrint(adjacentStep,legend,outFileName,"adjacent",save_method);
    DrawPrint(toZero,legend,outFileName,"toZero",save_method);
    DrawPrint(toMax,legend,outFileName,"toMax",save_method);
    DrawPrint(ratios,legend,outFileName,"ratios",save_method);
}