#include <TROOT.h>
#include <TFile.h>
#include <TObject.h>
#include <TChain.h>
#include <TH1.h>
#include <TH2.h>
#include <TH3.h>
#include <TTree.h>
#include <TDirectory.h>
#include <TDatabasePDG.h>
#include <TParticlePDG.h>
#include <THnSparse.h>
#include <TPrincipal.h>
#include "IRooTrackerSelector.hxx"
#include <string>
#include <cmath>
#include <iostream>
#include <sstream>
#include <cstdlib>
#include <ICOMETLog.hxx>

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 char* GetParticleName(int pid){
  static TNamed all("all","all particles");
  static TNamed undefined("undef","undefined particles");
  if(pid==0) return all.GetName();
  if(pid==1000020041) return "dynamic";
  TParticlePDG* particle=TDatabasePDG::Instance()->GetParticle(pid);
  if(particle) return particle->GetName();
  std::stringstream sstream;
  sstream<<pid;
  return sstream.str().c_str();
}

struct Histogram_t{
  TDirectory* directory;
  typedef std::map<int,TNamed*> IdToHistList;
  std::string name, title, axes, options;
  int nx, bx, sx;
  int ny, by, sy;
  int nVariables;
  IdToHistList histos;
  bool is2D, isPrincipal;
  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),options(), nx(aNx),bx(aBx),sx(aSx), ny(aNy),by(aBy),sy(aSy),is2D(true),isPrincipal(false){}
  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), options(),nx(aNx),bx(aBx),sx(aSx),is2D(false),isPrincipal(false){}
  Histogram_t( std::string aName, int nVar, const char* opts):directory(gDirectory),name(aName),options(opts),nVariables(nVar),is2D(false),isPrincipal(true){}

  void Write(){
      TDirectory* old=TDirectory::CurrentDirectory();
      directory->cd();
      for(IdToHistList::iterator i_hist=histos.begin();
              i_hist!=histos.end(); ++i_hist){
          i_hist->second->Write();
      }
      old->cd();
  }

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

     std::string tmp_title=title+" for ";
     tmp_title+=pid_name;
     tmp_title+=";" + axes;
     std::string full_name=Form("%s_%s",name.c_str(),pid_name.c_str());

     // Check this really doesn't yet exist:
     TObject* object=directory->Get(full_name.c_str());
     if(object) return (TH1*)object;

     TNamed* hist=NULL;
     if(is2D) hist=new TH2F(full_name.c_str(), tmp_title.c_str(), nx,bx,sx,ny,by,sy);
     else if(isPrincipal) {
         hist=new TPrincipal(nVariables,options.c_str());
         hist->SetName(full_name.c_str());
     }
     else hist=new TH1F(full_name.c_str(), tmp_title.c_str(), nx,bx,sx);
     histos[pid_1]=hist;
     return hist;
  }

  void Fill(int pid_1,const double* vals){
    GetPrincipal(0)->AddRow(vals);
    GetPrincipal(pid_1)->AddRow(vals);
  }
  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){
    if(pid_1>1000020040) pid_1=1000020041;
    IdToHistList::iterator i_1=histos.find(pid_1);
    if(i_1!=histos.end()) {
        return dynamic_cast<TH1*>(i_1->second);
    }
    if(!make) return NULL;
    //else need to make this histogram first
    return (TH1*) Make(pid_1);
  }
  TPrincipal* GetPrincipal(int pid_1, bool make=true){
    if(pid_1>1000020040) pid_1=1000020041;
    IdToHistList::iterator i_1=histos.find(pid_1);
    if(i_1!=histos.end()) {
        return dynamic_cast<TPrincipal*>(i_1->second);
    }
    if(!make) return NULL;
    //else need to make this histogram first
    return (TPrincipal*) 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);
  }
  static Histogram_t* MakePrincipalHistoList( std::string aName, int nVars, const char* opts){
    ChangeDirectory(aName);
    return new Histogram_t(aName,nVars,opts);
  }
};

class SummarizeRTSelector: public IRooTrackerSelector{
    public:
        SummarizeRTSelector(std::string inMonitor):
            IRooTrackerSelector(),fInputMonitorID(-1),fInputMonitorName(inMonitor),
            fMonitorNamesObjectName("MonitorNames"),fFile(NULL),fFilesRead(0),fTotalFilesToRead(0), fEventsRead(0){}

        int GetEventsRead() const{
            return fEventsRead;
        }

