#include <TH2.h>
#include <TF1.h>
#include <TFormula.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 <TList.h>
#include <TStyle.h>
#include <TColor.h>
#include <TROOT.h>
#include <iostream>
#include <vector>
#include "../scripts/palettes.C"
using namespace std;

class PlotLooper{
    public: 
        struct Plot{
            TH1* histogram;
            TH1* original_histogram;
            double value;
            Plot():histogram(NULL),original_histogram(NULL),value(1){}
        };

        PlotLooper(
                const char* file_names,
                const char* histograms,
                const char* normalisation,
                const char* value_format,
                double max_field,
                double min_field,
                const char* method
                )
            :fNormalisationVal(1),fNormalisation_hist(NULL),fInputFileNames(file_names),
            fHistogramNames(histograms), fNormalisations(normalisation),
            fValueFormat(value_format),fPlotMethod(method),fValueFunction(NULL){
                fMaxField=1.01*(max_field-0.99*min_field);
                fMinField=0.99*min_field;
            }

        bool Initiailise(){
            SetUserPalette(1,255,0.5);

            // Sort file names
            fInputFileNames_tkns=fInputFileNames.Tokenize(" ");
            fNHists=fInputFileNames_tkns->GetEntries();
            fDifferentInputFiles=(fNHists!=1);
            fUseFullNameForValue=fDifferentInputFiles;

            // Sort histogram names
            fHistogramNames_tkns=TokenizeAndCheck(fHistogramNames,fNHists,fDifferentHists, "histogram name");
            if (!fHistogramNames_tkns) return false;
            if(fNHists==1){
                fNHists=fHistogramNames_tkns->GetEntries();
                fUseFullNameForValue=fInputFileNames.Contains(":");
            }

            // sort out the normalisation
            fNormalisations_tkns=TokenizeAndCheck(fNormalisations,fNHists,fNormPerPlot, "normalisation");
            if (!fNormalisations_tkns) return false;

            if(!fNormPerPlot) {
                if(fNormalisations.Contains("%")){
                    fNormalisations_tkns=BuildNormalisationTokens(fNHists,fNormalisations);
                    cout<<"Normalisation tokens:"<<endl;
                    fNormalisations_tkns->Print();
                    fNormPerPlot=true;
                }else{
                    SortOneNormalisationToken(0);
                }
            }

            cout<<std::boolalpha<<"normPerPlot: "<<fNormPerPlot<<endl;
            cout<<std::boolalpha<<"differentHists: "<<fDifferentHists<<endl;
            cout<<std::boolalpha<<"differentInputFiles: "<<fDifferentInputFiles<<endl;

            return true;

        }

        void SetValueFunction(const TString function){
            if(function=="") return;
            fValueFunction=new TFormula("value_function",function);
        }

        TObjArray* BuildNormalisationTokens(int n_tokens,TString format){
            TObjArray* array=new TObjArray();
            for(int i=0; i<n_tokens; ++i){
                TString input_name=GetInputName(i);
                cout<<"input name: "<<input_name<<endl;
                double value=ExtractValue(input_name);
                cout<<"value: "<<value<<endl;
                array->Add(new TObjString(Form(format,value)));
            }
            return array;
        }

        Plot GetPlot(int i_plot){
            Plot return_plot=GetPlotWithoutNorm(i_plot);
            if(!return_plot.histogram) return return_plot;
            double norm_value=1;
            TH1* norm_hist=NULL;
            GetNormalisation(i_plot,norm_value,norm_hist);
            ApplyNormalisation(return_plot.histogram,norm_value,norm_hist);
            return return_plot;
        }

        Plot GetPlotWithoutNorm(int i_plot){
            cout<<"Processing plot: "<<i_plot<<endl;

            // Sort out the file / histogram name
            TString input_name=GetInputName(i_plot);

            // Extract the actual histogram
            TH1* hist=GetHistogramFromFile(input_name);
            if(!hist) return Plot();

            Plot return_plot;
            return_plot.original_histogram=hist;

            // Convert TH2 and TH3 to a TH1 and normalise
            return_plot.histogram=GetFlattenedHistogram( hist);

            // Extract the label value for this histogram
            return_plot.value=ExtractValue(input_name);

            // Set plot colour
            SetPlotColour(ConvertValue(return_plot.value),return_plot.histogram);

            return return_plot;
        }

