import ROOT,os,sys,getopt
import rootUtils as ut
import shipunit as u
import numpy as n
import collections

from ShipGeoConfig import ConfigRegistry
import shipDet_conf

import BookHistos

from MCTrackInfo  import MCTrackInfo
from StrawHits    import StrawHits
from FitTrackInfo import FitTrackInfo

import RecoSettings


class LoopAnalyse(object):  
  def __init__(self, tree, pref, geoFile, debug=0, 
	       plotField          = False,
	       plotTracking       = False, 
	       plotVertex         = False,
	       plotVertexDetailed = False,
	       plotMomentum       = False,
	       compareFits        = False):
    self.tree    = tree
    self.debug   = debug

    self.modules   = None
    self.ShipGeo   = None    
    self.run       = None
    self.createGeom(pref, geoFile)
        
    self.stat = collections.OrderedDict()

    self.hmagnet   = {}
    self.htrck     = {}
    self.hvtx      = {}
    self.hmom      = {}
    self.hcmfits   = {}
    self.plotField    = False
    self.plotTracking = False
    self.plotVertex   = False
    self.plotVertexCh = False
    self.plotMomentum = False
    self.compareFits  = False
    
    
    if plotField:
      self.plotField = True
      BookHistos.bookMFieldHistos(self.hmagnet)
    if plotTracking:
      self.plotTracking = True
      BookHistos.bookTrackingHistos(self.htrck)
    if plotVertex:
      self.plotVertex = True
      BookHistos.bookVertexHistos(self.hvtx)
    if plotVertexDetailed:
      self.plotVertexCh = True
      BookHistos.bookVertexHistos(self.hvtx,"mUp")
      BookHistos.bookVertexHistos(self.hvtx,"mDn")
      
    if plotMomentum:
      self.plotMomentum = True
      BookHistos.bookMomentumHistos(self.hmom)
    if compareFits:
      self.compareFits = True
      BookHistos.bookCompareFitHistos(self.hcmfits)
      if plotVertex  : BookHistos.bookVertexHistos(self.hvtx, "NEW")
      if plotMomentum: BookHistos.bookMomentumHistos(self.hmom, "NEW")

    self.MCTracks  = MCTrackInfo(self.tree, 0)#self.debug)
    self.StrHits   = StrawHits(self.tree, self.modules, self.ShipGeo.straw.resol, self.debug, self.hmagnet, self.ShipGeo)
    self.FitTracks = FitTrackInfo(self.tree, 0)#self.debug)
    # FIXneeded - to introduce one more FixTracks!

      
  def createGeom(self,pref, geoFile):
    self.ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = pref )
    self.run = ROOT.FairRunSim()
    self.modules = shipDet_conf.configure(self.run,self.ShipGeo)
    tgeom = ROOT.TGeoManager("Geometry", "Geane geometry")
    gMan  = tgeom.Import(geoFile)
    geoMat =  ROOT.genfit.TGeoMaterialInterface()
    ROOT.genfit.MaterialEffects.getInstance().init(geoMat)
    volDict = {}
    i=0
    for x in ROOT.gGeoManager.GetListOfVolumes():
      volDict[i]=x.GetName()
      i+=1

    bfield = ROOT.genfit.BellField(self.ShipGeo.Bfield.max ,self.ShipGeo.Bfield.z,2, self.ShipGeo.Yheight/2.)
    fM = ROOT.genfit.FieldManager.getInstance()
    fM.init(bfield)



  def CleanStat(self):
    self.stat['total'] = 0
    self.stat['HNLnTr']   = 0
    self.stat['HNLzAcc']  = 0
    self.stat['HNLeAcc']  = 0
    self.stat['twoFitsO']  = 0
    self.stat['twoFitsN']  = 0
    self.stat['twoFitsNdfO']  = 0
    self.stat['twoFitsNdfN']  = 0
    
  def writeHistos(self):
    ut.writeHists(htrck,"LoopA.root")
  
  
  
  
  def hitMCtrackCorellations(self, weight=1):
    # FIXneeded - add HNL vertex distribution plots!
    if ( not self.plotTracking): 
      print "hitMCtrackCorellations does nothing since BookHistos.bookTrackingHistos() was not called!"
      return
    for tid in self.StrHits.getTrIDs():            # only multi-hits>25 tracks
      if ( self.StrHits.checkVetoHits(tid) > 0)  :  continue  # remove ones with vetoStattion
      vec3h = self.StrHits.getStartHit(tid)
      vec3p = self.MCTracks.getTrackPropagation(tid, vec3h.Z())
      if(self.debug>0):
	print "extr.", tid, "\t", vec3p.X(), vec3p.Y(), vec3p.Z(),
	print "\tcharge:", self.MCTracks.getCharge(tid), "\tdx=", vec3p.X()-vec3h.X(), "dy=", vec3p.Y()-vec3h.Y()
      dR       = ROOT.TMath.Sqrt( (vec3p.X()-vec3h.X())**2 + (vec3p.Y()-vec3h.Y())**2 )
      self.htrck['dRvsPz'].Fill(self.MCTracks.getMomentum(tid).Z(), dR, weight)
      self.htrck['dRvsZhit'].Fill(vec3h.Z(), dR, weight)
      charge   = self.MCTracks.getCharge(tid)
      theRatio = abs(vec3p.Y())/(vec3h.Z()-self.MCTracks.getHNLdecayVertex().Z())
      HNLflag  = tid in self.MCTracks.getHNLdecayTrIDs()
      if HNLflag:
	self.htrck['YMC2dZMCVer'].Fill(theRatio,weight)
      if charge<0:
	self.htrck['dxdyALLTrHm'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)  
	if HNLflag:
	  self.htrck['dxdyHNLTrHm'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)  
      elif charge>0:
	self.htrck['dxdyALLTrHp'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)  
	if HNLflag:
	  self.htrck['dxdyHNLTrHp'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)  
    
    
    
  def fitMCvertexCorellations(self, weight=1, old=True, cut=False, detailed=False):
    if( cut and detailed ) : print "fitMCvertexCorellations can't plot detailed vertex for cut! FIX me!"
    vMC = self.MCTracks.getHNLdecayVertex()
    if old :
      vFt = self.FitTracks.getVertex()
    else :
      vFt = self.StrHits.getReFitVertex()
    if (not vFt or not vMC) : return
    
    if(self.plotVertex or self.plotVertexCh):
      distanceToTracker = self.ShipGeo.TrackStation1['z']-vMC.Z()
      coord = ['x','y','z']
      iz = 1
      while ( (iz<11) and (distanceToTracker>iz*500) ):
	iz = iz+1
      
      pref=""
      prefCut=""
      if( cut )    : prefCut='ndf'
      if( not old ): pref='NEW'
      for i in range(0,3):
	delta = vMC(i)-vFt(i)
	if( self.plotVertex ):
	  hname = 'vd'+coord[i]+'2'+prefCut+pref
	  self.hvtx[hname].Fill(distanceToTracker, delta, weight)
	  if (not cut):
	    if (iz<11):
	      hname='vd'+coord[i]+'2_'+str(iz)+pref
	      self.hvtx[hname].Fill(delta, weight)
	
	if( self.plotVertexCh and old ):
	  sign = self.MCTracks.checkVertexUpDown()
	  if ( sign==0 ): continue
	  if ( sign<0) :
	    pref = "mUp"
	  else:
	    pref = "mDn"
	  hname = 'vd'+coord[i]+'2'+pref
	  self.hvtx[hname].Fill(distanceToTracker, delta, weight)
      if( pref=='NEW') :
	doca = self.StrHits.getStepDoca(RecoSettings.VertexExtrSteps-1)
	if ( doca!=None ) : self.hvtx['doca'+pref].Fill(doca)
	doca0 = self.StrHits.getStepDoca(0)
	if ( doca0!=None ) : self.hvtx['doca0'+pref].Fill(doca0)
	if ( doca0!=None and doca0!=0):
	  self.hvtx['docaRel'+pref].Fill(1,1.)
	  for ii in range (1, RecoSettings.VertexExtrSteps):
	    doca = self.StrHits.getStepDoca(ii)
	    #print "...", doca, doca/doca0
	    if ( doca!=None ) : 
	      self.hvtx['docaRel'+pref].Fill(ii+1, doca/doca0)
    return vFt





  def doCompareFits(self, oldStyle=True, ellCut=False, weight=1):
    ntracks = self.StrHits.FitTracks(oldStyle) # does re-fitting
    if (self.debug>0) : self.StrHits.PrintNewTracks()

    fitOldTrIDs = self.FitTracks.getTrIDs()
    fitNewTrIDs = self.StrHits.getReFitTrIDs()
    
    for tid in fitOldTrIDs:
      self.hcmfits['Chi2NDFold'].Fill(self.FitTracks.getChi2Ndf(tid))
      self.hcmfits['NDFold'].Fill(self.FitTracks.getNdf(tid))
    for tid in fitNewTrIDs:
      self.hcmfits['Chi2NDFnew'].Fill(self.StrHits.getReFitChi2Ndf(tid))
      self.hcmfits['NDFnew'].Fill(self.StrHits.getReFitNdf(tid))
      
    for tid in self.StrHits.getTrIDsALL(): # hists with MCTrID>0 
      told = ( tid in fitOldTrIDs)
      tnew = ( tid in fitNewTrIDs)
      if( (not told) and (not tnew) ):
	self.hcmfits['FitCorrespondence'].Fill(0)
      elif( (not told) and tnew):
	self.hcmfits['FitCorrespondence'].Fill(1)
      elif( told and (not tnew) ):
	self.hcmfits['FitCorrespondence'].Fill(2)
      else:
	eq = self.StrHits.compareFitTracks(tid, self.FitTracks)
	if eq:
	  self.hcmfits['FitCorrespondence'].Fill(5)
	else:
	  self.hcmfits['FitCorrespondence'].Fill(6)
    self.hcmfits['dNFoldFNnew'].Fill(len(fitOldTrIDs), len(fitNewTrIDs))
    return ntracks




  def doPlotMomentum(self, weight, old=True):
    tids = []
    pfix = ""
    if old:
      tids  = self.FitTracks.getTrIDs()
    else:
      tids  = self.StrHits.getReFitTrIDs()
      pfix  = "NEW" 
      
    for tid in tids:
      r,d,pFit = None, None, None
      chi2ndf,ndf  = None, None
      if old:
	r,d,pFit = self.FitTracks.getPosDirPval(tid)
	ndf      = self.FitTracks.getNdf(tid)
	chi2ndf  = self.FitTracks.getChi2Ndf(tid)
      else:
	r,d,pFit = self.StrHits.getReFitPosDirPval(tid)
	ndf      = self.StrHits.getReFitNdf(tid)
	chi2ndf  = self.StrHits.getReFitChi2Ndf(tid)
      if( not tid in self.MCTracks.getTrIDs() ): continue
      pMC = self.MCTracks.getMomentum(tid).Mag()
      dPP = (pMC - pFit)/pMC
      hname = 'delPOverP0'+pfix
      self.hmom[hname].Fill(chi2ndf, dPP, weight)
      hname = 'delPOverP1'+pfix
      self.hmom[hname].Fill(pMC, dPP, weight)
      
      if( ndf<RecoSettings.trackMinNofHits ): continue
      hname = 'delPOverP2'+pfix
      self.hmom[hname].Fill(pMC, dPP, weight)
  
      if( not tid in self.MCTracks.getHNLdecayTrIDs() ): continue
      hname = 'delPOverP3'+pfix
      self.hmom[hname].Fill(pMC, dPP, weight)

      iz = 1
      while ( (iz<6) and (pMC>iz*10) ):
	iz = iz+1
      if (iz<6):
	hname = 'delPOverP3_'+str(iz)+pfix
	self.hmom[hname].Fill(dPP, weight)













  def PrintStat(self):
    for stage in self.stat:
      print "{0:15}{1:10}".format(stage, self.stat[stage])




