import ROOT,argparse,random,math

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

points=[]
mcpoints=[]
'''
for i in range(5):
    x=random.randrange(-2,2)
    y=random.randrange(-2,2)
    points.append(ROOT.TVector2(x,y))
    mcpoints.append(ROOT.TVector2(x+random.gauss(0,1),y+random.gauss(0,1)))
'''

points.append(ROOT.TVector2(0.15,0))
points.append(ROOT.TVector2(0.075,-0.1299))
points.append(ROOT.TVector2(-0.075,-0.1299))
points.append(ROOT.TVector2(-0.15,0))
points.append(ROOT.TVector2(-0.075,0.1299))
points.append(ROOT.TVector2(0.075,0.1299))
points.append(ROOT.TVector2(0.15 ,0))

points.append(ROOT.TVector2(0.15,0))
points.append(ROOT.TVector2(0.075,-0.1299))
points.append(ROOT.TVector2(-0.075,-0.1299))
points.append(ROOT.TVector2(-0.15,0))
points.append(ROOT.TVector2(-0.075,0.1299))
points.append(ROOT.TVector2(0.075,0.1299))
points.append(ROOT.TVector2(0.15 ,0))


'''
points.append(ROOT.TVector2(0.15,0))
points.append(ROOT.TVector2(-0.15,0))
points.append(ROOT.TVector2(0,0.1289))
points.append(ROOT.TVector2(0,-0.1289))
'''

'''
points.append(ROOT.TVector2(1,1))
#points.append(ROOT.TVector2(1,1))
points.append(ROOT.TVector2(1,-1))
points.append(ROOT.TVector2(-1,1))
points.append(ROOT.TVector2(-1,-1))
'''
mcpoints=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(random.randrange(-5,5),random.randrange(-5,5))
mcp=ROOT.TVector2(mcx,mcy)

#pmean*=1./len(points)

cov=ROOT.TMatrixD(2,2)
for p in points:
    c=ROOT.TMatrixD(2,1)
    c[0][0]=(p.X()-pmean.X())
    c[1][0]=(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]
    p.Print()
    acov.Print()

cov*=1./len(points)
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])
teigenvector.Print()
eigenvector_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,eigenvector)
print "cov eigenwert:"
eigenwert.Print()
print 'cov eigenvect:'
eigenvector.Print()
print 'cov:'
cov.Print()

#claculate ellipse
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]**2)
                S1[i][j]=eigenwert[i][j]
            else:
                S1[i][j]=0
        else:
            S1[i][j]=0
#print 'S matrix:'
#S.Print()
temp=ROOT.TMatrixD(S1,ROOT.TMatrixD.kMult,eigenvector_i)
ellipse=ROOT.TMatrixD(eigenvector,ROOT.TMatrixD.kMult,temp)

#ellipse=cov
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])
print "ellipse eigenwert:"
ellipse_eigenwert.Print()
print 'ellipse eigenvect:'
ellipse_eigenvect.Print()
print 'ellipse:'
ellipse.Print()

#claculate intersection
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=1./math.sqrt(num[0][0])
num=num[0][0]
if num==0:
    num=1
mp_diff_num=ROOT.TVector2(mp_diff.X()*1./math.sqrt(num),mp_diff.Y()*1./math.sqrt(num))

print 'num=',num,1./num
#end claculate intersection

#calculate residual on main axis
if (eigenwert[1][1]!=0):
    proj=[teigenvector*mp_diff/eigenwert[0][0],teigenvector2*mp_diff/eigenwert[1][1]]
else:
    proj=0

#proj=[teigenvector*mp_diff,1]
print 'the projection:',proj

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]=1/math.sqrt(ellipse_eigenwert[i][j]) #same
#                S2[i][j]=(eigenwert[1-i][1-j])*abs(proj[1-i]) #scale axes
#                S2[i][j]=(eigenwert[1-i][1-j]*math.sqrt(num))#scale ellipse
            else:
                S2[i][j]=0
        else:
            S2[i][j]=0
#print 'S matrix:'
#S.Print()
temp=ROOT.TMatrixD(S2,ROOT.TMatrixD.kMult,ellipse_eigenvect_i)
temp2=ROOT.TMatrixD(ellipse_eigenvect,ROOT.TMatrixD.kMult,temp)
ellipse2=temp2