        TString GetInputName(int i_plot){
            TString input_name;
            cout<<"fDifferentInputFiles: "<<fDifferentInputFiles<<endl;
            if(fDifferentInputFiles) {
                input_name=((TObjString*)fInputFileNames_tkns->At(i_plot))->GetString();
                input_name+=":";
            }else{
                input_name=fInputFileNames;
            }
            cout<<"input_name: "<<input_name<<endl;

            // Add or complete the histogram name
            cout<<"fDifferentHists: "<<fDifferentHists<<endl;
            if(fDifferentHists){
                input_name+=((TObjString*)fHistogramNames_tkns->At(i_plot))->GetName();
            }else{
                input_name+=fHistogramNames;
            }
            cout<<"input_name: "<<input_name<<endl;
            return input_name;
        }

        void GetNormalisation(int i_plot, double& norm_value,TH1*& norm_hist){
            if(fNormPerPlot) {
                // reset the norms from the last time
                fNormalisationVal=1;
                fNormalisation_hist=NULL;
                // Get this plot's norm
                SortOneNormalisationToken(i_plot);
                cout<<"Set normalisation to: "<<fNormalisationVal<<endl;
                cout<<"Set normalisation hist to: "<<(norm_hist?norm_hist->GetName():"NULL")<<endl;
            }
            norm_value=fNormalisationVal;
            norm_hist=fNormalisation_hist;
        }

    private:

        TH1* GetHistogramFromFile(TString input_name){
            TObjArray* tokens=input_name.Tokenize(":");
            if(tokens->GetEntries() !=2){
                cout<<"Error: Wrong number of tokens for input name which should be of the form: <input_file>:<hist_path> (and contain a single colon, ':' )"<<endl;
                cout<<"Received: "<<input_name<<endl;
                return NULL;
            }
            // Open the file
            TString file=((TObjString*)tokens->At(0))->GetName();
            const TString hist_name=((TObjString*)tokens->At(1))->GetName();
            if(file=="<>"){
                file=fCurrentFile->GetName();
            }else{
                fCurrentFile=NULL;
                fCurrentFile=TFile::Open(file.Data());
                if(!fCurrentFile) {
                    cout<<"Error: Cannot open file: "<<file<<endl;
                    return NULL;
                }
            }

            // Get the source histogram
            TH1* profile=(TH1*)fCurrentFile->Get(hist_name);
            if(!profile) {
                cout<<"Error: Cannot find histogram ("<<hist_name<<") in file: "<<file<<endl;
                return NULL;
            }
            cout<<"Found "<<hist_name<<" in file "<<file<<endl;

            return profile;
        }

        TH1* GetFlattenedHistogram( TH1* hist){
            if(!hist){
                cout<<"NUll hist passed  to GetFlattenedHistogram()"<<endl;
                return NULL;
            }
            TH2* hist2d=dynamic_cast<TH2*>(hist);
            if(hist2d){
                cout<<"2D histogram, make profile."<<endl;
                if (fPlotMethod=="projectionX"){
                    hist=hist2d->ProjectionX();
                }else if (fPlotMethod=="projectionY"){
                    hist=hist2d->ProjectionY();
                }else if (fPlotMethod=="profileY"){
                    hist=hist2d->ProfileY();
                }else if (fPlotMethod=="profileX"){
                    hist=hist2d->ProfileX();
                }else{
                    cout<<"Bad 2D --> 1D method: "<<fPlotMethod<<". Should be one of: profileX,profileY, projectionX,projectionY"<<endl;
                    return NULL;
                }
            }
            return hist;
        }

        void ApplyNormalisation( TH1* hist, double norm_value, TH1* norm_hist){
            cout<<"applying normalisation"<<endl;
            // Apply normalisation
            if(norm_value!=1) hist->Scale(1/norm_value);
            if(norm_hist) hist->Divide(norm_hist);
            cout<<"applied normalisation"<<endl;
        }

        void SortOneNormalisationToken(int i_plot,TObjArray* normalisations=NULL){
            if(normalisations==NULL) normalisations=fNormalisations_tkns;
            // Get the token we need
            TString token=((TObjString*)normalisations->At(i_plot))->GetString();

            TObjArray* float_and_hist=token.Tokenize("*");
            if(float_and_hist->GetEntries()>=2){
                for(int i=0; i<float_and_hist->GetEntries();++i_plot){
                    SortOneNormalisationToken(i,float_and_hist);
                }
                return;
            }

            // Is it only a number?
            if(token.IsFloat()){
                fNormalisationVal*=token.Atof();
                return;
            }

            // Treat as a histogram
            TH1* hist=GetHistogramFromFile(token);
            if(!fNormalisation_hist){
                fNormalisation_hist=(TH1*)hist->Clone(Form("norm_%d",i_plot));
            }else{
                fNormalisation_hist->Multiply(hist);
            }
        }

