#include <TROOT.h>
#include <cometEventLoop.hxx>

#include "IHandle.hxx"
#include "IG4VHit.hxx"
#include "IG4HitCalo.hxx"
#include "IG4HitSegment.hxx"
#include "IG4Trajectory.hxx"
#include <TFile.h>
#include <TTree.h>
#include <TH1F.h>
#include <TH2F.h>
#include <globals.hh>
#include <G4ios.hh>
#include <iostream>
#include <cassert>
using std::cout;
using std::endl;

class ICheckFits: public COMET::ICOMETEventLoopFunction {
 private:
  TH1F* fEnergy;
  TH1F* fESurface;
  TH1F* fNumXtal;
  TH1F* fNumHit;
  TH1F* fNumClus;
  TH1F* fXtalE;
  TH1F* fXtalT;
  TH1F* fClusT;
  TH1F* fClusdT;
  TH1F* fClusEoP;
  TH1F* fClusE;
  TH2F* fClusTvE;
  TH2F* fXtalTvE;

public:

  ICheckFits() {
    fMuon = false;
  }
  
  virtual ~ICheckFits() {};
  
  void Usage(void) {
        std::cout << "    -O quiet    Don't call TObject::ls()"
                  << std::endl;
        std::cout << "    -O ofile    List events in the file"
                  << std::endl;
        std::cout << "    -O muon     The primary is a muon"
                  << std::endl;
  }
  virtual bool SetOption(std::string option,std::string value="") {
    if (option == "quiet") fQuiet = true;
    else if (option == "ofile") fOutputFile = value;
    else if (option == "muon") fMuon = true;
    return true;
  }
  // using COMET::ICOMETEventLoopFunction::Finalize;

  void Finalize(COMET::ICOMETOutput* const output) {
    std::string rootname = fOutputFile+".root";
    //TFile* histfile = new TFile(rootname.c_str(),"RECREATE");
    if (!fHistFile->IsOpen()) { std::cout << " no output file " << rootname << std::endl ; return;}
    //    if (!fHistFile->IsOpen()) return;
    fHistFile->cd();
    fEnergy->Write();
    fESurface->Write();
    fNumXtal->Write();
    fNumHit->Write();
    fNumClus->Write();
    fXtalE->Write();
    fXtalT->Write();
    fClusEoP->Write();
    fClusE->Write();
    fClusT->Write();
    fClusdT->Write();
    fXtalTvE->Write();
    fClusTvE->Write();
    outTree->Write();
    //    fHistFile->Close();
    //    histfile->Close();

  }

   void Initialize() {

     //     if (!gFile->IsWritable()){
     //  cout << "Can't find suitable file for putting output trees in" << endl;
     //  assert(0);
     // }
    std::string rootname = fOutputFile+".root";
    fHistFile = new TFile(rootname.c_str(),"RECREATE");
    if (!fHistFile->IsOpen()) { std::cout << " couldn't open output file " << rootname << std::endl ; return;}
    outTree = new TTree("CaloHits","Hits in COMET Calorimeter");
    outTree->Branch("eHit",&eHit);
    outTree->Branch("tHit",&tHit);
    outTree->Branch("pHit",&pHit);
    outTree->Branch("idHit",&idHit);
    outTree->Branch("PDGHit",&PDGHit);
    outTree->Branch("XtalIDHit",&XtalIDHit);
    outTree->Branch("sumE",&sumE,"sumE/D");
    outTree->Branch("wavX",&wavX,"wavX/D");
    outTree->Branch("wavY",&wavY,"wavY/D");
    outTree->Branch("wavZ",&wavZ,"wavZ/D");
    outTree->Branch("iniPx",&iniPx,"iniPx/D");
    outTree->Branch("iniPy",&iniPy,"iniPy/D");
    outTree->Branch("iniPz",&iniPz,"iniPz/D");
    outTree->Branch("iniX",&iniX,"iniX/D");
    outTree->Branch("iniY",&iniY,"iniY/D");
    outTree->Branch("iniZ",&iniZ,"iniZ/D");
    outTree->Branch("strPx",&strPx,"strPx/D");
    outTree->Branch("strPy",&strPy,"strPy/D");
    outTree->Branch("strPz",&strPz,"strPz/D");

    fEnergy   = new TH1F("Energy", "Energy measured in Calorimeter",
                         350, 0., 350.);
    fESurface = new TH1F("ESurface", "Energy of Particle Entering Calorimeter",
                         350, -0., 350.);
    fNumHit   = new TH1F("NumHit", "Number of Hits in Event",
                         50, 0., 50.);
    fNumClus  = new TH1F("NumClus", "Number of Clusters in Event",
                         20, 0., 20.);
    fNumXtal  = new TH1F("NumXtal", "Number of Crystals Hit in Event",
                         50, 0., 50.);
    fXtalE    = new TH1F("XtalE", "Energy in Individual Crystals",
                         200, 0., 200.);
    fXtalT    = new TH1F("XtalT", "Time in Individual Crystals",
                         1000, 0., 200.);
    fXtalTvE  = new TH2F("XtalTvE", "Time vs Energy in Individual Crystals",
                         500, 0., 200., 50, 0., 300.);
    fClusE    = new TH1F("ClusE", "Energy in Cluster",
                         300, 0., 300.);
    fClusEoP  = new TH1F("ClusEoP", "Energy/Momentum for Cluster",
                         500, 0., 2.);
    fClusT    = new TH1F("ClusT", "Time for Cluster",
                         1000, 0., 200.);
    fClusdT   = new TH1F("ClusdT", "Time for Cluster",
                         500, -5., 5.);
    fClusTvE  = new TH2F("ClusTvE", "Time vs Energy for Clusters",
                         500, 0., 200., 50, 0., 300.);
         // Do something
  }


