#include <cmath>
#include <string>

#include <IReconTrack.hxx>
#include <ITrackState.hxx>
#include <IGeomInfo.hxx>
#include <IRecPackManager.hxx>

#include <TrackTruthInfo.hxx>

#include "IReconSMRDModule.hxx"

#include "ITFBDigit.hxx"
#include "ITFBChannelId.hxx"

#include <IRealDatum.hxx>

#include <TrackingUtils.hxx>

ClassImp(COMET::IReconSMRDModule);
ClassImp(COMET::IReconSMRDModule::IReconSMRDHit);
COMET::IReconSMRDModule::IReconSMRDHit::~IReconSMRDHit() {}
ClassImp(COMET::IReconSMRDModule::ISmrdIsoTrack);
COMET::IReconSMRDModule::ISmrdIsoTrack::~ISmrdIsoTrack() {}

#define CVSTAG "\
  $Name: v5r19 $"
#define CVSID  "\
"  

COMET::IReconSMRDModule::IReconSMRDModule(const char *name, const char *title){
  SetNameTitle(name, title);

    fIsEnabled   = kTRUE;
    fDescription = "SMRD Recon Output";
    fCVSTagName  = CVSTAG;
    fCVSID       = CVSID;

    fNSmrdReconHits = 0;
    fSmrdReconHits  = new TClonesArray("COMET::IReconSMRDModule::IReconSMRDHit", 200);
    fNSmrdIsoTracks = 0;
    fSmrdIsoTracks  = new TClonesArray("COMET::IReconSMRDModule::ISmrdIsoTrack", 200);
}

COMET::IReconSMRDModule::~IReconSMRDModule() {}

Bool_t COMET::IReconSMRDModule::ProcessFirstEvent(COMET::ICOMETEvent &event){
    return true;
}

void COMET::IReconSMRDModule::InitializeModule() {}

void COMET::IReconSMRDModule::InitializeBranches(){
    fOutputTree->Branch("NSmrdReconHits",
          &fNSmrdReconHits, "NSmrdReconHits/I",
          fBufferSize);
    fOutputTree->Branch("NSmrdIsoTracks",
          &fNSmrdIsoTracks, "NSmrdIsoTracks/I",
          fBufferSize);

   fOutputTree->Branch("SmrdReconHits",
          &fSmrdReconHits,
          fBufferSize, fSplitLevel);
    fOutputTree->Branch("SmrdIsoTracks",
          &fSmrdIsoTracks,
          fBufferSize, fSplitLevel);
}