        double ExtractValue(TString input_name){
            TString string_for_value;
            if( fUseFullNameForValue ){
                if( fDifferentHists ){
                    string_for_value=input_name;
                }else{
                    string_for_value=input_name.Tokenize(":")->At(0)->GetName();
                }
            }
            else if( fDifferentHists) {
                string_for_value=input_name.Tokenize(":")->At(1)->GetName();
            }
            cout<<"Extracting field value from: "<<string_for_value<<" using "<<fValueFormat<<endl;
            double value;
            sscanf(string_for_value.Data(),fValueFormat,&value);
            return value;
        }

        double ConvertValue(double value){
            if(fValueFunction){
                value=fValueFunction->Eval(value);
            }
            return value;
        }

        void SetPlotColour(double value,TH1* hist){
            hist->SetTitle(Form("%g",value));

            // Set the profile's line and marker colour
            double colorIndex=(value-fMinField)/fMaxField*gStyle->GetNumberOfColors();
            int colour=gStyle->GetColorPalette(colorIndex);
            hist->SetLineColor(colour);
            hist->SetLineStyle(0);
            hist->SetMarkerColor(colour);

            cout<<"Colours: "<<value<<", "<<colorIndex<<", "<<colour<<endl;
        }

        int GetNumPlots()const{return fNHists;}


        TObjArray* TokenizeAndCheck(const TString input,int n_expected_tokens, bool& many_values,const char* purpose){
            TObjArray* tokens=input.Tokenize(" ");
            if(tokens->GetEntries()==1){
                many_values=false;
            }else if(tokens->GetEntries()==0){
                cout<<"Error: No "<<purpose<<" tokens provided"<<endl;
                return NULL;
            }else {
                if (n_expected_tokens!=1 && tokens->GetEntries()!=n_expected_tokens){
                    cout<<"Error: Number of "<<purpose<<" tokens must be either 1 (a single value for all histograms), or the same as the number of histograms"<<endl;
                    return NULL;
                }
                many_values=true;
            }
            return tokens;
        }

    private:
        double fNormalisationVal;
        TH1* fNormalisation_hist;
        TString fCurrentHistName;
        int fNHists;
        TFile* fCurrentFile;

        TString fInputFileNames;
        TString fHistogramNames;
        TString fNormalisations;
        TString fValueFormat;
        TString fPlotMethod;
        TFormula* fValueFunction;

        TObjArray* fInputFileNames_tkns;
        TObjArray* fHistogramNames_tkns;
        TObjArray* fNormalisations_tkns;

        bool fUseFullNameForValue;
        bool fDifferentInputFiles;
        bool fDifferentHists;
        bool fNormPerPlot;
        double fMaxField;
        double fMinField;
};

void PlotHeightVsBeamline(
        TString histogram_files,
        const char* title="Signal Particle Height as a Function of Dipole Field",
        const char* histograms="truth_IPhaseIISignalAcceptanceAnalyzer/hBeamHeightPrimary",
        const char* normalisation_str="1",
        const char* outFileName="overlaid",
        double y_high=0,
        double y_low=0,
        const char* value_format="-Tor2_1.50_profile",
        const char* value_function="",
        double max_field=2.5,
        const char* method="profileX",
        int rebin=1,
        double min_field=0,
        double xmax=0,
        double xmin=0
        ){

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

    THStack stack("stack",title);
    TLegend* legend=new TLegend(0.1,0.7,0.3,0.9);

    for(int i_hist=0; i_hist< plotter.GetNumPlots(); ++i_hist){
        PlotLooper::Plot plot=plotter.GetPlot(i_hist);
        TH1* hist=plot.histogram;
        if(!hist) {
            cout<<"no profile returned"<<endl;
            continue;
        }

        if(rebin!=1) hist->Rebin(rebin);

        cout<<"Adding profile: "<<hist->GetName()<<", "<<hist->GetTitle()<<endl;
        stack.Add(hist);
        legend->AddEntry(hist,hist->GetTitle(),"l");
    }

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

    //TCanvas* c1=new TCanvas("c1","c1",1400,1000);
    TCanvas* c1=new TCanvas("c1","c1",899,972);
    stack.Draw("L nostack ");
    if(xmax!=0) stack.GetHistogram()->GetXaxis()->SetRangeUser(xmin,xmax);
    if(y_high!=y_low){
        if(y_high < y_low) {
            double tmp=y_low; y_high=y_low; y_low=tmp;
        }
        stack.GetHistogram()->GetYaxis()->SetRangeUser(y_low,y_high);
    }
    legend->SetFillColor(kWhite);
    legend->SetNColumns(2);
    legend->SetMargin(0.75);
    //legend->SetHeader("Momentum At Target (MeV/c)");
    legend->Draw();
    c1->SetGridy();
    c1->SetGridx();

    c1->Print(Form("%s.png" ,outFileName));
    c1->Print(Form("%s.svg" ,outFileName));
    c1->Print(Form("%s.root",outFileName));
}