  bool operator () (COMET::ICOMETEvent& event) {

    eHit.clear();
    pHit.clear();
    idHit.clear();
    tHit.clear();
    XtalIDHit.clear();    
    PDGHit.clear();

    std::map <int, double> xtal_list; // could sum energy over xtal here if we want
    std::map <int, double> track_list; // energy from a given track
    std::map <int, double> cluster_time; // could sum energy over xtal here if we want
    std::map <int, double> cluster_energy; // could sum energy over xtal here if we want
    std::map <int, int> cluster_count; // could sum energy over xtal here if we want

    COMET::IHandle<COMET::IG4HitContainer> tgtSt =
      event.Get<COMET::IG4HitContainer>("truth/g4Hits/StrawTrk");
    strPx=strPy=strPz=0;
    if(tgtSt) {
      Double_t strP = 0;
      for (COMET::IG4HitContainer::const_iterator hits = tgtSt->begin();
           hits!= tgtSt->end(); ++hits) {
        COMET::IG4HitSegment *theHit  = dynamic_cast<COMET::IG4HitSegment*>(*hits);
        if (strP<theHit->GetMomentumMag()) {
          strP = theHit->GetMomentumMag();
          strPx = theHit->GetMomentumX();
          strPy = theHit->GetMomentumY();
          strPz = theHit->GetMomentumZ();
        }
      }
    }

    COMET::IHandle<COMET::IG4HitContainer> target =
      event.Get<COMET::IG4HitContainer>("truth/g4Hits/ECAL");
    if(target) {
      Double_t iniP = 0;
      sumE=0;
      wavX=wavY=wavZ=0;
      iniPx=iniPy=iniPz=0;
      iniX=iniY=iniZ=0;
      int nHits=target->size();
      double deposited=0;
      double primaryMomentum=0;
      int pdg=-1;
      //
      // Loop over all hits
      //

      for (COMET::IG4HitContainer::const_iterator hits = target->begin();
           hits!= target->end(); ++hits) {
        COMET::IG4HitCalo *theHit  = dynamic_cast<COMET::IG4HitCalo*>(*hits);
        //theHit->ls();
        //    COMET::IHandle<COMET::IG4HitContainer> target =
        COMET::IHandle<COMET::IG4Trajectory> traj = theHit->GetTrajectory();
        if (traj) {
          std::string name = traj->GetParticleName();
          TLorentzVector initialP = traj->GetInitialMomentum();
          if (iniP<initialP.P()) {
            iniPx = initialP.Px();
            iniPy = initialP.Py();
            iniPz = initialP.Pz();
            iniP  = initialP.P();
            iniX  = (traj->GetInitialPosition()).X();
            iniY  = (traj->GetInitialPosition()).Y();
            iniZ  = (traj->GetInitialPosition()).Z();
          }
          pdg = traj->GetPDGEncoding();
          //COMETLog("Trajectory is " << name
          //            << "PDG Encoding is  " << pdg
          //           << " Incident PID is " << theHit->GetIncidentPID()
          //           << " Hadronic E is " << theHit->GetDepositedHadEnergy()
          //           << " em E is " << theHit->GetDepositedEMEnergy()
          //           << "Initial Momentum is  (" 
          //           << " x=" << initialP.Px() 
          //           << " y= " << initialP.Py() 
          //           << " z= " << initialP.Pz() 
          //           << ")"
          //);
          //          traj->ls();
          // can get the information from the trajectory here about what the primary particle was associated with the hit
        }
        // add total hit energy

        //        tree_hits.push_back(theHit);
        double time = theHit->GetTime();
        double energy = theHit->GetDepositedEnergy();
        primaryMomentum = theHit->GetMomentumMag(); // This is the magnitude of the momentum of particle on cal surface
        unsigned int xtalID = theHit->GetCrystalID();
        int incidentID = theHit->GetPrimaryId();
        int sourcePID = theHit->GetIncidentPID();
        eHit.push_back(energy);
        tHit.push_back(time);
        XtalIDHit.push_back(xtalID);
        idHit.push_back(incidentID);
        pHit.push_back(primaryMomentum);
        //        PDGHit.push_back(pdg);
        PDGHit.push_back(sourcePID);
        sumE += energy;
        wavX += energy*theHit->GetPosX();
        wavY += energy*theHit->GetPosY();
        wavZ += energy*theHit->GetPosZ();

        xtal_list[xtalID]=primaryMomentum;
        
        fXtalE->Fill(energy);
        fXtalT->Fill(time);
        fXtalTvE->Fill(time,energy);

        int slice = (int)time/10.0; // this is a hack should be a parameter

        cluster_time[slice] += time;
        ++cluster_count[slice];
        cluster_energy[slice] +=energy;

        //        std::cout << "Primary Momentum at surface = " << primaryMomentum
        //          << " Deposited Energy = " << energy
        //          << " Primary = " << theHit->GetPrimaryId()
        //          << " Xtal ID = " << theHit->GetXtalID()
        //          << " Time = " << time 
        //          << " (slice = " << slice << " ) "
        //          << std::endl;
        //          (*hits)->ls();++<
      } // end of for loop over hits
      if (sumE) {
        wavX /= sumE;
        wavY /= sumE;
        wavZ /= sumE;
      } else {
        wavX = -1e9;
        wavY = -1e9;
        wavZ = -1e9;
      }

      outTree->Fill();
      //std::cout << "number of distinct crystals is "<< xtal_list.size()
      //          << " number of distinct hits is "<< nHits << std:: endl; 

      fNumXtal->Fill(xtal_list.size());
      fNumHit->Fill(nHits);
 
      std::map<int,double>::iterator iter;
      for (iter = cluster_time.begin(); iter != cluster_time.end();++iter) {
          int islice = iter->first;
          double ctime = iter->second;
          double cenergy = cluster_energy[islice];
          int ccount = cluster_count[islice];

          double clusterTime = ctime/(double)ccount;

          //  std::cout << " time information for slice " << islice
          //          << " time is " << clusterTime        
          //          << " energy is " << cenergy
          //          << " count is " << ccount 
          //          << std::endl ;
        
          fNumClus->Fill(ccount);
          fClusT->Fill(clusterTime);        
          fClusdT->Fill(clusterTime-14.0);
          fClusE->Fill(cenergy);
          fClusEoP->Fill(cenergy/primaryMomentum); // hack get as input, this is P for 100 MeV KE Muon
          fClusTvE->Fill(clusterTime,cenergy);
      }

      //
      // Now fill the sums for the clusters...
      //
      fESurface->Fill(primaryMomentum);
      fEnergy->Fill(sumE);
      //std::cout << "Momentum at surface = " << primaryMomentum 
      //        << " Deposited = " << sumE
      //        << " n Hits  = " << nHits
      //        << " n Crystals Hit  = " << xtal_list.size()
      //        << std::endl;
      //target->ls();
    } else {
      COMETLog("No IDatum called exists for event number "<<event.GetEventId());
    }
    //COMETLog("Event " << event.GetContext());
    return true;
  }

private:
  std::string fOutputFile;
  bool fQuiet;
  bool fMuon;
  TFile* fHistFile;
  TTree* outTree;
  std::vector<double>eHit;
  std::vector<double>tHit;
  std::vector<double>pHit;
  std::vector<int>idHit;
  std::vector<int>PDGHit;
  std::vector<unsigned int>XtalIDHit;
  double sumE;  // Summed Energy
  double wavX;  // Weighted average position X in global coordinate
  double wavY;  // Weighted average position Y in global coordinate
  double wavZ;  // Weighted average position Z in global coordinate
  double iniPx; // Initial Momentum X
  double iniPy; // Initial Momentum Y
  double iniPz; // Initial Momentum Z
  double iniX; // Initial Position X
  double iniY; // Initial Position Y
  double iniZ; // Initial Position Z
  double strPx; // Momentum X from 1st Str
  double strPy; // Momentum Y from 1st Str
  double strPz; // Momentum Z from 1st Str
};

int main(int argc, char **argv) {
  ICheckFits userCode;
  COMET::cometEventLoop(argc,argv,userCode,1);
}