#ellipse2=ellipse
#ellipse2*=1./30
#ellipse2*=mag

ellipse2_eigenprob=ROOT.TMatrixDEigen(ellipse2)
ellipse2_eigenwert=ellipse2_eigenprob.GetEigenValues()
ellipse2_eigenvect=ellipse2_eigenprob.GetEigenVectors()
tellipse2_eigenvect=ROOT.TVector2(ellipse2_eigenvect[0][0],ellipse2_eigenvect[1][0])
print 'ellipse2 eigenwert:'
ellipse2_eigenwert.Print()
print 'ellipse2. eigenvect:'
ellipse2_eigenvect.Print()
#make new ellipse from old one, including mcp


#end make new ellipse
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)

#pellipse=ROOT.TEllipse(pmean.X(),pmean.Y(),1./math.sqrt(abs(ellipse_eigenwert[0][0])),1./math.sqrt(abs(ellipse_eigenwert[1][1])),0,360,tellipse_eigenvect.Phi()*ROOT.TMath.RadToDeg())
pellipse=ROOT.TEllipse(pmean.X(),pmean.Y(),(ellipse_eigenwert[0][0]),(ellipse_eigenwert[1][1]),0,360,tellipse_eigenvect.Phi()*ROOT.TMath.RadToDeg())
pellipse.SetFillStyle(4000)

#pmcellipse=ROOT.TEllipse(pmean.X(),pmean.Y(),1./math.sqrt(abs(ellipse2_eigenwert[0][0])),1./math.sqrt(abs(ellipse2_eigenwert[1][1])),0,360,tellipse2_eigenvect.Phi()*ROOT.TMath.RadToDeg())
pmcellipse=ROOT.TEllipse(pmean.X(),pmean.Y(),ellipse2_eigenwert[0][0],ellipse2_eigenwert[1][1],0,360,tellipse2_eigenvect.Phi()*ROOT.TMath.RadToDeg())

#pellipse.SetFillColor(1)
pmcellipse.SetFillStyle(4000)
pmcellipse.SetLineColor(3)
#pellipse.SetLineWidth(1)

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

ppoints.SetMarkerStyle(20)
ppoints.SetMarkerSize(1)

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

pellipse_axis1=ROOT.TLine(pmean.X(),pmean.Y(),
                          pmean.X()+(ellipse_eigenvect[0][0]*1./math.sqrt(ellipse_eigenwert[0][0])),
                          pmean.Y()+(ellipse_eigenvect[1][0]*1./math.sqrt(ellipse_eigenwert[0][0])))
pellipse_axis1.SetLineColor(4)

pellipse_axis2=ROOT.TLine(pmean.X(),pmean.Y(),
                          pmean.X()+ellipse_eigenvect[0][1]*1./math.sqrt(ellipse_eigenwert[1][1]),
                          pmean.Y()+ellipse_eigenvect[1][1]*1./math.sqrt(ellipse_eigenwert[1][1]))
pellipse_axis2.SetLineColor(4)


pproj1=ROOT.TLine(pmean.X(),pmean.Y(),
                  pmean.X()+eigenvector[0][0]*eigenwert[0][0]*proj[0],
                  pmean.Y()+eigenvector[1][0]*eigenwert[0][0]*proj[0])
pproj1.SetLineColor(6)
pproj2=ROOT.TLine(pmean.X(),pmean.Y(),
                  pmean.X()+eigenvector[0][1]*eigenwert[1][1]*proj[1],
                  pmean.Y()+eigenvector[1][1]*eigenwert[1][1]*proj[1])
pproj2.SetLineColor(6)


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)

hsize=6
hist=ROOT.TH2D('hist','hist',100,-hsize,hsize,100,-hsize,hsize)
c1=ROOT.TCanvas('matrix','matrix',1000,1000)
c1.cd()
hist.Draw()
ppoints.Draw('same')
ppmean.Draw('same')
pellipse.Draw('same')
pcov.Draw("same")
pmcellipse.Draw('same')
pmcp.Draw('same')
pmcnum.Draw('same')
pellipse_axis1.Draw('same')
pellipse_axis2.Draw('same')
pproj1.Draw('same')
pproj2.Draw('same')

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