bool COMET::IReconSMRDModule::FillTree(COMET::ICOMETEvent &event) {
    fNSmrdReconHits = 0;
    fSmrdReconHits->Clear();

    fNSmrdIsoTracks = 0;
    fSmrdIsoTracks->Clear();

    COMET::IHandle<COMET::IAlgorithmResult> SmrdResult = event.GetFit("ReconSMRD");
    COMET::IHandle<COMET::IHitSelection> ReconSMRDHitsMatched;
  COMET::IHandle<COMET::IHitSelection> ReconSMRDHitsUnused;
    COMET::IHandle<COMET::IHitSelection> ReconSMRDHitsUsed;
    COMET::IHandle<COMET::IReconObjectContainer> xyzReconTracks;
  
    if (SmrdResult){
      ReconSMRDHitsUnused          = SmrdResult->GetHitSelection("unused");
    ReconSMRDHitsMatched          = SmrdResult->GetHitSelection("matchedSmrdReconHits");
    ReconSMRDHitsUsed          = SmrdResult->GetHitSelection("used");
      xyzReconTracks         = SmrdResult->GetResultsContainer("final");
    }
  
    //-------------------------------------------------------------------------------------------
    //Fill hits information
    //-------------------------------------------------------------------------------------------  
  // ReconSMRDHitsUnused
    if(ReconSMRDHitsUnused){
      for (COMET::IHitSelection::iterator h = ReconSMRDHitsUnused->begin();h != ReconSMRDHitsUnused->end(); h++) {
          COMET::IHandle<COMET::IHit> hit = *h;
          if (!hit) continue;
      
      
      //get hit contributors
      //only deal with "fully-calibrated" two-sided hits 
      if ((*h)->GetContributorCount()!=2){
        COMETWarn("warning: more or less than two smrd single hits contributed  to this IReconHit; skipping this IReconHit");
        continue;
                         } 
       
      //******************get two single hits-contributors
                  COMET::IHandle<COMET::IHit> contrhit1=(*h)->GetContributor(0);
                        COMET::IHandle<COMET::IHit> contrhit2=(*h)->GetContributor(1);
      
            
      if(!contrhit1 || ! contrhit2){
        COMETWarn("warning: failed get constituents for SMRD TRecon hit; skip this IReconHit");
        continue;
       }
      
      
          IReconSMRDHit *reconHit = new((*fSmrdReconHits)[fNSmrdReconHits++]) IReconSMRDHit;

                      
      reconHit->Position            = TLorentzVector(hit->GetPosition(),hit->GetTime());
          reconHit->PositionUncertainty = TLorentzVector(hit->GetUncertainty(),hit->GetTimeUncertainty());
          reconHit->Charge              = hit->GetCharge();   
      
      reconHit->dZ = (*h)->GetPosition().Z()-contrhit1->GetPosition().Z();
      reconHit->dT = contrhit1->GetTime()-contrhit2->GetTime();
            reconHit->ContribCharge[0] = contrhit1->GetCharge();
      reconHit->ContribCharge[1] = contrhit2->GetCharge(); 
    
          bool isTopWall    = COMET::IGeomInfo::SMRD().IsSMRDTopWall(hit);
          bool isBottomWall = COMET::IGeomInfo::SMRD().IsSMRDBottomWall(hit);
          bool isRightWall  = COMET::IGeomInfo::SMRD().IsSMRDRightWall(hit);
          bool isLeftWall   = COMET::IGeomInfo::SMRD().IsSMRDLeftWall(hit);
    
          if(isTopWall)            
            reconHit->Wall = kTop;    
          if(isBottomWall)            
            reconHit->Wall = kBottom;    
          if(isRightWall)            
            reconHit->Wall = kRight;    
          if(isLeftWall)            
            reconHit->Wall = kLeft;    		
  
          reconHit->Yoke    = COMET::IGeomInfo::SMRD().GetYokeRingNumber(hit);
          reconHit->Layer   = COMET::IGeomInfo::SMRD().GetSMRDRingNumber(hit);
          reconHit->Tower   = COMET::IGeomInfo::SMRD().GetSMRDTowerNumber(hit);        
          reconHit->Counter = COMET::IGeomInfo::SMRD().GetSMRDScintNumber(hit);            
      
      
      //fill tfb and rmm info
      COMET::IChannelId chan = contrhit1->GetChannelId();
                        COMET::ITFBChannelId tfbid(chan);
            
      reconHit->RMM = tfbid.GetRMM();
      reconHit->TFB = tfbid.GetTFB();
                        
      reconHit->IsUsed=false;
      
      reconHit->IsInnerMatched=false;
      
      //should be always unmatched but let`s double check
      if (ReconSMRDHitsMatched){
        if (std::find(ReconSMRDHitsMatched->begin(), ReconSMRDHitsMatched->end(), *h) != ReconSMRDHitsMatched->end())
          reconHit->IsInnerMatched=true;
                            
      }	
      
      
      }//end of hits loop
    }//end of if(ReconSMRDHitsUnused)
  
  //ReconSMRDHitsUsed
    if(ReconSMRDHitsUsed){
      for (COMET::IHitSelection::iterator h = ReconSMRDHitsUsed->begin();h != ReconSMRDHitsUsed->end(); h++) {
          COMET::IHandle<COMET::IHit> hit = *h;
          if (!hit) continue;
      
            //get hit contributors
      //only deal with "fully-calibrated" two-sided hits 
      if ((*h)->GetContributorCount()!=2){
        COMETWarn("warning: more or less than two smrd single hits contributed  to this IReconHit; skipping this IReconHit");
        continue;
                         } 
       
      //******************get two single hits-contributors
                  COMET::IHandle<COMET::IHit> contrhit1=(*h)->GetContributor(0);
                        COMET::IHandle<COMET::IHit> contrhit2=(*h)->GetContributor(1);
      
      if(!contrhit1 || ! contrhit2){
        COMETWarn("warning: failed get constituents for SMRD TRecon hit; skip this IReconHit");
        continue;
       }
      
      
          IReconSMRDHit *reconHit = new((*fSmrdReconHits)[fNSmrdReconHits++]) IReconSMRDHit;

          reconHit->Position            = TLorentzVector(hit->GetPosition(),hit->GetTime());
          reconHit->PositionUncertainty = TLorentzVector(hit->GetUncertainty(),hit->GetTimeUncertainty());
          reconHit->Charge              = hit->GetCharge();    
      
      reconHit->dZ = (*h)->GetPosition().Z()-contrhit1->GetPosition().Z();
      reconHit->dT = contrhit1->GetTime()-contrhit2->GetTime();
            
      reconHit->ContribCharge[0] = contrhit1->GetCharge();
      reconHit->ContribCharge[1] = contrhit2->GetCharge();
        
      
    
          bool isTopWall    = COMET::IGeomInfo::SMRD().IsSMRDTopWall(hit);
          bool isBottomWall = COMET::IGeomInfo::SMRD().IsSMRDBottomWall(hit);
          bool isRightWall  = COMET::IGeomInfo::SMRD().IsSMRDRightWall(hit);
          bool isLeftWall   = COMET::IGeomInfo::SMRD().IsSMRDLeftWall(hit);
    
          if(isTopWall)            
            reconHit->Wall = kTop;    
          if(isBottomWall)            
            reconHit->Wall = kBottom;    
          if(isRightWall)            
            reconHit->Wall = kRight;    
          if(isLeftWall)            
            reconHit->Wall = kLeft;    		
  
          reconHit->Yoke    = COMET::IGeomInfo::SMRD().GetYokeRingNumber(hit);
          reconHit->Layer   = COMET::IGeomInfo::SMRD().GetSMRDRingNumber(hit);
          reconHit->Tower   = COMET::IGeomInfo::SMRD().GetSMRDTowerNumber(hit);        
          reconHit->Counter = COMET::IGeomInfo::SMRD().GetSMRDScintNumber(hit);            
      
              //fill tfb and rmm info
      COMET::IChannelId chan = contrhit1->GetChannelId();
                        COMET::ITFBChannelId tfbid(chan);
            
      reconHit->RMM = tfbid.GetRMM();
      reconHit->TFB = tfbid.GetTFB();

      
      reconHit->IsUsed=true;
      
      reconHit->IsInnerMatched=false;
      
      if (ReconSMRDHitsMatched){
        if (std::find(ReconSMRDHitsMatched->begin(), ReconSMRDHitsMatched->end(), *h) != ReconSMRDHitsMatched->end())
          reconHit->IsInnerMatched=true;
                            
      }	
      
      
      
      }//end of hits loop
    }//end of if(ReconSMRDHitsUsed)

  //-------------------------------------------------------------------------------------------
    //Fill tracks information
    //-------------------------------------------------------------------------------------------	
    if(xyzReconTracks){  
  
        bool isMC = event.GetContext().IsMC();	
    
      for (COMET::IReconObjectContainer::iterator t = xyzReconTracks->begin(); t != xyzReconTracks->end(); t++) {
          COMET::IHandle<COMET::IReconTrack> track = *t;
          if (!track) continue;
      
          COMET::IHandle<COMET::ITrackState> trackState = track->GetState();
          if (!trackState) continue;

          ISmrdIsoTrack *isoTrack = new((*fSmrdIsoTracks)[fNSmrdIsoTracks++]) ISmrdIsoTrack;
      
          COMET::IReconNodeContainer& nodes = track->GetNodes();    
      
      COMET::IHandle<COMET::ITrackState> frontState = TrackingUtils::GetFirstState(*track);
          if( !frontState ) continue;

          COMET::IHandle<COMET::ITrackState> backState = TrackingUtils::GetLastState(*track);
          if( !backState ) continue;			
      
          COMET::IHandle<COMET::IHitSelection> hits = track->GetHits();			
      
      if(!hits) continue;				
            
          isoTrack->AlgorithmName = track->GetAlgorithmName();
          isoTrack->Chi2      = track->GetQuality();
      
      //some of all contributed hits charges	
      isoTrack->EDeposit      = -9999.0;
            COMET::IHandle<COMET::IRealDatum> totalCharge =  track->Get< COMET::IRealDatum >("totalCharge");
      if(totalCharge)
        isoTrack->EDeposit = track->Get< COMET::IRealDatum >("totalCharge")->GetValue();
        
        
      //hits averaged time
      isoTrack->avgtime = -9999.0;
          COMET::IHandle<COMET::IRealDatum> smrdTimeRD = track->Get<COMET::IRealDatum>("averageHitTime");
            if (smrdTimeRD) 
        isoTrack->avgtime = smrdTimeRD->GetValue();
        
        
      
                  isoTrack->NHits 	= hits->size();
          isoTrack->NNodes 	= nodes.size();
      isoTrack->Status	= track->CheckStatus(track->kSuccess);
      isoTrack->NDOF		= track->GetNDOF();     
                        isoTrack->UniqueID	= track->GetUniqueID(); 
      
      //Check the Kalman refitting status
      std::string name("RECPACK_KalmanFilter");
      if((track->GetAlgorithmName().find(name) != std::string::npos) && (track->CheckStatus(track->kSuccess)))
        isoTrack->KalmanStatus = 1;
      else
          isoTrack->KalmanStatus = 0;
        
          isoTrack->FrontPos 	   = frontState->GetPosition();
          isoTrack->FrontPosVariance = frontState->GetPositionVariance();
    
          isoTrack->BackPos 	   = backState->GetPosition();
          isoTrack->BackPosVariance  = backState->GetPositionVariance();
    
          TVector3 backpos  	   = backState->GetPosition().Vect();				
          TVector3 frontpos 	   = frontState->GetPosition().Vect();
        
          //isoTrack->Direction = (backpos - frontpos).Unit();
          isoTrack->Direction         = trackState->GetDirection();    
          isoTrack->DirectionVariance = trackState->GetDirectionVariance();    
            
      //Calculate length of the track assuming straight line
      TVector3 diff = backpos-frontpos;
          isoTrack->Range     = diff.Mag();
    
          //Angles...
          isoTrack->ThetaAngle = trackState->GetDirection().Theta();
          isoTrack->PhiAngle   = trackState->GetDirection().Phi();
      
      //Set truth info
      int trueId 		= -999;
      TLorentzVector trueInitPos(-9999,-9999,-9999,-9999);
      TLorentzVector trueFinalPos(-9999,-9999,-9999,-9999);
      TLorentzVector trueInitMom(-9999,-9999,-9999,-9999);
      int truePDG 		=  0;
      int trueParentId 	= -999;
      double truePurity 	= -9999.0;					
      double trueEff 	= -9999.0;		
      
      if(isMC){
        trueId = TrackTruthInfo::GetG4TrajIDHits(*hits, trueEff, truePurity);
        
        COMET::IHandle<COMET::IG4TrajectoryContainer> trajectoryContainer 	= event.Get<COMET::IG4TrajectoryContainer>("truth/G4Trajectories");
        COMET::IHandle<COMET::IG4Trajectory> g4traj(0);
        if(trajectoryContainer)
          g4traj = trajectoryContainer->GetTrajectory(trueId);
                              
        if(g4traj){
          trueInitPos 	= g4traj->GetInitialPosition(); 
          trueFinalPos 	= g4traj->GetFinalPosition();
          trueInitMom 	= g4traj->GetInitialMomentum();
          truePDG		= g4traj->GetPDGEncoding();
          trueParentId	= g4traj->GetParentId();				
                         
        }
      }
      
      isoTrack->TrueInitPos 	= trueInitPos;
         
      gGeoManager->FindNode(trueInitPos.X(), trueInitPos.Y(), trueInitPos.Z());
      gGeoManager->GetCurrentNodeId();			
      isoTrack->TrueInitDet 	= COMET::IoaAnalysisUtils::PathToSubdetector(gGeoManager->GetPath());
      
      isoTrack->TrueFinalPos 	= trueFinalPos;
         
      gGeoManager->FindNode(trueFinalPos.X(), trueFinalPos.Y(), trueFinalPos.Z());
      gGeoManager->GetCurrentNodeId();
      isoTrack->TrueFinalDet 	= COMET::IoaAnalysisUtils::PathToSubdetector(gGeoManager->GetPath());								
      
      isoTrack->TrueInitMom 	= trueInitMom;
      isoTrack->TrueId 	= trueId;
      isoTrack->TruePDG 	= truePDG;																		   		
      isoTrack->TrueParentId  = trueParentId;
      //isoTrack->TrueParentPDG =;
      isoTrack->TrueHitPurity = truePurity;
      isoTrack->TrueHitEff = trueEff;
      
          
      }//end of tracks loop
    }//end of if(xyzTracks)

  return true;
}