import sys, os, copy, array
from math import sin, cos, atan, sqrt
pandapath = os.environ.get('PANDAPATH')
sys.path.append(pandapath + '/macro/tpc/FOPI/python/argparse-1.2.1')
sys.path.append(pandapath + '/macro/tpc/FOPI/mberger')
import argparse, ROOT
from functions import set_palette, openTree, thisIsTheEnd, getSlices
import random

ROOT.gROOT.ProcessLine(".x rootlogon.C")  # @UndefinedVariable
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")') # @UndefinedVariable
ROOT.gROOT.LoadMacro("stlPYROOT.h+")  # @UndefinedVariable
set_palette()


def calcPOCA(thisrefpos,mom,pos):
    tmp1=ROOT.TVector3(mom) # @UndefinedVariable
    tmp1.SetMag(1)
    tmp2=ROOT.TVector3(pos-thisrefpos) # @UndefinedVariable
    fac=tmp2*mom
    fac/=mom.Mag()
    tmp1*= fac
    poca=ROOT.TVector3(thisrefpos+tmp1) # @UndefinedVariable
    return poca

def projOnVect(vect,vectToBeProj):
    tmp=ROOT.TVector3(vectToBeProj) # @UndefinedVariable
    if vect.Mag()!=0:
        tmp=tmp*vect
        tmp*=1./vect.Mag()
        tmp2=ROOT.TVector3(vect) # @UndefinedVariable
        tmp2*=tmp
        tmp2*=1./vect.Mag()
    else:
        tmp2=ROOT.TVector3(999,999,999) # @UndefinedVariable
    return tmp2
def projOnPlane(plane1,plane2,vectToBeProj): 
    tmp1=ROOT.TVector3(vectToBeProj) # @UndefinedVariable
    if plane1.Mag()!=0:
        tmp1=tmp1*plane1
        tmp1*=1./plane1.Mag()
    tmp2=ROOT.TVector3(vectToBeProj) # @UndefinedVariable
    if plane2.Mag()!=0:
        tmp2=tmp2*plane2
        tmp2*=1./plane2.Mag()
        
    return ROOT.TVector3(tmp1,tmp2,0) # @UndefinedVariable

def plotPlaneHist(c,h):
    c.Divide(2,2)
    for i in range(2):
        c.cd(i+1)
        h[i].Draw('colz')
        h[i].FitSlicesY(0,1,-2,0,'QNRG3S')
        c.cd(i+3)
        ROOT.gDirectory.Get(h[i].GetName()+'_2').Draw()# @UndefinedVariable

Dx=2
Dy=2
Dz=4

clposHist=[]
residualHist=[]
pullHist=[]
for i in range(3):
    clposHist.append(ROOT.TH1D('clpos'+str(i),'clpos '+str(i),100,-10,10))  # @UndefinedVariable
    residualHist.append(ROOT.TH2D("residual"+str(i),'residual '+str(i),90,0,90,100,-10,10))  # @UndefinedVariable
    pullHist.append(ROOT.TH2D("pull"+str(i),"pull "+str(i),90,0,90,100,-10,10)) # @UndefinedVariable

clposUVHist=[]
residualUVHist=[]
pullUVHist=[]
errorUVHist=[]
for i in range(2):
    clposUVHist.append(ROOT.TH2D('clposUV'+str(i),'clposUV '+str(i),90,0,90,100,-10,10)) # @UndefinedVariable
    residualUVHist.append(ROOT.TH2D("residualUV"+str(i),'residual UV '+str(i),90,0,90,100,-10,10))  # @UndefinedVariable
    pullUVHist.append(ROOT.TH2D("pullUV"+str(i),"pull UV "+str(i),90,0,90,100,-10,10)) # @UndefinedVariable
    errorUVHist.append(ROOT.TH2D("errorUV"+str(i),"error UV "+str(i),90,0,90,1000,0,10)) # @UndefinedVariable

