#!/bin/env python

from ROOT import *
import hddm_s
import math
import numpy

EMC = numpy.array([0] * 5, dtype=float)
thetaMC = numpy.array([0] * 5, dtype=float)
phiMC = numpy.array([0] * 5, dtype=float)
origMC = numpy.array([0] * 4, dtype=float) 
nTP = numpy.array([0], dtype=int)
ETP = numpy.array([0] * 32, dtype=float)
dEdxTP = numpy.array([0] * 32, dtype=float)
phiTP = numpy.array([0] * 32, dtype=float)
primaryTP = numpy.array([0] * 32, dtype=int)
ptypeTP = numpy.array([0] * 32, dtype=int)
pthetaTP = numpy.array([0] * 32, dtype=float)
pphiTP = numpy.array([0] * 32, dtype=float)
rTP = numpy.array([0] * 32, dtype=float)
tTP = numpy.array([0] * 32, dtype=float)
trackTP = numpy.array([0] * 32, dtype=int)
nhitTP = numpy.array([0] * 32, dtype=int)
dETP = numpy.array([0] * 32, dtype=float)
dEhitTP = [numpy.array([0] * 32, dtype=float),
           numpy.array([0] * 32, dtype=float),
           numpy.array([0] * 32, dtype=float),
           numpy.array([0] * 32, dtype=float),
           numpy.array([0] * 32, dtype=float)]
sectTP = [numpy.array([0] * 32, dtype=int),
          numpy.array([0] * 32, dtype=int),
          numpy.array([0] * 32, dtype=int),
          numpy.array([0] * 32, dtype=int),
          numpy.array([0] * 32, dtype=int)]

ftree = TFile("tpoltree.root", "recreate")
tpoltree = TTree("tpoltree", "")
tpoltree.Branch("EMC", EMC, "EMC[5]/D")
tpoltree.Branch("thetaMC", thetaMC, "thetaMC[5]/D")
tpoltree.Branch("phiMC", phiMC, "phiMC[5]/D")
tpoltree.Branch("origMC", origMC, "origMC[4]/D")
tpoltree.Branch("nTP", nTP, "nTP/L")
tpoltree.Branch("ETP", ETP, "ETP[nTP]/D")
tpoltree.Branch("dEdxTP", dEdxTP, "dEdxTP[nTP]/D")
tpoltree.Branch("phiTP", phiTP, "phiTP[nTP]/D")
tpoltree.Branch("primaryTP", primaryTP, "primaryTP[nTP]/L")
tpoltree.Branch("ptypeTP", ptypeTP, "ptypeTP[nTP]/L")
tpoltree.Branch("pthetaTP", pthetaTP, "pthetaTP[nTP]/D")
tpoltree.Branch("pphiTP", pphiTP, "pphiTP[nTP]/D")
tpoltree.Branch("rTP", rTP, "rTP[nTP]/D")
tpoltree.Branch("tTP", tTP, "tTP[nTP]/D")
tpoltree.Branch("trackTP", trackTP, "trackTP[nTP]/L")
tpoltree.Branch("nhitTP", nhitTP, "nhitTP[nTP]/L")
tpoltree.Branch("dETP", dETP, "dETP[nTP]/D")
for i in range(0,5):
   tpoltree.Branch("dE{0}TP".format(i), dEhitTP[i], "dE{0}TP[nTP]/D".format(i))
   tpoltree.Branch("sect{0}TP".format(i), sectTP[i], "sect{0}TP[nTP]/D".format(i))

mElectron = 0.51099894621e-3

def hist(flist, maxcount=1e30):
   evcount = 0
   elcount = 0
   for fin in flist:
      for r in hddm_s.istream(fin):
         o = r.getOrigins()
         origMC[0] = o[1].vx
         origMC[1] = o[1].vy
         origMC[2] = o[1].vz + 1518.20
         origMC[3] = o[1].t
         m = r.getMomenta()
         for i in range(0,len(m)):
            EMC[i] = m[i].E
            thetaMC[i] = math.atan2((m[i].px**2 + m[i].py**2)**0.5, m[i].pz)
            phiMC[i] = math.atan2(m[i].py, m[i].px)
         for i in range(len(m),5):
            EMC[i] = 0
            thetaMC[i] = 0
            phiMC[i] = 0
         t = r.getTpolTruthPoints()
         for i in range(0,len(t)):
            ETP[i] = t[i].E
            dEdxTP[i] = t[i].dEdx
            phiTP[i] = t[i].phi
            primaryTP[i] = t[i].primary
            ptypeTP[i] = t[i].ptype
            pthetaTP[i] = math.atan2((t[i].px**2 + t[i].py**2)**0.5, t[i].pz)
            pphiTP[i] = math.atan2(t[i].py, t[i].px)
            rTP[i] = t[i].r
            tTP[i] = t[i].t
            trackTP[i] = t[i].track
            nhitTP[i] = 0
            dETP[i] = 0
            for j in range(0,len(dEhitTP)):
               dEhitTP[j][i] = 0
               sectTP[j][i] = -1
         nTP[0] = len(t)
         for h in r.getTpolTruthHits():
            for i in range(0,len(t)):
               if h.itrack == trackTP[i]:
                  dETP[i] += h.dE
                  if nhitTP[i] < len(dEhitTP):
                     dEhitTP[nhitTP[i]][i] = h.dE
                     sectTP[nhitTP[i]][i] = h.getAttribute("sector")
                     nhitTP[i] += 1
         tpoltree.Fill()
         for p in r.getProducts():
            if p.pdgtype == 11:
               elcount += 1
         evcount += 1
         if evcount == maxcount:
            break
   tpoltree.Write()
   return evcount, elcount

def asym(var, par):
   phi = var[0] * asym_scale
   return par[0] * (1 + par[1] * math.cos(2*(phi - par[2])))

def fasym(sectors=0):
   global asym_scale
   if sectors:
      f = TF1("fasym", asym, 1, 33, 3)
      asym_scale = math.pi / 16
   else:
      f = TF1("fasym", asym, -math.pi, math.pi, 3)
      asym_scale = 1
   f.SetParameter(0, 100)
   f.SetParameter(1, 0.2)
   f.SetParameter(2, 0.1)
   return f