import ROOT,os,sys,getopt
import rootUtils as ut
import shipunit as u

class FitTrackInfo(object):
  
  def __init__(self, tree, debug=0):
    self.tree     = tree
    self.debug    = debug
    self.count         = 0
    self.Momentum      = {}
    self.Direction     = {}
    self.Position      = {}
    self.__info          = {}
    self.Vertex        = None
    self.Doca          = None
    # 0 orig, 1 refit, -1 failed refit
    self.vertexEFlag    = 0 

  def clean(self): 
    self.count = 0
    self.Momentum.clear()
    self.Direction.clear()
    self.Position.clear()
    self.__info.clear()
    self.Vertex = None
    self.Doca   = None

  def Print(self):
    print "FitTrackInfo:"
    for tid in self.__info:
      print "\t", tid, "{:10.4f}".format(self.__info[tid]['Ndf']), "{:10.4f}".format(self.__info[tid]['Chi2']), 
      print "  pos:", "  ".join("{:10.4f}".format(self.Position[tid](ii)) for ii in range(0,3)),
      print "  mom:", "  ".join("{:10.4f}".format(self.Direction[tid](ii)*self.Momentum[tid]) for ii in range(0,3)), 
      print "  P:", "{:10.4f}".format(self.Momentum[tid])
      
  ## \brief returns list of keys #__info .
  #  \return list of MCtrackIDs of fitted tracks.
  def getTrIDs(self):
    trID = []
    for tid in self.__info:
      trID.append(tid)
    return trID
  
  def getNtracks(self) : 
    return len(self.__info)

  def getChi2Ndf(self, tid):
    return self.__info[tid]['Chi2']/self.__info[tid]['Ndf']
  
  def getNdf(self, tid):
    return self.__info[tid]['Ndf']
  
  def compareTracks(self, tid, Pos, Dir, Pval):
    false2 = (Pos==None or Dir==None or Pval==None)
    false1 = (not tid in self.__info)
    if (false2 and false1): return True
    if (false2  or false1): return False
    return ( (self.Direction[tid]==Dir) and (self.Position[tid]==Pos) and (self.Momentum[tid]==Pval) )

  
  def getPosDirPval(self, tid):
    if not tid in self.__info: return None, None, None
    return self.Position[tid], self.Direction[tid], self.Momentum[tid]
  
  def getVertex(self)  : 
    return self.Vertex
  
  def getDoca(self):
    return self.Doca
  
  def myVertex2(self, t1,t2):
    deltaPos =(self.Position[t1]-self.Position[t2])       # A1-A2
    dotDir   = self.Direction[t1].Dot(self.Direction[t2])   # a1.a2
    crossDir = self.Direction[t1].Cross(self.Direction[t2]) # a1xa2
    uPerpend = crossDir*(1./crossDir.Mag())     # (a1xa2)/|a1xa2| from a1 to a2

    minDist = deltaPos.Dot(uPerpend) # (A1-A2).(a1xa2)/|a1xa2|
    
    # A1 + a1*t1 + (minDist * uPerpend)  is  (A1 + a1*t1) projected to the plane: 
    #                             1) A2+a2*t2 belons to the plane,  
    #                             2) A1+a1*t1 is parallel to the plane
    # cross at t1,t2: A1+a1*t1+(minDist*uPerpend) = A2+a2*t2
    t2X   = self.Direction[t2].X()
    if (t2X == 0) : t2X = 0.00000000001
    a2a   = self.Direction[t2].Y()/t2X
    alpha = deltaPos - minDist*uPerpend
    nomin = alpha.Y() - a2a*alpha.X()
    denom = a2a*self.Direction[t1].X() - self.Direction[t1].Y()
    s1    = nomin/denom
    s2    = ( alpha.X() + self.Direction[t1].X()*s1 ) / t2X#self.Direction[t2].X()
    vec1 = self.Position[t1]+s1*self.Direction[t1]
    vec2 = self.Position[t2]+s2*self.Direction[t2]
    ave  = (vec1+vec2)*0.5
    dif  = vec1-vec2
    debugNeed = False
    if(abs(abs(minDist)-dif.Mag())>0.00000001):
      print "myVertex2 - problem:"
      debugNeed = True
    if(self.debug>2 or debugNeed):
      for tid in (t1,t2):
	print str(tid)+": Pos : ", "".join(str(self.Position[tid](ii)) +"   " for ii in range(0,3)),
	print "\t\tMom : ",       "".join(str(self.Direction[tid](ii))+"   " for ii in range(0,3))
      print "uPerpend: ","".join(str(uPerpend(ii))+"  " for ii in range(0,3))
    if(self.debug>0 or debugNeed):
      print "fit vertex: -> 1st poing : ", vec1.X(), vec1.Y(), vec1.Z()
      print "fit vertex: -> 2nd point : ", vec2.X(), vec2.Y(), vec2.Z()
      print "fit vertex: -> average   : ", ave.X(),  ave.Y(),  ave.Z()
      print "distance", abs(minDist), dif.Mag()
    return ave, abs(minDist)


  def readEvent(self):
    self.clean()
    indx = -1
    for atrack in self.tree.FitTracks:
    #  kill tracks outside fiducial volume
    #  if not checkFiducialVolume(sTree,key,dy): continue      
      fitStatus   = atrack.getFitStatus()
      if not fitStatus.isFitConverged() : continue    
      
      indx+=1
      mcTrID = self.tree.fitTrack2MC[indx]
      
      self.__info[mcTrID]      = {}
      self.__info[mcTrID]['Ndf']  = fitStatus.getNdf()
      self.__info[mcTrID]['Chi2'] = fitStatus.getChi2()

      fittedState = atrack.getFittedState()
      self.Momentum[mcTrID]  = fittedState.getMomMag()
      self.Direction[mcTrID] = fittedState.getDir()
      self.Position[mcTrID]  = fittedState.getPos()
    
    if(indx>0):
      if(indx==1):
	self.createVertex(self.__info.keys()[0], self.__info.keys()[1])
      else:
	pass
	#print "More than 2 fitterd tracks"
    return len(self.__info) 
  
  def createVertex(self, tid1, tid2, flag=0):
    if( (tid1 in self.__info) and (tid2 in self.__info) ):
      self.Vertex, self.Doca = self.myVertex2(tid1, tid2)
      self.vertexEFlag = flag
    
    
  def addNewTrack(self, mcTrID, position, direction, Pval, ndf, chi2, verb = True):
    if (verb and mcTrID in self.Momentum):
      print "FotTrackInfo WARNING - trID ", mcTrID, "already filled! Will rewrite all records for this trID!"
    
    self.__info[mcTrID]={}
    self.__info[mcTrID]['Ndf']  = ndf
    self.__info[mcTrID]['Chi2'] = chi2

    self.Momentum[mcTrID]  = Pval
    self.Direction[mcTrID] = direction
    self.Position[mcTrID]  = position
      
  
  def deleteTrack(self, mcTrID):
    if mcTrID in self.__info:
      del self.__info[mcTrID]
    if mcTrID in self.Momentum:
      del self.Momentum[mcTrID]
    if mcTrID in self.Direction:
      del self.Direction[mcTrID]
    if mcTrID in self.Position:
      del self.Position[mcTrID]