################################################################################################
  def Loop(self, stat=-1):
    self.CleanStat()
    
    if(stat==-1):
      nEvents = self.tree.GetEntries()
    else:
      nEvents = min(self.tree.GetEntries(),stat)

    # --------------- cut related values ------------------------------------------------------#
    # FIXneeded -  should be done depending on RecoSettings.trackMinNofStations!
    # self.MCTracks.setAccPlaneZ(3380*u.cm)  # - after 3rd Station
    self.MCTracks.setAccPlaneZ(2500*u.cm)  #- before 1st Station
    
    
    
    # -----------------------------------------------------------------------------------------#
    # -----------------------------------------------------------------------------------------#
    for n in range(nEvents):
      
      rc = self.tree.GetEntry(n)
      print
      print " HO GETTATO L'ENTRY ", n
      print
      self.stat['total'] +=1
            
      # ------------------- MCTracks -----------------------------------------------------------#
      nMCPrimTracks = self.MCTracks.readEvent()
      if( self.debug>0) : 
	print "\n\nevent ", n, "===================================================="
	self.MCTracks.PrintHNL()
      
      #=>
      if( nMCPrimTracks!=2 ) : continue
      self.stat['HNLnTr']+=1
      
      #=>
      # FIXneeded - should be setting from ShipGeo here!  
      if( (self.MCTracks.getHNLdecayVertex().Z()<-1900*u.cm) \
	or (self.MCTracks.getHNLdecayVertex().Z()>RecoSettings.VertexMaxZcut)  ) : continue
      self.stat['HNLzAcc']+=1
      
      #=>
      if( (not self.MCTracks.checkEllipticAcc()) ) : continue 
      self.stat['HNLeAcc'] +=1
      if( self.debug>1 ): print "MC event accepted"
            
            
      weight = self.MCTracks.getEventWeight()
      #weight = 1

      # ------------------- FitTracks -----------------------------------------------------------#
      nFitTracks = self.FitTracks.readEvent()
      if(self.plotMomentum):self.doPlotMomentum(weight, old=True)
      if(self.debug>0) : self.FitTracks.Print()
      
      printMCVert = False
      if (nFitTracks==2):
	self.stat['twoFitsO']+=1 
	v=self.fitMCvertexCorellations(weight, old=True, cut=False, detailed=self.plotVertexCh)
	printMCVert = True
	print "\n\nOld vertex ", v.X(), v.Y(), v.Z()
	
	ndfok = True
	for tid in self.FitTracks.getTrIDs() :
	  ndfok = ( ndfok and (self.FitTracks.getNdf(tid)>RecoSettings.trackMinNofHits) )
	  if ( not ndfok ) : break
	if ndfok : # FIXneeded - not optimal to call it once more :(
	  self.stat['twoFitsNdfO']+=1
	  v=self.fitMCvertexCorellations(weight, old=True, cut=True, detailed=False)
	  print "the old vertex passes ndf cut!"

      # -------------------- StrawHits --------------------------------------------------------------#
      if ( self.compareFits or self.plotTracking ) :
	self.StrHits.readEvent()
	
	
      # -------------------- ReFit --------------------------------------------------------------#
      nNewFitTracks = 0
      if self.compareFits : 	
	nNewFitTracks = self.doCompareFits(oldStyle=True, ellCut=False, weight=weight)
	if(self.plotMomentum):self.doPlotMomentum(weight, old=False)
	
      if (nNewFitTracks==2): 
	self.stat['twoFitsN']+=1 
	v=self.fitMCvertexCorellations(weight, old=False)
	print "New vertex ", v.X(), v.Y(), v.Z()
	printMCVert = True

	ndfok = True
	for tid in self.StrHits.getReFitTrIDs() :
	  ndfok = ( ndfok and (self.StrHits.getReFitNdf(tid)>RecoSettings.trackMinNofHits) )
	  if ( not ndfok ) : break
	if ndfok : # FIXneeded - not optimal to call it once more :(
	  self.stat['twoFitsNdfN']+=1
	  v=self.fitMCvertexCorellations(weight, old=False, cut=True, detailed=False)
	  print "the new vertex passes ndf cut!"
      

      if printMCVert : self.MCTracks.PrintHNL()



      if(self.plotTracking) : self.hitMCtrackCorellations(weight)
      
      
	
	
    # end event loop
    #=======================================
    print "here -> ", self.plotField
    if self.plotField    : BookHistos.plotMFieldHistos(self.hmagnet)
    if self.plotTracking : BookHistos.plotMCTrackHitCorrelation(self.htrck)
    if self.plotVertex   : BookHistos.plotVertexHistos(self.hvtx)
    if self.plotVertexCh :  
      BookHistos.plotVertexHistos(self.hvtx,"mUp")
      BookHistos.plotVertexHistos(self.hvtx,"mDn")
    if self.plotMomentum : BookHistos.plotMomentumHistos(self.hmom)
    if self.compareFits  :
      BookHistos.plotCompareFitHistos(self.hcmfits)
      if self.plotVertex   : BookHistos.plotVertexHistos(self.hvtx,'NEW')
      if self.plotMomentum : BookHistos.plotMomentumHistos(self.hmom, 'NEW')
      
    #=======================================
    self.PrintStat()