#include <TROOT.h>
#include <TFile.h>
#include <TH1.h>
#include <TH2.h>
#include <TH3.h>
#include <TTree.h>
#include <TDirectory.h>
#include <TDatabasePDG.h>
#include <TParticlePDG.h>
#include "IRooTrackerSelector.hxx"
#include <string>
#include <cmath>
#include <iostream>

namespace{

template <typename T>
double Magnitude(const T&x, const T& y, const T z){
  return sqrt(x*x + y*y + z*z);
}

void ChangeDirectory(const std::string& directory){
  static TDirectory* top_dir=gDirectory;
  top_dir->mkdir(directory.c_str());
  top_dir->cd(directory.c_str());
}

const TNamed* GetParticleName(int pid){
  static TNamed all("all","all particles");
  static TNamed undefined("undef","undefined particles");
  if(pid==0) return &all;
  TParticlePDG* particle=TDatabasePDG::Instance()->GetParticle(pid);
  if(particle) return particle;
  return &undefined;
}

struct Histogram_t{
  TDirectory* directory;
  typedef std::map<int,TH1*> IdToHistList;
  std::string name, title, axes;
  int nx, bx, sx;
  int ny, by, sy;
  IdToHistList histos;
  bool is2D;
  Histogram_t( ){};
  Histogram_t( std::string aName, std::string aTitle, std::string aAxes, int aNx, int aBx, int aSx,int aNy, int aBy, int aSy):
      directory(gDirectory),name(aName),title(aTitle),axes(aAxes), nx(aNx),bx(aBx),sx(aSx), ny(aNy),by(aBy),sy(aSy),is2D(true){}
  Histogram_t( std::string aName, std::string aTitle, std::string aAxes, int aNx, int aBx, int aSx):
      directory(gDirectory),name(aName),title(aTitle),axes(aAxes), nx(aNx),bx(aBx),sx(aSx),is2D(false){}

  TH1* Make(int pid_1){
     directory->cd();
     // make the name for the first particle
     const TNamed* pid_name=GetParticleName(pid_1);

     std::string tmp_title=title+" for ";
     tmp_title+=pid_name->GetTitle();
     tmp_title+=";" + axes;
     TH1* hist;
     if(!is2D) hist=new TH1F(Form("%s_%s",name.c_str(),pid_name->GetName()),
                                     tmp_title.c_str(), nx,bx,sx);
     else hist=new TH2F(Form("%s_%s",name.c_str(),pid_name->GetName()),
                                     tmp_title.c_str(), nx,bx,sx,ny,by,sy);
     histos[pid_1]=hist;
     return hist;
  }

  void Fill(int pid_1,double x ,double y=1){
    GetHistogram(0)->Fill(x,y);
    GetHistogram(pid_1)->Fill(x,y);
  }
  TH1* GetHistogram(int pid_1, bool make=true){
    IdToHistList::iterator i_1=histos.find(pid_1);
    if(i_1!=histos.end()) {
        return i_1->second;
    }
    if(!make) return NULL;
    //else need to make this histogram first
    return Make(pid_1);
  }

  static Histogram_t* Make1DHistoList( std::string aName, std::string aTitle, std::string aAxes, int aNx, int aBx, int aSx){
    ChangeDirectory(aName);
    return new Histogram_t(aName,aTitle,aAxes,aNx,aBx,aSx);
  }
  static Histogram_t* Make2DHistoList( std::string aName, std::string aTitle, std::string aAxes, int aNx, int aBx, int aSx,int aNy, int aBy, int aSy){
    ChangeDirectory(aName);
    return new Histogram_t(aName,aTitle,aAxes,aNx,aBx,aSx,aNy,aBy,aSy);
  }
};

class SummarizeRTSelector: public IRooTrackerSelector{
    public:
        SummarizeRTSelector():IRooTrackerSelector(){}

        virtual void SlaveBegin(TTree *tree){
            hMomentum=Histogram_t::Make1DHistoList("hMomentum","Momentum","Momentum (MeV/c)",400,0,1e3);
            hEnergy=Histogram_t::Make1DHistoList("hEnergy","Energy","Energy (MeV)",400,0,1e3);
            hTime=Histogram_t::Make1DHistoList("hTime","Time","Time (s)",400,0,5);
            hPosZX=Histogram_t::Make2DHistoList("hPosZX","Position in the ZX plane","Z position (cm); X position (cm)",400,-300,1e3,400, -300, 1e3);
            hPosXY=Histogram_t::Make2DHistoList("hPosXY","Position in the XY plane","X position (cm); Y position (cm)",400,-300,1e3,400, -300, 1e3);
            hPosYZ=Histogram_t::Make2DHistoList("hPosYZ","Position in the YZ plane","Y position (cm); Z position (cm)",400,-300,1e3,400, -300, 1e3);
            ChangeDirectory("//");
            hMultiplicity=new TH1F("hMultiplicity","Number of particles;Event No.; Number of Particles",tree->GetEntries(),0,tree->GetEntries());
            hDirection=new TH3F("hDirection","Direction Cosines of All Particles;X;Y;Z",300,-1,1,300,-1,1,300,-1,1);
        }

