import sys, os, copy, array
from math import sin, cos, atan
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") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")
set_palette()

def calcPOCA(thisrefpos,mom,pos):
    tmp1=ROOT.TVector3(mom)
    tmp1.SetMag(1)
    tmp2=ROOT.TVector3(pos-thisrefpos)
    fac=tmp2*mom
    fac/=mom.Mag()
    tmp1*= fac
    poca=ROOT.TVector3(thisrefpos+tmp1)
    return poca
def projOnVect(vect,vectToBeProj):
    tmp=ROOT.TVector3(vectToBeProj)
    if vect.Mag()!=0:
        tmp=tmp*vect
        tmp*=1./vect.Mag()
        tmp2=ROOT.TVector3(vect)
        tmp2*=tmp
        tmp2*=1./vect.Mag()
    else:
        tmp2=ROOT.TVector3(999,999,999)
    return tmp2

dir=ROOT.TVector3(0,1,0)

print 'dir phi',dir.Phi()
print 'dir theta',dir.Theta()
DXY=0.02
DZ=0.03

mcpoints=[]
diffpoints=[]
absres=[]
pocares=[]
poca=[]
theta=[]
err=[]
errproj=[]
errScal=[]
errPocaScal=[]
betas=[]
diffrad=[]
for itheta in range(1,91):
    print 'at theta:',itheta
    itheta_rad=itheta*ROOT.TMath.DegToRad()
    dir.SetTheta(itheta_rad)
    dir.SetPhi(90*ROOT.TMath.DegToRad())    
    print 'dir phi',dir.Phi()
    print 'dir theta',dir.Theta()
    for i in range(10):
        pos=ROOT.TVector3(dir)
        pos*=ROOT.gRandom.Uniform(0,15)
        mcpoints.append(pos)
        #diffX=ROOT.gRandom.Gaus(0,DXY)
        #diffY=ROOT.gRandom.Gaus(0,DXY)
        #diffZ=ROOT.gRandom.Gaus(0,DZ)
        diffX=DXY
        diffY=DXY
        diffZ=DZ
        diffpoints.append(ROOT.TVector3(pos.X()+diffX,pos.Y()+diffY,pos.Z()+diffZ))
        absres.append(ROOT.TVector3(diffX,diffY,diffZ))
        poca.append(calcPOCA(ROOT.TVector3(0,0,0),dir,diffpoints[-1]))
        tmp=ROOT.TVector3(poca[-1])
        tmp-=diffpoints[-1]
        pocares.append(ROOT.TVector3(tmp))
        theta.append(itheta)
        tmp=ROOT.TVector3(0,0,DZ)
        err.append(ROOT.TVector3(tmp))
        '''
        if(pocares[-1].Mag()!=0):
            tmp=tmp*pocares[-1]
            tmp*=1./pocares[-1].Mag()
            tmp2=ROOT.TVector3(pocares[-1])
            tmp2*=tmp
            tmp2*=1./pocares[-1].Mag()
        else:
            tmp2=ROOT.TVector3(999,999,999)
            '''
        #errproj.append(ROOT.TVector3(tmp2))
        errproj.append(ROOT.TVector3(projOnVect(pocares[-1],err[-1])))
        errScal.append(ROOT.TVector3(0,DZ*sin(itheta_rad)*cos(itheta_rad),DZ*(sin(itheta_rad))**2))
        errPocaScal.append(ROOT.TVector3(projOnVect(pocares[-1],errScal[-1])))
        betas.append(atan(DXY/DZ))
    
    
rfile=ROOT.TFile("./test_residual.root","recreate")
mytuple=ROOT.TNtuple("points","","mcx:mcy:mcz:diffpX:diffpY:diffpZ:diffX:diffY:diffZ:diffMag:pocaX:pocaY:pocaZ:pocaresX:pocaresY:pocaresZ:pocaresMag:theta:beta:errX:errY:errZ:errpX:errpY:errpZ:errScalX:errScalY:errScalZ:errScalMag:errPocaScalX:errPocaScalY:errPocaScalZ:errPocaScalMag")
for i in range(len(mcpoints)):
    tofill=[mcpoints[i].X(),
            mcpoints[i].Y(),
            mcpoints[i].Z(),
            diffpoints[i].X(),
            diffpoints[i].Y(),
            diffpoints[i].Z(),
            absres[i].X(),
            absres[i].Y(),
            absres[i].Z(),
            absres[i].Mag(),
            poca[i].X(),
            poca[i].Y(),
            poca[i].Z(),
            pocares[i].X(),
            pocares[i].Y(),
            pocares[i].Z(),
            pocares[i].Mag(),
            theta[i],
            betas[i],
            err[i].X(),
            err[i].Y(),
            err[i].Z(),
            errproj[i].X(),
            errproj[i].Y(),
            errproj[i].Z(),
            errScal[i].X(),
            errScal[i].Y(),
            errScal[i].Z(),
            errScal[i].Mag(),
            errPocaScal[i].Z(),
            errPocaScal[i].Y(),
            errPocaScal[i].Z(),
            errPocaScal[i].Mag()
           ]
    mytuple.Fill(array.array("f",tofill))
       

mytuple.Write()             
rfile.Close()