for thetaDeg in range(1,89):
    print "at theta:",thetaDeg
    theta=thetaDeg*ROOT.TMath.DegToRad()    # @UndefinedVariable

    thecov=ROOT.TMatrixDSym(3)  # @UndefinedVariable
    thecov[0][0]=Dx*Dx
    thecov[1][1]=Dy*Dy
    thecov[2][2]=Dz*Dz*2
    
    normal=ROOT.TVector3(1,0,0)  # @UndefinedVariable
    normal.SetTheta(theta)
    normal.SetPhi(90*ROOT.TMath.DegToRad())  # @UndefinedVariable

    planeU=normal.Orthogonal()
    planeV=normal.Cross(planeU)
    planeM=ROOT.TMatrixD(3,2)  # @UndefinedVariable
    planeU.SetMag(1)
    planeV.SetMag(1)
    planeM[0][0]=planeU[0]
    planeM[1][0]=planeU[1]
    planeM[2][0]=planeU[2]
    planeM[0][1]=planeV[0]
    planeM[1][1]=planeV[1]
    planeM[2][1]=planeV[2]
    
    for iclus in range(500):
        mcpos=ROOT.TVector3(0,0,0) # @UndefinedVariable
        clpos=ROOT.TVector3(mcpos) # @UndefinedVariable
        diff=ROOT.TVector3(random.gauss(0,Dx),random.gauss(0,Dy),random.gauss(0,Dz)) # @UndefinedVariable
        clpos+=diff
        for i in range(3):
            clposHist[i].Fill(clpos[i])
        residual=ROOT.TVector3() # @UndefinedVariable
        poca=calcPOCA(ROOT.TVector3(0,0,0),normal,clpos) # @UndefinedVariable
        residual=clpos-poca

        for i in range(3):
            residualHist[i].Fill(thetaDeg,residual[i])
            pullHist[i].Fill(thetaDeg,residual[i]/thecov[i][i])
        
        planeUD=ROOT.TVectorD(3) # @UndefinedVariable
        for i in range(3):
            planeUD[i]=planeU[i]
    
        projCov=ROOT.TMatrixDSym(3,thecov.GetMatrixArray()) # @UndefinedVariable
        projCov.SimilarityT(planeM)
        
        clposPlane=ROOT.TVector3(projOnPlane(planeU,planeV,clpos)) # @UndefinedVariable
        residualPlane=ROOT.TVector3(projOnPlane(planeU, planeV, residual)) # @UndefinedVariable
        eigenprob=ROOT.TMatrixDSymEigen(projCov) # @UndefinedVariable
        eigenval=eigenprob.GetEigenValues()
        eigenvect=eigenprob.GetEigenVectors()
        teigenvect1=ROOT.TVector3(eigenvect[0][0],eigenvect[1][0],0)# @UndefinedVariable
        teigenvect2=ROOT.TVector3(eigenvect[0][1],eigenvect[1][1],0)    # @UndefinedVariable 
        residualPlaneAxis=ROOT.TVector3(projOnPlane(teigenvect1, teigenvect2, residualPlane))# @UndefinedVariable
        clposPlaneAxis=ROOT.TVector3(projOnPlane(teigenvect1,teigenvect2,clpos)) # @UndefinedVariable
        for i in range(2):
            clposUVHist[i].Fill(thetaDeg,clposPlane[i])
            residualUVHist[i].Fill(thetaDeg,residualPlaneAxis[i])
            pullUVHist[i].Fill(thetaDeg,residualPlaneAxis[i]/sqrt(eigenval[i]))
            errorUVHist[i].Fill(thetaDeg,sqrt(eigenval[i]))
    
clpos_canvas=ROOT.TCanvas() # @UndefinedVariable
clpos_canvas.Divide(2,2)
for i in range(3):
    clpos_canvas.cd(i+1)
    clposHist[i].Draw()

residual_canvas=ROOT.TCanvas('residualC','residuals') # @UndefinedVariable
residual_canvas.Divide(2,2)
for i in range(3):
    residual_canvas.cd(i+1)
    residualHist[i].Draw('colz')
    
pull_canvas=ROOT.TCanvas("pullC","pulls")  # @UndefinedVariable
pull_canvas.Divide(2,2)
for i in range(3):
    pull_canvas.cd(i+1)
    pullHist[i].Draw('colz')
    
clposPlane_canvas=ROOT.TCanvas('cposPlaneC','clpos plane')
plotPlaneHist(clposPlane_canvas,clposUVHist)
residualPlane_canvas=ROOT.TCanvas("residualPlaneC",'residual plane')
plotPlaneHist(residualPlane_canvas,residualUVHist)
pullPlane_canvas=ROOT.TCanvas("pullPlaneC",'pull plane')
plotPlaneHist(pullPlane_canvas, pullUVHist)
errorPlane_canvas=ROOT.TCanvas("errorPlaneC",'error plane')
plotPlaneHist(errorPlane_canvas, errorUVHist)

thisIsTheEnd()