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 plotHists(c,h,num,fit=0):
    c.Divide(num,2)
    for i in range(num):
        c.cd(i+1)
        h[i].Draw('colz')
        if fit:
            h[i].FitSlicesY(0,1,-2,0,'QNRG3S')
            c.cd(i+1+num)
            ROOT.gDirectory.Get(h[i].GetName()+'_2').Draw()# @UndefinedVariable

def getCorrMatrix2(m):
    tmp=ROOT.TMatrixD(2,2)
    tmp[0][0]=1./sqrt(m[0][0])
    tmp[1][1]=1./sqrt(m[1][1])
    tmp2=ROOT.TMatrixD(m,ROOT.TMatrixD.kMult,tmp)
    corrm=ROOT.TMatrixD(tmp,ROOT.TMatrixD.kMult,tmp2)
    return corrm
def binPos(pos):
    bw=1
    #print "po before: "
    #pos.Print()
    pos.SetXYZ(int(pos.X()/bw),int(pos.Y()/bw),int(pos.Z()/(bw)))      
    #print 'pos after:'
    #pos.Print()

Dx=2
Dy=2
Dz=4
dobinning=1

digipos_hist=[]
clusterpos_hist=[]
for i in range(3):
    digipos_hist.append(ROOT.TH1D("digipos"+str(i),'digipos '+str(i),100,-10,10)) #@UndefinedVariable
    clusterpos_hist.append(ROOT.TH1D("clusterpos"+str(i),'clusterpos '+str(i),100,-2,2)) #@UndefinedVariable

digiposCorr_hist=[]
clusterposCorr_hist=[]
for i in range(3):
    digiposCorr_hist.append(ROOT.TH1D("digiposCorr"+str(i),'digipos correlated '+str(i),200,-20,20)) #@UndefinedVariable
    clusterposCorr_hist.append(ROOT.TH1D("clusterposCorr"+str(i),'clusterpos correlated '+str(i),100,-2,2)) #@UndefinedVariable

correlationGraph1=ROOT.TGraph() #@UndefinedVariable
correlationGraph2=ROOT.TGraph() #@UndefinedVariable
ncluster=1000
ndigi=10
#correlated, along track:

for thetaDeg in range(44,45):
    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
    
    normal=ROOT.TVector3(1,0,0)  # @UndefinedVariable
    normal.SetTheta(theta)
    normal.SetPhi(90*ROOT.TMath.DegToRad())  # @UndefinedVariable
    normal.SetMag(1)

    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]
    
    corelMatrix1=ROOT.TMatrixD(3,3) #@UndefinedVariable
    corelMatrix1_cl=ROOT.TMatrixD(3,3) #@UndefinedVariable
    corelMatrix1_plane=ROOT.TMatrixD(2,2) #@UndefinedVariable
    corelMatrix1_cl_plane=ROOT.TMatrixD(2,2) #@UndefinedVariable
    
    corelMatrix2=ROOT.TMatrixD(3,3) #@UndefinedVariable
    corelMatrix2_cl=ROOT.TMatrixD(3,3) #@UndefinedVariable
    corelMatrix2_plane=ROOT.TMatrixD(2,2) #@UndefinedVariable
    corelMatrix2_cl_plane=ROOT.TMatrixD(2,2) #@UndefinedVariable
    

    #uncorrelated i.e no track
    for i in range(ncluster):
        digipos=[]
        print 'at:',i,' \r',
        clpos=ROOT.TVector3(0,0,0)#@UndefinedVariable 
        clpos_plane=ROOT.TVector3(0,0,0)  #@UndefinedVariable 
        for j in range(ndigi):
            digipos.append(ROOT.TVector3(random.gauss(0,Dx),random.gauss(0,Dy),random.gauss(0,Dz))) #@UndefinedVariable
            if dobinning:
                binPos(digipos[-1])
            clpos+=digipos[-1]
            m_digi=ROOT.TMatrixD(3,1) #@UndefinedVariable
            m_digiPlane=ROOT.TMatrixD(2,1)#@UndefinedVariable
            for k in range(3):
                digipos_hist[k].Fill(digipos[-1][k])
                m_digi[k][0]=digipos[-1][k] 
            digiposPlane=projOnPlane(planeU,planeV,digipos[-1])
            clpos_plane+=digiposPlane
            m_digiPlane[0][0]=digiposPlane[0]
            m_digiPlane[1][0]=digiposPlane[1]
                
            m_digi_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_digi) #@UndefinedVariable
            m_digiPlane_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_digiPlane) #@UndefinedVariable
            corelMatrix1+=ROOT.TMatrixD(m_digi,ROOT.TMatrixD.kMult,m_digi_T) #@UndefinedVariable
            corelMatrix1_plane+=ROOT.TMatrixD(m_digiPlane,ROOT.TMatrixD.kMult,m_digiPlane_T)
        
        clpos*=1./ndigi
        clpos_plane*=1./ndigi
        
        m_clus=ROOT.TMatrixD(3,1) #@UndefinedVariable
        for k in range(3):
            m_clus[k][0]=clpos[k]
        m_clus_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_clus) #@UndefinedVariable
        corelMatrix1_cl+=ROOT.TMatrixD(m_clus,ROOT.TMatrixD.kMult,m_clus_T) #@UndefinedVariable

        m_clusPlane=ROOT.TMatrixD(2,1) #@UndefinedVariable
        m_clusPlane[0][0]=clpos_plane[0]
        m_clusPlane[1][0]=clpos_plane[1]
        m_clusPlane_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_clusPlane) #@UndefinedVariable
        corelMatrix1_cl_plane+=ROOT.TMatrixD(m_clusPlane,ROOT.TMatrixD.kMult,m_clusPlane_T) 

        corelMatrix1*=1./ndigi
        corelMatrix1_plane*=1./ndigi
        
        for k in range(3):
            clusterpos_hist[k].Fill(clpos[k])
            
        
    
   
        clpos=ROOT.TVector3(0,0,0) # @UndefinedVariable
        clpos_plane=ROOT.TVector3(0,0,0)  #@UndefinedVariable 
        
        for j in range(ndigi):
            trackpos=ROOT.TVector3(normal) #@UndefinedVariable
            trackpos*=ROOT.gRandom.Uniform(-15,15) #@UndefinedVariable
            diffFactor=sqrt( abs(trackpos.Z()+15) )
            #digipos[j]*=diffFactor
            digipos[j]+=trackpos
            
            #diff=ROOT.TVector3(random.gauss(0,Dx*diffFactor,random.gauss(0,Dy*diffFactor),random.gauss(0,Dz*diffFactor)) # @UndefinedVariable    
            clpos+=digipos[j]
            #digipos.Print()
            m_digi=ROOT.TMatrixD(3,1) #@UndefinedVariable
            m_digiPlane=ROOT.TMatrixD(2,1)#@UndefinedVariable
            for k in range(3):
                digiposCorr_hist[k].Fill(digipos[j][k])
                m_digi[k][0]=digipos[j][k] 
            m_digi_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_digi) #@UndefinedVariable
            corelMatrix2+=ROOT.TMatrixD(m_digi,ROOT.TMatrixD.kMult,m_digi_T) #@UndefinedVariable
            
            digiposPlane=projOnPlane(planeU,planeV,digipos[j])
            clpos_plane+=digiposPlane
            m_digiPlane[0][0]=digiposPlane[0]
            m_digiPlane[1][0]=digiposPlane[1]
            m_digiPlane_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_digiPlane) #@UndefinedVariable
            corelMatrix2_plane+=ROOT.TMatrixD(m_digiPlane,ROOT.TMatrixD.kMult,m_digiPlane_T)

        clpos*=1./ndigi
        clpos_plane*=1./ndigi
        
        m_clus=ROOT.TMatrixD(3,1)  #@UndefinedVariable
        for k in range(3):
            m_clus[k][0]=clpos[k]
        m_clus_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_clus) #@UndefinedVariable
        corelMatrix2_cl+=ROOT.TMatrixD(m_clus,ROOT.TMatrixD.kMult,m_clus_T) #@UndefinedVariable
        
        corelMatrix2*=1./ndigi
        corelMatrix2_plane*=1./ndigi
        
        for k in range(3):
            clusterposCorr_hist[k].Fill(clpos[k])
        
        m_clusPlane=ROOT.TMatrixD(2,1) #@UndefinedVariable
        m_clusPlane[0][0]=clpos_plane[0]
        m_clusPlane[1][0]=clpos_plane[1]
        m_clusPlane_T=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,m_clusPlane) #@UndefinedVariable
        corelMatrix2_cl_plane+=ROOT.TMatrixD(m_clusPlane,ROOT.TMatrixD.kMult,m_clusPlane_T) 

    print 
    corelMatrix1_cl*=1./ncluster
    corelMatrix1_cl_plane*=1./ncluster
    corrMatrix1_cl_plane=getCorrMatrix2(corelMatrix1_cl_plane)
    correlationGraph1.SetPoint(correlationGraph1.GetN(),thetaDeg,corrMatrix1_cl_plane[0][1])
    corelMatrix2_cl*=1./ncluster
    corelMatrix2_cl_plane*=1./ncluster
    corrMatrix2_cl_plane=getCorrMatrix2(corelMatrix2_cl_plane)
    correlationGraph2.SetPoint(correlationGraph2.GetN(),thetaDeg,corrMatrix2_cl_plane[0][1])
        