        virtual void SlaveBegin(TTree *tree){
            fInputChain=dynamic_cast<TChain*>(tree);
            fTotalFilesToRead=fInputChain->GetNtrees();

            double min_pid       = -5e3 , max_pid       = 5e3           ; int n_pid       = 10e3+1                      ; 
            double min_mom       = 0    , max_mom       = 1e3           ; int n_mom       = 400                         ; 
            double min_ene       = 0    , max_ene       = 1e3           ; int n_ene       = 400                         ; 
            double min_time      = 0    , max_time      = 6e4           ; int n_time      = 1e3                         ; 
            double min_dir       = -1   , max_dir       = 1             ; int n_dir       = 400                         ; 
            double min_pos_x     = -600 , max_pos_x     = 3e3           ; int n_pos_x     = 1e3                         ; 
            double min_pos_y     = -600 , max_pos_y     = 600           ; int n_pos_y     = 1e3                         ; 
            double min_pos_z     = 4e3  , max_pos_z     = 8e3           ; int n_pos_z     = 1e3                         ; 
            double min_pos_phi   = 0    , max_pos_phi   = 2*TMath::Pi() ; int n_pos_phi   = max_pos_phi-min_pos_phi     ; 
            double min_pos_theta = 0    , max_pos_theta =   TMath::Pi() ; int n_pos_theta = max_pos_theta-min_pos_theta ; 
            hMomentum=Histogram_t::Make1DHistoList("hMomentum","Momentum","Momentum (MeV/c)",n_mom,min_mom,max_mom);
            hEnergy=Histogram_t::Make1DHistoList("hEnergy","Energy","Energy (MeV)",n_ene,min_ene,max_ene);
            hTime=Histogram_t::Make1DHistoList("hTime","Time","Time (s)",n_time,min_time,max_time);
            hPosZX=Histogram_t::Make2DHistoList("hPosZX","Position in the ZX plane","Z position (cm); X position (cm)",n_pos_z,min_pos_z,max_pos_z,n_pos_x,min_pos_x,max_pos_x);
            hPosXY=Histogram_t::Make2DHistoList("hPosXY","Position in the XY plane","X position (cm); Y position (cm)",n_pos_x,min_pos_x,max_pos_x,n_pos_y,min_pos_y,max_pos_y);
            hPosYZ=Histogram_t::Make2DHistoList("hPosYZ","Position in the YZ plane","Y position (cm); Z position (cm)",n_pos_y,min_pos_y,max_pos_y,n_pos_z,min_pos_z,max_pos_z);
            hDirX=Histogram_t::Make1DHistoList("hDirX","Distribution over X Direction","Direction",n_dir,min_dir,max_dir);
            hDirY=Histogram_t::Make1DHistoList("hDirY","Distribution over Y Direction","Direction",n_dir,min_dir,max_dir);
            hDirZ=Histogram_t::Make1DHistoList("hDirZ","Distribution over Y Direction","Direction",n_dir,min_dir,max_dir);
            hPrincipal=Histogram_t::MakePrincipalHistoList("hPrincipal",8,"N");
            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",n_dir,min_dir,max_dir,n_dir,min_dir,max_dir,n_dir,min_dir,max_dir);
            ChangeDirectory("//");


            const int n_bins[]={n_pid , n_mom   /*, n_ene  */ , n_time   , n_pos_x   , n_pos_y   , n_pos_z   , n_dir   , n_dir   , n_dir   }; 
            const double min[]={min_pid , min_mom /*, min_ene*/ , min_time , min_pos_x , min_pos_y , min_pos_z , min_dir , min_dir , min_dir }; 
            const double max[]={max_pid , max_mom /*, max_ene*/ , max_time , max_pos_x , max_pos_y , max_pos_z , max_dir , max_dir , max_dir }; 
            const int dim=sizeof(n_bins)/sizeof(int);
            hFullCorrelation=new THnSparseS("hFullCorrelation","Multi-dimensional correlation histogram",dim,n_bins,min,max);
            for(int i=0; i<9;++i){
                std::string axis;
                switch(i){
                    case 0:axis="PID"         ; break;
                    case 1:axis="Momentum"    ; break;
                    case 2:axis="Time"        ; break;
                    case 3:axis="Position X"  ; break;
                    case 4:axis="Position Y"  ; break;
                    case 5:axis="Position Z"  ; break;
                    case 6:axis="Direction X" ; break;
                    case 7:axis="Direction Y" ; break;
                    case 8:axis="Direction Z" ; break;
                }
                hFullCorrelation->GetAxis(i)->SetTitle(axis.c_str());
            }

        }

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

