import ROOT,sys,math
sys.path.append('$PANDAPATH/macro/tpc/FOPI/mberger')
from functions import thisIsTheEnd
from random import gauss

class cluster():
    def __init__(self):
        self.pos=ROOT.TVector3(0,0,0)
        self.dir=ROOT.TVector3(0,0,0)
        self.digis=[]
        self.planeO=ROOT.TVector3(0,0,0)
        #self.planeU=ROOT.TVector3()
        self.planeV=ROOT.TVector3(0,0,0)
        self.graph=ROOT.TGraph()
        self.nDigis=0
        self.digisProj=[]
        self.poca=ROOT.TVector3(0,0,0)
        self.Residual=ROOT.TVector3(0,0,0)
        self.ResidualV=-100000

    def setPos(self,pos):
        self.pos=pos
    def setDir(self,dir):
        self.dir=dir
    def setPlane(self,origin,norm):
        self.planeO=origin
        self.planeV=norm.Orthogonal()
        #self.planeV=planU.Cross(self.planeU)

        
    def addDigi(self,digi):
        self.digis.append(digi)
        self.digisProj.append(-10000)
        self.graph.SetPoint(self.graph.GetN(),digi.X(),digi.Y())
        self.nDigis+=1
            
    def getGraph(self):
        return self.graph
    def getNdigis(self):
        return self.nDigis
    def getDigi(self,i):
        return self.digis[i]
    def getDigiProj(self,i):
        return self.digisProj[i]

    def makePos(self):
        for i in range(self.nDigis):
            self.pos+=self.digis[i]
        self.pos*=1./self.nDigis
        self.planeO=ROOT.TVector3(self.pos)
        self.calcPoca()
        self.Residual=ROOT.TVector3(self.pos)
        self.Residual-=self.poca
        self.ResidualV=self.Residual*self.planeV
        self.ResidualV/=self.planeV.Mag()
        
    def projDigis(self):
        for i in range(self.nDigis):
            tmp=ROOT.TVector3(self.digis[i])
            tmp-=self.planeO
            Vpos=tmp*self.planeV
            Vpos/=self.planeV.Mag()
            self.digisProj[i]=Vpos
    def calcPoca(self):
        ptPos=ROOT.TVector3(0,0,0)
        tmp1=ROOT.TVector3(self.pos)
        tmp1-=ptPos
        tmp1=tmp1*self.dir
        tmp2=ROOT.TVector3(self.dir)
        tmp2*=tmp1
        poca=ROOT.TVector3(ptPos)
        poca+=tmp2
        self.poca=ROOT.TVector3(poca)

clusters=[]
#generate dataSample
#loop over theta
t=45
for i in range(t,t+1):
    theta=i*ROOT.TMath.DegToRad()
    print theta
    trackdir=ROOT.TVector3(ROOT.TMath.Cos(theta),ROOT.TMath.Sin(theta),0)
    for k in range(1000):
        print "at cluster",k
        clusters.append(cluster())
        clusters[-1].setDir(trackdir)
        clusters[-1].setPlane(ROOT.TVector3(0,0,0),trackdir)
        #number of digis        self.nDigis+=1
        for j in range(1000):
            dpos=ROOT.TVector3(trackdir)
            dpos*=ROOT.gRandom.Uniform(0,5)
            #print 'dpos'
            #dpos.Print()
            diff=ROOT.TVector3(gauss(0,2),gauss(0,10),0)
            #print 'diff'
            #diff.Print()
            dpos+=diff
            clusters[-1].addDigi(dpos)
        clusters[-1].makePos()
        clusters[-1].projDigis()


hproj=ROOT.TH1D('hproj','proj',100,0,0)
#for i in range(clusters[-1].getNdigis()):
#   hproj.Fill(clusters[-1].getDigiProj(i))
for i in range(len(clusters)):
    print "cluster",i
    #clusters[i].poca.Print()
    #clusters[i].pos.Print()
    #clusters[i].Residual.Print()
    hproj.Fill(clusters[i].ResidualV)
c1=ROOT.TCanvas()
hproj.Draw()

diffCov=ROOT.TMatrixD(2,2)
diffCov[0][0]=2./math.sqrt(500.)
diffCov[1][1]=10./math.sqrt(500.)
plane=ROOT.TVectorD(2)
plane[0]=clusters[-1].planeV.X()
plane[1]=clusters[-1].planeV.Y()
projected=diffCov.Similarity(plane)
print "projected:",projected
projGaus=ROOT.TF1("projgaus","gaus",-10,10)
projGaus.SetParameter(0,2)
projGaus.SetParameter(1,0)
projGaus.SetParameter(2,projected)
projGaus.FixParameter(2,projected)
projGaus.SetParLimits(2,projected,projected)
hproj.Fit(projGaus,"B+")
diffCov.Invert()
cutted=1./diffCov.Similarity(plane)
print "cutted:",cutted
cutGaus=ROOT.TF1("projgaus","gaus",-10,10)
cutGaus.SetParameter(0,2)
cutGaus.SetParameter(1,0)
cutGaus.SetParameter(2,cutted)
cutGaus.FixParameter(2,cutted)
cutGaus.SetParLimits(2,cutted,cutted)
cutGaus.SetLineColor(3)
hproj.Fit(cutGaus,"B+")

print "chi^2: projected:",projGaus.GetChisquare()," cutted:",cutGaus.GetChisquare()
c2=ROOT.TCanvas()
clusters[-1].getGraph().Draw("AP")

thisIsTheEnd()