import ROOT,os,sys,getopt
import rootUtils as ut
import shipunit as u
import RecoSettings
########################################################################
## For a single event stores HNL, its decay products and several other tracks (if requested) 
#  from ShipMCTracks collection of a MCtrack branch (see $FAIRSHIP/shipdata/ShipMCTrack.h)\n
#
#
class MCTrackInfo(object):
  """ More description here"""
  def __init__(self, tree, debug=0):
    ## root tree to be read.
    self.__tree        = tree
    ## debug level [0,3].
    self.__debug       = debug
    ## z of the last full acceptance plane (to be used in #checkEllipticAcc), must be set later.
    self.__zAcc          = 3380*u.cm
    ## event weight. Read in #readEvent as weight of HNL.
    self.__weight        = 0
    ## HNL decay vertex coordinates (TVector3). Assigned in #readEvent()
    self.__HNLdecayCoord = None
    ## MCTrID (index of the track in MCtrack list) for \b HNL \b daughter products. 
    self.__decayProdTrID = []
    ## Production vertices of several stored products (\b not \b only HNL daughter) as {MCTrID : TVector3}.
    self.__productVertex = {}
    ## Momentum of several stored products (\b not \b only HNL daughter) as {MCTrID : TVector3}. Correspond to $__productVertex.
    self.__productMoment = {}
    ## Charge of several stored products (\b not \b only HNL daughter) as {MCTrID : value}. Correspond to $__productVertex.
    self.__productCharge = {}
    ## Basic information on the particle (pdg, mother pdg,... to be added - z decay!)
    self.__productInfo   = {}
  
  
  
  
  
  ########################################################################
  ##\brief clears all dictionaries and lists. To be called inside #readEvent.
  def __clean(self):
    self.__weight        = 0
    self.__HNLdecayCoord = None
    self.__decayProdTrID = []
    self.__productVertex.clear()
    self.__productMoment.clear()
    self.__productCharge.clear()
    self.__productInfo.clear()

  ########################################################################
  ##\brief set z of the last full acceptance plane (to be used in #checkEllipticAcc), depends on #RecoSettings. trackMinNofStations.
  # \param z - the value [cm].
  def setAccPlaneZ(self, z):
    self.__zAcc = z
    

  ########################################################################
  ##\brief prints HNL decay products
  def PrintHNL(self):
    print type(self).__name__, " : "
    print "\tdecay vertex: ", "  ".join("{:10.4f}".format(self.__HNLdecayCoord(ii)) for ii in range(0,3)),
    print "\n\tdecay products (trID, pdg, charge, momentum):"
    for tid in self.__decayProdTrID :
      print "\t", "{0:6}{1:10}{2:6.0f}".format(tid, self.__productInfo[tid]['pdg'], self.__productCharge[tid]),
      print "\t", " ".join("{:10.4f}".format(self.__productMoment[tid](ii)) for ii in range(0,3))


        
  ########################################################################
  ## \brief returns event weight.
  #  \return value of #__weight
  def getEventWeight(self):
    return self.__weight


  ########################################################################
  ## \brief returns \b HNL \b decay \b vertex coordinates (TVector3 #__HNLdecayCoord).
  #  \return TVector3 #__HNLdecayCoord 
  def getHNLdecayVertex(self):
    return self.__HNLdecayCoord
  
  
  ########################################################################
  ## \brief returns list of MCtrIDs (index in MCTrack collection) fot \b HNL \b decay \b products.
  #  \return list of MCtrIDs (index in MCTrack collection) fot \b HNL \b decay \b products.
  def getHNLdecayTrIDs(self):
    trID = []
    for tid in self.__decayProdTrID:
      trID.append(tid)
    return trID
  
  
  ########################################################################
  ## \brief returns list of all MCtrIDs read by #readEvent and #readTrack - \b not \b only HNL decay products.
  #  \return list of MCtrIDs (index in MCTrack collection) fot all tracks read by #readEvent and #readTrack.
  def getTrIDs(self):
    trID = []
    for tid in self.__productMoment:
      trID.append(tid)
    return trID


  ########################################################################
  ## \brief returns charge of a particle with the given trID (index in MCTrack collection). May return None.
  #  Tries to get charge from #__productCharge; 
  #  if it is not available, tries to add the track calling #readTrack. If fails, returns None.
  #  \param  trID - index of the MCTrack
  #  \return charge of the track or None
    def getCharge(self, trID):
      if ( not trID in self.__productCharge ): self.readTrack(trID)
      return self.__productCharge[trID]
  
  

  ########################################################################
  ## \brief returns momentum (TVector3) for a particle with the given trID (index in MCTrack collection). May return None.
  #  Tries to get momentum from #__productMoment; 
  #  if it is not available, tries to add the track calling #readTrack. If fails, returns None.
  #  \param  trID - index of the MCTrack
  #  \return momentum of the track (TVector3) or None
  def getMomentum(self, trID):
    if ( not trID in self.__productMoment ): self.readTrack(trID)
    return self.__productMoment[trID]
  
  

  
  ########################################################################
  ## \brief checks if HNL decay vertex (and both HNL daughter tracks, if tight cut) are in elliptic acceptance at #__zAcc.
  #  For tight cut each HNL daughter is propagated to #__zAcc (see #getTrackPropagation) 
  #  and checked if it in the acceptance (see #RecoSettings .checkEllipticAcc). Also checks HNL decay vertex.
  #  \param tight - if true, not only HNL vertex but also propagated tracks are checked.
  #  \return True if the vertex (and both tracks, if tight cut) are in the acceptance.
  def checkEllipticAcc(self, tight=True):
    ok = RecoSettings.checkEllipticAcc(self.__HNLdecayCoord)
    if ( not ok ) : return False
    if ( tight ) :
      for tid in self.__decayProdTrID :
	v3 = self.getTrackPropagation(tid, self.__zAcc)
	ok = ( ok and RecoSettings.checkEllipticAcc(v3) )
	if ( not ok ) : break
    return ok
  
  
  
  
  ########################################################################
  ## \brief for Tracker Performance studies. Returns -1, 0, 1 depending on vertex topology.
  #  defines direction in y: 0 if both momY the same direction, -1 if (negative up, positive down).
  #  \return 0 if both momY the same direction, -1 if (negative up, positive down).
  def checkVertexUpDown(self):
    theCase = 0 # undefined
    if ( (not self.__productVertex) or (len(self.__decayProdTrID)<2) ): return theCase
    tid1 = self.__decayProdTrID[0]
    tid2 = self.__decayProdTrID[1]
    if ( self.__productMoment[tid1].Y()*self.__productMoment[tid2].Y() >=0 ): return theCase
    
    theCase = self.__productCharge[tid1]*self.__productMoment[tid1].Y()/abs(self.__productMoment[tid1].Y())
    
    return theCase


  ########################################################################
  ## \brief for Tracker Performance studies. Returns TVector3.
  #  Linear propagation of a given track to a plane (0,0,z) perendicular to Z-axis.
  #  \param trID - track index (MCtrID).
  #  \param    z - coordinate of the plane. 
  #  \return coordinates of the point of track crossing the given plalne (TVector3). 
  def getTrackPropagation(self, trID, z):
    if ( not trID in self.__productMoment ): self.readTrack(trID)
    xy=[0.,0.]
    for ii in [0,1]:
      xy[ii] = self.__productVertex[trID][ii]+self.__productMoment[trID][ii]/self.__productMoment[trID][2]*(z-self.__productVertex[trID][2])
    vec3 = ROOT.TVector3(xy[0],xy[1],z)
    return vec3




  ########################################################################
  ## \brief reads one additional track with trID from MCTrack collection. Returns 1 if success, 0 otherwise.
  #  fills #__productVertex, #__productMoment, #__productCharge dictionaries.
  #  \param trID - track index (MCtrID).
  #  \return 1 if successful, 0 if trID is not found and the dictionaries are not modified.
  def readTrack(self, trID):
    x = None
    try:
      x=self.__tree.MCTrack[trID]
    except:
      pass
    if (not b) : return 0
    
    vec3 = ROOT.TVector3(x.GetPx(),x.GetPy(),x.GetPz())
    self.__productMoment[trID] = vec3
    vec3 = ROOT.TVector3(x.GetStartX(),x.GetStartY(),x.GetStartZ())
    self.__productVertex[trID] = vec3
    self.__productCharge[trID] = RecoSettings.chargePDG(x.GetPdgCode())
    self.__productInfo[trID] ={}
    self.__productInfo[trID]['pdg'] = x.GetPdgCode()
    mid  = x.GetMotherId()
    if(mid>=0) : self.__productInfo[trID]['mpdg'] = self.__tree.MCTrack[mid].GetPdgCode()
    else       : self.__productInfo[trID]['mpdg'] = 0
    if( self.__debug>0 ):
	print "additional MC track for trID", trID
	print "\tpdgid/charge/mother: ", x.GetPdgCode(), self.__productCharge[trID], mid, "(pdf=",self.__productInfo[trID]['mpdg'], ")"
	print "\torigin:   ", self.__productVertex[trID].X(), self.__productVertex[trID].Y(), self.__productVertex[trID].Z()
	print "\tmomentum: ", self.__productMoment[trID].X(), self.__productMoment[trID].Y(), self.__productMoment[trID].Z()
    return 1




    
  ########################################################################
  ## \brief reads \b HNL \b decay \b products from MCTrack collection (\b two at the moment!). Returns number of daughters.
  #  Fills #__HNLdecayCoord and #__productVertex, #__productMoment, #__productCharge dictionaries for \b HNL \b daughters.
  #  \return number of daughters (loop stops at daughter==2 at the moment!).
  def readEvent(self):
    self.__clean()
    if (self.__debug>0 ):
      print "sTree.MCTrack", self.__tree.MCTrack, "size:", len(self.__tree.MCTrack)
    daughter = 0  
    tid      = 0
    #/home/kkuzn/SHIP/ShipSoft/tmp/FairShip/shipdata/ShipMCTrack.h
    for x in self.__tree.MCTrack:
      if daughter>1:
	break
      mid = x.GetMotherId()
      Z   = x.GetStartZ()
      if( Z>5000 ): Z=5000
      if mid==0: # #NB# here I suppose HNL is record #1 always!
	self.__weight = x.GetWeight()
	if self.__debug>0: print "HNL\t\t", tid,"\t", x.GetStartX(), x.GetStartY(), x.GetStartZ(), "\t\t", x.GetPx(), x.GetPy(), x.GetPz() 
      elif( mid==1 ): # HNL decay products
	self.__decayProdTrID.append(tid)
	vec3 = ROOT.TVector3(x.GetPx(),x.GetPy(),x.GetPz())
	self.__productMoment[tid] = vec3
	vec3 = ROOT.TVector3(x.GetStartX(),x.GetStartY(),x.GetStartZ())
	self.__productVertex[tid] = vec3
	self.__productCharge[tid] = RecoSettings.chargePDG(x.GetPdgCode())
	self.__productInfo[tid] ={}
	self.__productInfo[tid]['pdg']  = x.GetPdgCode()
	self.__productInfo[tid]['mpdg'] = self.__tree.MCTrack[mid].GetPdgCode()
	if( self.__debug>0 ): print  x.GetPdgCode(), "\t\t", tid,"\t", x.GetStartX(), x.GetStartY(), x.GetStartZ(), "\t\t", x.GetPx(), x.GetPy(), x.GetPz() 
	if daughter<1: 
	  self.__HNLdecayCoord = ROOT.TVector3(x.GetStartX(), x.GetStartY(), x.GetStartZ())
	else:
	  if self.__HNLdecayCoord.Z()!=x.GetStartZ():
	    self.__HNLdecayCoord.SetXYZ(0,0,-8899)
	daughter+=1
      tid+=1
      
    return daughter