            hMultiplicity->Fill(StdHepN);
            for(int i=0; i< StdHepN; ++i){
                // Is this in the requested monitor?
                if (fInputMonitorID>-1 && MonitorID[i]!=fInputMonitorID)
                    continue;

                // 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);
                hDirX->Fill(pid,mom[0]/mag_mom);
                hDirY->Fill(pid,mom[1]/mag_mom);
                hDirZ->Fill(pid,mom[2]/mag_mom);
                ++fParticleTally[pid];

                const double values[]={pid   , mag_mom , time   , pos[0]  , pos[1]  , pos[2]  , mom[0]/mag_mom , mom[1]/mag_mom , mom[2]/mag_mom};
                hFullCorrelation->Fill(values);
                hPrincipal->Fill(pid,&values[1]);
            }
            // Add the event as read
            ++fEventsRead;
            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);
                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;
            }

            ChangeDirectory("//");
            hFullCorrelation->Write();
            hPrincipal->Write();

        }

        Bool_t Notify(){
            ++fFilesRead;
            fFile=fInputChain->GetFile();
            if(fTotalFilesToRead>1) COMETLog("Loaded file "<<fFilesRead<<": "<<fFile->GetName());

            // Find the MonitorNames array in the currently open TFile
            TObject* object=fFile->Get(fMonitorNamesObjectName.c_str());
            if(!object ){
                COMETError("Cannot find the Monitor name to ID mapping ("<<fMonitorNamesObjectName<<")");
                fInputMonitorID=-1;
                return true;
            }
            const TObjArray* map=(const TObjArray*)object;
            if(!map){
                COMETError("No suitable object containing the Monitor name to ID mapping.  Found "<<fMonitorNamesObjectName<<" which is "<<object->ClassName());
                fInputMonitorID=-1;
                return true;
            }

            // Convert the array into the monitor map and set the
            // fInputMonitorID
            int count=1;
            fMonitorNameMap.push_back("");
            for(TIter nextName(map);
                    (object=nextName());
                    ++count){
                TObjString* name=(TObjString*)object;
                if(!name)continue;
                fMonitorNameMap.push_back(name->GetName());
                if(!fInputMonitorName.empty() && name->GetName()==fInputMonitorName){
                    fInputMonitorID=count;
                }
            }
            return true;
        }

    private:
        typedef std::map<int,long int> ParticleTallyMap;
        ParticleTallyMap fParticleTally;
        ParticleTallyMap fMonitorTally;
        Histogram_t *hMomentum, *hEnergy, *hTime;
        Histogram_t *hPosZX, *hPosXY, *hPosYZ;
        Histogram_t *hDirX, *hDirY, *hDirZ;
        Histogram_t *hPrincipal;
        TH1 *hMultiplicity;
        TH3 *hDirection;
        THnSparse *hFullCorrelation;
        int fInputMonitorID;
        std::string fInputMonitorName;
        std::string fMonitorNamesObjectName;
        typedef std::vector<std::string> MonitorNameMap;
        MonitorNameMap fMonitorNameMap;
        TFile* fFile;
        TChain* fInputChain;
        int fFilesRead,fTotalFilesToRead, fEventsRead;
};

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

void PrintUsage(const char* progName){
    std::cout<<"Usage: "<<progName<<": [options] InputFile [InputFile2 ...]\n"
               "Summarize a RooTracker file's contents.\n"
               "If OutputFile is not provided, histograms are written to '"<<defaultOutFileName<<"'\n"
               "InputFile will be loaded as a TChain, so it can contain wildcarding if desired, but \n"
               "take care that the shell doesn't expand the glob if you wish ROOT to do this instead.\n"
               "Options:\n"
               "-t TREENAME      Specify the name of the tree in the file [defaults to "<<defaultTreeName <<"]\n"
               "-o OUTFILE       Specify the name of the output file [defaults to "<<defaultOutFileName <<"]\n"
               "-m monitor       Specify the name of monitor to filter particles\n"
               "-h               Print this help and exit\n"
               <<std::endl;
}

} // end anonymous namespace

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

    // Open the input file(s)
    TChain* inChain=new TChain(inTreeName.c_str());
    int n_added=0;
    for(size_t i=0; i<inFileNames.size(); ++i){
        int added=inChain->Add(inFileNames.at(i).c_str());
        if(added<=0){
            std::cerr<<"Filename (or glob): "<<inFileNames.at(i)<<" is invalid."<<std::endl;
        }
        n_added+=added;
    }
    if(n_added<=0 || inChain->GetNtrees()<=0 || inChain->LoadTree(0)){
        std::cerr<<"No valid input filename(s) provided."<<std::endl;
        return 2;
    }
    COMETLog(inChain->GetNtrees()<<" files will be studied");
    COMETLog(inChain->GetEntries()<<" events will be studied");

    // 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 3;
    }

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

    // Print the number of events read
    COMETLog("Total Events Read : " << selector.GetEventsRead());

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

    return 0;
}