import ROOT,argparse,random,math

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')

#get random points
points=[]
mcpoints=[]

xpos=[1,-1,0,0]
ypos=[0,0,1,-1]
xpos=[-1.06,-2,-1.03,-0.2,-0.4,1.67]
ypos=[-2,-1.33,0.25,1.56,-0.87,1.7]
for i in range(4):
    #x=random.randrange(-2,2)
    #y=random.randrange(-2,2)
    xpos[i]=random.randrange(-2,2)
    ypos[i]=random.randrange(-2,2)
    x=xpos[i]
    y=ypos[i]
    points.append(ROOT.TVector2(x,y))
    mcpoints.append(ROOT.TVector2(x+random.gauss(0,1),y+random.gauss(0,1)))

#calculate mean positions of points and mc points
mx=0
my=0
mcx=0
mcy=0
for p in points:
    mx+=p.X()
    my+=p.Y()
for p in mcpoints:
    mcx+=p.X()
    mcy+=p.Y()


mx/=len(points)
my/=len(points)
mcx/=len(mcpoints)
mcy/=len(mcpoints)
pmean=ROOT.TVector2(mx,my)
mcp=ROOT.TVector2(mcx,mcy)
#mcp=ROOT.TVector2(0.2,0.2)

stdErr=[0,0]
for p in points:
    stdErr[0]+=(p.X()-pmean.X())**2
    stdErr[1]+=(p.Y()-pmean.Y())**2

#calculate covarianz matrix
cov=ROOT.TMatrixD(2,2)
for p in points:
    c=ROOT.TMatrixD(2,1)
    c[0][0]=abs(p.X()-pmean.X())
    c[1][0]=abs(p.Y()-pmean.Y())
    c_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,c)
    acov=ROOT.TMatrixD(c,ROOT.TMatrixD.kMult,c_t)
    cov+=acov

cov*=1./len(points)


#get eigenvectors and eigenvalues of cov
eigenproblem=ROOT.TMatrixDEigen(cov)
eigenwert=eigenproblem.GetEigenValues()
eigenvector=eigenproblem.GetEigenVectors()
teigenvector=ROOT.TVector2(eigenvector[0][0],eigenvector[1][0])
teigenvector2=ROOT.TVector2(eigenvector[0][1],eigenvector[1][1])
eigenvector_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,eigenvector)

mp_diff=mcp-pmean
mag=math.sqrt(mp_diff.X()**2+mp_diff.Y()**2)


#calculate sigma along cluster-MC direction:
cov_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,cov)
residual=ROOT.TMatrixD(2,1)
residual[0][0]=mp_diff.X()/mag
residual[1][0]=mp_diff.Y()/mag
residual_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,residual)
residual_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,residual)
print 'residual='
residual.Print()

temp2=ROOT.TMatrixD(cov_i,ROOT.TMatrixD.kMult,residual)
#temp2=ROOT.TMatrixD(cov,ROOT.TMatrixD.kMult,residual)

sigma=ROOT.TMatrixD(residual_t,ROOT.TMatrixD.kMult,temp2)

print 'sigma=',sigma[0][0],1./math.sqrt(sigma[0][0])
print 'mag/sqrt(sigma)=',mag/math.sqrt(sigma[0][0])

mp_diff_num=ROOT.TVector2(mp_diff)
mp_diff_num.SetMagPhi(1./math.sqrt(sigma[0][0]),mp_diff.Phi())


cov2=ROOT.TMatrixD(cov)

scale=ROOT.TMatrixD(2,2)
scale[0][0]=2
scale[0][1]=0
scale[1][0]=0
scale[1][1]=1
#cov2=ROOT.TMatrixD(scale,ROOT.TMatrixD.kMult,cov2)
cov2*=(mag**2 * sigma[0][0])
print 'cov2='
cov2.Print()
#get eigenvectors and eigenvalues of cov2
eigenproblem_2=ROOT.TMatrixDEigen(cov2)
eigenwert_2=eigenproblem_2.GetEigenValues()
eigenvector_2=eigenproblem_2.GetEigenVectors()
teigenvector_2=ROOT.TVector2(eigenvector_2[0][0],eigenvector_2[1][0])
teigenvector2_2=ROOT.TVector2(eigenvector_2[0][1],eigenvector_2[1][1])
eigenvector_2_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,eigenvector_2)

#calculate sigma along cluster-MC direction for cov2:
cov2_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,cov2)
temp3=ROOT.TMatrixD(cov2_i,ROOT.TMatrixD.kMult,residual)
sigma2=ROOT.TMatrixD(residual_t,ROOT.TMatrixD.kMult,temp3)
print 'mag/sqrt(sigma2)=',mag/math.sqrt(sigma2[0][0]),sigma2[0][0],mag
error=1./math.sqrt(sigma2[0][0])
print 'error=',1./error,1./math.sqrt(sigma2[0][0])

print 'pull=',mag/error
print 'mag:',mag,
print 'error:',error,'1./error',1./error

#mp_diff_num.SetMagPhi(1./math.sqrt(sigma2[0][0]),mp_diff.Phi())

#create drawable objects
pcov=ROOT.TEllipse(pmean.X(),pmean.Y(),math.sqrt(eigenwert[0][0]),math.sqrt(eigenwert[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector.Phi())
pcov.SetLineColor(6)
pcov.SetFillStyle(4000)
pcov_axis1=ROOT.TLine(pmean.X(),pmean.Y(),
                      pmean.X()+teigenvector.X()*math.sqrt(eigenwert[0][0]),
                      pmean.Y()+teigenvector.Y()*math.sqrt(eigenwert[0][0]))
pcov_axis1.SetLineColor(3)
pcov_axis2=ROOT.TLine(pmean.X(),pmean.Y(),
                      pmean.X()+teigenvector2.X()*math.sqrt(eigenwert[1][1]),
                      pmean.Y()+teigenvector2.Y()*math.sqrt(eigenwert[1][1]))
pcov_axis2.SetLineColor(3)



pcov2=ROOT.TEllipse(pmean.X(),pmean.Y(),math.sqrt(eigenwert_2[0][0]),math.sqrt(eigenwert_2[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector_2.Phi())
pcov2.SetLineColor(7)
pcov2.SetFillStyle(4000)


pmcnum=ROOT.TLine(pmean.X(),pmean.Y(),pmean.X()+mp_diff_num.X(),pmean.Y()+mp_diff_num.Y())

ppmean=ROOT.TPolyMarker()
ppmean.SetNextPoint(pmean.X(),pmean.Y())
ppmean.SetMarkerColor(2)
ppmean.SetMarkerStyle(20)
ppmean.SetMarkerSize(1)

pmcp=ROOT.TPolyMarker()
pmcp.SetNextPoint(mcp.X(),mcp.Y())
pmcp.SetMarkerColor(3)
pmcp.SetMarkerStyle(20)
pmcp.SetMarkerSize(1)

ppoints=ROOT.TPolyMarker()
for p in points:
    ppoints.SetNextPoint(p.X(),p.Y())
ppoints.SetMarkerStyle(20)
ppoints.SetMarkerSize(1)

#draw 
hsize=6
hist=ROOT.TH2D('hist','hist',100,-hsize,hsize,100,-hsize,hsize)
c1=ROOT.TCanvas('matrix','matrix',1000,1000)
c1.cd()
hist.Draw()
pcov.Draw("same")
pcov_axis1.Draw("same")
pcov_axis2.Draw("same")
pcov2.Draw("same")
pmcnum.Draw('same')
ppmean.Draw('same')
pmcp.Draw('same')
ppoints.Draw('same')

c1.Update()
u=raw_input('done?')