dpos_canv=ROOT.TCanvas("dposC","digipos") #@UndefinedVariable
cpos_canv=ROOT.TCanvas("cposC","cluster pos") #@UndefinedVariable
plotHists(dpos_canv,digipos_hist,3)
plotHists(cpos_canv,clusterpos_hist, 3)

print "digi correlation matrix:"
corelMatrix1.Print()
corelMatrix1_plane.Print()
print 'cluster correlatiopn matric:'
corelMatrix1_cl.Print()
corelMatrix1_cl_plane.Print()
eigenprob1=ROOT.TMatrixDEigen(corelMatrix1_cl_plane) # @UndefinedVariable
eigenval1=eigenprob1.GetEigenValues()
eigenvect1=eigenprob1.GetEigenVectors()
eigenval1.Print()

for i in range(3):
    print "rms Ratio {0}: {1}".format(i,digipos_hist[i].GetRMS()/clusterpos_hist[i].GetRMS())
    
            
dposC_canv=ROOT.TCanvas("dposCC","digipos correlated") #@UndefinedVariable
cposC_canv=ROOT.TCanvas("cposCC","cluster pos correlated") #@UndefinedVariable
plotHists(dposC_canv,digiposCorr_hist,3)
plotHists(cposC_canv,clusterposCorr_hist, 3)

print "digi correlation matrix:"  
corelMatrix2.Print()
corelMatrix2_plane.Print()
print 'cluster correlatiopn matric:'
corelMatrix2_cl.Print()
corelMatrix2_cl_plane.Print()
eigenprob2=ROOT.TMatrixDEigen(corelMatrix2_cl_plane) # @UndefinedVariable
eigenval2=eigenprob2.GetEigenValues()
eigenvect2=eigenprob2.GetEigenVectors()
eigenval2.Print()


print 'correlation matrices:'
corrMatrix1_cl_plane.Print()
corrMatrix2_cl_plane.Print()

for i in range(3):
    print "rms Ratio Corr {0}: {1}".format(i,digiposCorr_hist[i].GetRMS()/clusterposCorr_hist[i].GetRMS())    
        
   
cCorr=ROOT.TCanvas('cCrorr','correlation.')
correlationGraph1.Draw('APL')
correlationGraph2.SetMarkerColor(2)
correlationGraph2.SetLineColor(2)
correlationGraph2.Draw('PL')        
        
        
thisIsTheEnd()