        virtual Bool_t Process(Long64_t entry){
            // Load the entry
            GetEntry(entry);

            hMultiplicity->SetBinContent(entry,StdHepN);
            for(int i=0; i< StdHepN; ++i){
                // Get variables
                const int& pid=StdHepPdg[i];
                const Double_t* mom=StdHepP4[i];
                const Double_t& energy=StdHepP4[i][3];
                const Double_t* pos=StdHepX4[i];
                const Double_t& time=StdHepX4[i][3];
                const Double_t mag_mom=Magnitude(mom[0],mom[1],mom[2]);

                // Fill the histograms
                hMomentum->Fill(pid,mag_mom);
                hEnergy->Fill(pid,energy);
                hTime->Fill(pid,time);
                hPosZX->Fill(pid,pos[2],pos[0]);
                hPosXY->Fill(pid,pos[0],pos[1]);
                hPosYZ->Fill(pid,pos[1],pos[2]);
                hDirection->Fill(mom[0]/mag_mom,mom[1]/mag_mom,mom[2]/mag_mom);
                ++fParticleTally[pid];
            }
            return kTRUE;
        }

        virtual void SlaveTerminate(){
            int size=fParticleTally.size();
            ChangeDirectory("//");
            TH1* p_tally=new TH1F("hParticles","Tally of Particle Types",size,0,size);
            p_tally->SetStats(false);
            int i=1;
            for(ParticleTallyMap::const_iterator i_part=fParticleTally.begin();
                    i_part!=fParticleTally.end(); ++i_part){
                std::string name=GetParticleName(i_part->first)->GetName();
                std::cout<<name<<"("<<i_part->first<<") = "<<i_part->second<<std::endl;
                p_tally->GetXaxis()->SetBinLabel(i, name.c_str());
                p_tally->Fill(i-1, i_part->second);
                ++i;
            }
        }
    private:
        typedef std::map<int,long int> ParticleTallyMap;
        ParticleTallyMap fParticleTally;
        ParticleTallyMap fMonitorTally;
        Histogram_t *hMomentum, *hEnergy, *hTime;
        Histogram_t *hPosZX, *hPosXY, *hPosYZ;
        TH1 *hMultiplicity;
        TH3 *hDirection;
};

const std::string defaultTreeName("gRooTracker");
const std::string defaultOutFileName("summary.root");

void PrintUsage(const char* progName){
    std::cout<<"Usage: "<<progName<<": [options] InputFile [OutputFile]\n"
               "Summarize a RooTracker file's contents.\n"
               "If OutputFile is not provided, histograms are written to '"<<defaultOutFileName<<"'\n"
               "Options:\n"
               "-t TREENAME      Specify the name of the tree in the file [defaults to "<<defaultTreeName <<"]\n"
               "-h               Print this help and exit\n"
               <<std::endl;
}

} // end anonymous namespace

int main(int argc, char** argv){
    std::string inFileName;
    std::string outFileName=defaultOutFileName;
    std::string inTreeName=defaultTreeName;
    for (;;) {
        int c = getopt(argc, argv, "ht:");
        if (c<0) break;
        switch (c) {
            case 't':
                inTreeName=std::string(optarg);
                break;
            default:
                std::cout<<"Error: Unknown option "<<c <<std::endl<<std::endl;
                PrintUsage(argv[0]);
                return 1;
            case 'h':
                PrintUsage(argv[0]);
                return 0;
        }
    }
    if (argc<=optind) {
        std::cerr << "Error: No input file" << std::endl << std::endl;
        PrintUsage(argv[0]);
        return 1;
    }else {
        inFileName=argv[optind];
        ++optind;
        if(argc>optind){
            outFileName=argv[optind];
        }
    }

    // Open the input file
    TFile* inFile=TFile::Open(inFileName.c_str(),"READ");
    if(!inFile || !inFile->IsOpen()){
        std::cerr<<"Unable to open input file, '"<<inFileName<<"'"<<std::endl;
        return 2;
    }

    // Get the tree
    TTree* inTree=NULL;
    inFile->GetObject(inTreeName.c_str(),inTree);
    if(!inTree) {
        std::cerr<<"Input file, '"<<inFileName<<"' doesn't contain a tree called '"<<inTreeName<<"'"<<std::endl;
        return 3;
    }

    // Open the output file
    TFile* outFile=TFile::Open(outFileName.c_str(),"RECREATE");
    if(!outFile || !outFile->IsOpen()){
        std::cerr<<"Unable to open output file, '"<<outFileName<<"'"<<std::endl;
        return 4;
    }

    // Run the selector
    SummarizeRTSelector selector;
    inTree->Process(&selector);

    // Close the output file
    outFile->Write();
    outFile->Close();

    return 0;
}