import ROOT
class electron():
    def __init__(self,eid):
        self.positions=[]
        self.velocities=[]
        self.wFields=[]
        self.id=eid
        self.induced=[]
        self.times=[]
        self.force=[]
        self.eforce=[]

    def addState(self,x,y,z,vx,vy,vz,wx,wy,wz,t):
        self.positions.append([x,y,z])
        self.velocities.append([vx,vy,vz])
        self.wFields.append([wx,wy,wz])
        self.times.append(t)

    def addForce(self,fx,fy,fz):
        self.force.append([fx,fy,fz])

    def addEForce(self,fx,fy,fz):
        self.eforce.append([fx,fy,fz])

    def calcRamo(self):
        for i in range(len(self.wFields)):
            field=self.getWField(i)
            vel=self.getVel(i)
            if field!=field:
                print 'field NAN (calc ramo)'
            if vel!=vel:
                print 'vel NAN'
            #vel.SetX(0)
            #vel.SetY(0)
            self.induced.append(-field*vel)
    def calcRamoLast(self):
            field=self.getWField(-1)
            vel=self.getVel(-1)
            #vel.SetX(0)
            #vel.SetY(0)
            self.induced.append(-field*vel)
            return self.induced[-1]
        
    def getWField(self,i):
        for j in range(3):
            if self.wFields[i][j]!=self.wFields[i][j]:
                #print 'field NAN',j
                self.wFields[i][j]=0

        field=ROOT.TVector3(self.wFields[i][0],self.wFields[i][1],self.wFields[i][2])
                 
        return field

    def getVel(self,i):
        for j in range(3):
            if self.velocities[i][j]!=self.velocities[i][j]:
                self.velocities[i][j]=0

        vel=ROOT.TVector3(self.velocities[i][0],self.velocities[i][1],self.velocities[i][2])

        return vel
    def getForce(self,i):
        force=ROOT.TVector3(self.force[i][0],self.force[i][1],self.force[i][2])
        return force
    def getEForce(self,i):
        force=ROOT.TVector3(self.eforce[i][0],self.eforce[i][1],self.eforce[i][2])
        return force
    def getPos(self,i):
        pos=ROOT.TVector3(self.positions[i][0],self.positions[i][1],self.positions[i][2])
        return pos

class electronCollection():
    def __init__(self,sx,sy):
        self.electrons=[]
        self.ramoHist=ROOT.TH1D()
        self.current={}
        self.maxt=0
        self.startX=sx
        self.startY=sy
        
    def addElec(self,electron):
        self.electrons.append(electron)

    def readFromFile(self,fname,n=-1,force=False):
        driftpos=-1
        old_driftpos=-1
        counter=-1

        for line in open(fname,'read'):
            if line[0]=='#' or line[0]=='\n':
                continue

            words=line.split()
            if words[0]=='Particle:':
                self.addElec(electron(int(words[1])))
                #print words[1]
                counter+=1
                continue
            if counter>=n and n!=-1:
                break


            driftpos=float(words[1])
            if old_driftpos!=-1 and driftpos==old_driftpos:
                old_driftpos=driftpos
                continue
            for iw in range(len(words)):
                if words[iw]=='NaN':
                    words[iw]='0'

            self.electrons[-1].addState(float(words[0]),float(words[1]),float(words[2]),
                                        float(words[6]),float(words[7]),float(words[8]),
                                        float(words[3]),float(words[4]),float(words[5]),
                                        float(words[9]))
            if force:
                self.electrons[-1].addForce(float(words[10]),float(words[11]),float(words[12]))
                self.electrons[-1].addEForce(float(words[13]),float(words[14]),float(words[15]))

            self.maxt=max(float(words[9]),self.maxt)
                

            ramo=self.electrons[-1].calcRamoLast()
            if ramo!=ramo:
                print 'ALARM NAN'
                
            #print 'ramo=',ramo

            if self.current.get(words[9],None)==None:
                self.current[words[9]]=ramo
            else:
                self.current[words[9]]+=ramo

            old_driftpos=driftpos
            
    def getCurrentInt(self,posOnly=False):
        integral=0
        for t in self.current:
            if posOnly and self.current[t]<0:
                continue
            integral+=self.current[t]
        return integral

    def getStart(self):
        start=ROOT.TVector3(self.startX,self.startY,0)
        return start