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
    #print 'acov:',c[0][0],c[1][0]
    #acov.Print()

'''
for i in range(2):
    for j in range(2):
        if i==j:
            cov[i][j]=stdErr[i]
        else:
            cov[i][j]=0
'''
cov*=1./len(points)

#print 'cov='
#cov.Print()

#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)


#claculate ellipse with axis length = 1/eigenwert^2 (the axis length of the ellipsoid corresponing to the cov is 1/sqrt(eigenvalue))
S1=ROOT.TMatrixD(2,2)
for i in range(0,2):
    for j in range(0,2):
        if i==j:
            if eigenwert[i][j]!=0:
                S1[i][j]=1./(eigenwert[i][j])
            else:
                S1[i][j]=0
        else:
            S1[i][j]=0
temp=ROOT.TMatrixD(S1,ROOT.TMatrixD.kMult,eigenvector_i)
ellipse=ROOT.TMatrixD(eigenvector,ROOT.TMatrixD.kMult,temp)
ellipse_eigenprob=ROOT.TMatrixDEigen(ellipse)
ellipse_eigenwert=ellipse_eigenprob.GetEigenValues()
ellipse_eigenvect=ellipse_eigenprob.GetEigenVectors()
ellipse_eigenvect_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,ellipse_eigenvect)
tellipse_eigenvect=ROOT.TVector2(ellipse_eigenvect[0][0],ellipse_eigenvect[1][0])

#claculate intersection with ellipse by using vT M v = 1/l^2
c=ROOT.TMatrixD(2,1)
mp_diff=mcp-pmean
mag=math.sqrt(mp_diff.X()**2+mp_diff.Y()**2)
c[0][0]=mp_diff.X()
c[1][0]=mp_diff.Y()
c_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,c)
mult1=ROOT.TMatrixD(ellipse,ROOT.TMatrixD.kMult,c)
num=ROOT.TMatrixD(c_t,ROOT.TMatrixD.kMult,mult1)
num=(num[0][0])
#calculate scaled vector
#mp_diff_num=ROOT.TVector2(mp_diff.X()*1./math.sqrt(num),mp_diff.Y()*1./math.sqrt(num))
#print 'num=',num, 1./math.sqrt(num),mag/math.sqrt(num)#, math.sqrt(mp_diff_num.X()**2+mp_diff_num.Y()**2), math.sqrt(mp_diff.X()**2+mp_diff.Y()**2)*1./math.sqrt(num)



#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())
#mp_diff_num.SetMagPhi(mag*math.sqrt(sigma[0][0]),mp_diff.Phi())
#mp_diff_num.SetMagPhi(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())

#calculate projections of residual
proj1=mp_diff*teigenvector/math.sqrt(eigenwert[0][0])
proj2=mp_diff*teigenvector2/math.sqrt(eigenwert[1][1])

print 'projs: ',proj1,proj2
#generate axis scaled matrix:
#claculate ellipse with axis length = 1/eigenwert^2 (the axis length of the ellipsoid corresponing to the cov is 1/sqrt(eigenvalue))
S2=ROOT.TMatrixD(2,2)
for i in range(0,2):
    for j in range(0,2):
        if i==j:
            if eigenwert[i][j]!=0:
                S2[i][j]=(eigenwert[i][j])
            else:
                S1[i][j]=0
        else:
            S1[i][j]=0
temp=ROOT.TMatrixD(S1,ROOT.TMatrixD.kMult,eigenvector_i)
ellipse=ROOT.TMatrixD(eigenvector,ROOT.TMatrixD.kMult,temp)
ellipse_eigenprob=ROOT.TMatrixDEigen(ellipse)
ellipse_eigenwert=ellipse_eigenprob.GetEigenValues()
ellipse_eigenvect=ellipse_eigenprob.GetEigenVectors()
ellipse_eigenvect_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,ellipse_eigenvect)
tellipse_eigenvect=ROOT.TVector2(ellipse_eigenvect[0][0],ellipse_eigenvect[1][0])


#check cut of matrix with own eigenvectors
vec=ROOT.TVectorD(2)
vec[0]=eigenvector[0][0]
vec[1]=eigenvector[1][0]

minerr=cov_i.Similarity(vec)
print 'minerr:',minerr,'(',1./minerr,')','eigenwert',eigenwert[0][0],

#create drawable objects
#pcov=ROOT.TEllipse(pmean.X(),pmean.Y(),1/math.sqrt(eigenwert[0][0]),1/math.sqrt(eigenwert[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector2.Phi())
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=ROOT.TEllipse(pmean.X(),pmean.Y(),eigenwert[0][0],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)


pellipse=ROOT.TEllipse(pmean.X(),pmean.Y(),1/math.sqrt(ellipse_eigenwert[0][0]),1/math.sqrt(ellipse_eigenwert[1][1]),0,360,tellipse_eigenvect.Phi()*ROOT.TMath.RadToDeg())
pellipse.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")
#pellipse.Draw('same')
pmcnum.Draw('same')
ppmean.Draw('same')
pmcp.Draw('same')
ppoints.Draw('same')

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