import os
import pickle
import numpy as np
import sys
import ROOT as r

mother_ID=['Ds','Dplus']
meson_ID =['pi','X']
l_flv = ['e','mu']
data_type = ['MC','data']
mag_status =['Up','Down'] 

masses = {
			'Ds': 1.968,
			'Dplus':1.869,
			'pi':0.140,
			'mu_plus': 0.105,
			'mu_minus': 0.105,
			'e_plus': 0.0005,
			'e_minus': 0.0005,
}

mother_index=None
l_index=None


def get_4_momenta(data, particle):
    
    mass = masses[particle]
       
    p_x = np.array(data[particle+"_PX"])/1000
    p_y = np.array(data[particle+"_PY"])/1000
    p_z = np.array(data[particle+"_PZ"])/1000        
        
    p_spatial = np.array([ p_x, p_y, p_z ])
    
    mod_p_spatial = np.sqrt(np.power(p_spatial,2).sum(axis=0))
    E = np.sqrt(np.power(mod_p_spatial,2)+np.power(mass,2)).reshape(1,p_spatial.shape[1])
    
    p = np.concatenate((E, p_spatial),axis=0)
    
    return p, p_x, p_y, p_z


def get_TLorentzVector(p, i):
    
    p_r = r.TLorentzVector(p[:,i][1], p[:,i][2], p[:,i][3], p[:,i][0])
    return p_r


def get_sp(p1, p2, i=None):
    
    if i:
        spatial_sp2 = p1[:,i][1]*p2[:,i][1] + p1[:,i][2]*p2[:,i][2] + p1[:,i][3]*p2[:,i][3]
        E2 = p1[:,i][0]*p2[:,i][0]
        full_sp = E2 - spatial_sp2
        
    else:
        spatial_sp2 = p1[:,:][1]*p2[:,:][1] + p1[:,:][2]*p2[:,:][2] + p1[:,:][3]*p2[:,:][3]
        E2 = p1[:,:][0]*p2[:,:][0]
        full_sp = E2 - spatial_sp2
        
    return full_sp, E2, spatial_sp2

def get_costheta(p_l, p_pi):

	num = p_pi.Px()*p_l.Px()+p_pi.Py()*p_l.Py()+p_pi.Pz()*p_l.Pz()

	mod_p_pi = np.sqrt(p_pi.Px()*p_pi.Px()+p_pi.Py()*p_pi.Py()+p_pi.Pz()*p_pi.Pz())
	mod_p_l = np.sqrt(p_l.Px()*p_l.Px()+p_l.Py()*p_l.Py()+p_l.Pz()*p_l.Pz())

	costheta = num/(mod_p_pi*mod_p_l)
	return costheta

def get_boost_D_to_lab(data, mother_index=mother_index):
    
    p, p_x, p_y, p_z =  get_4_momenta(data, mother_ID[mother_index])
    
    p_spatial= np.array([ p_x, p_y, p_z ])
    
    mod_p_spatial = np.sqrt(np.power(p_spatial,2).sum(axis=0))
    
    E = p[0,:]
    
    beta_X=(p_x/E).reshape(p.shape[1])
    beta_Y=(p_y/E).reshape(p.shape[1])
    beta_Z=(p_z/E).reshape(p.shape[1])
    
    beta = (mod_p_spatial/E).reshape(p.shape[1])
    
    beta_vec= np.array([beta_X,beta_Y,beta_Z])

    return beta, beta_vec

def get_boost_D_to_q(data, l_index=l_index, mother_index=mother_index):
    
    l_m = masses[l_flv[l_index]+'_plus']
    pi_m = masses['pi']
    D_m = masses[mother_ID[mother_index]]
    
    p_plus , p_plus_x, p_plus_y, p_plus_z = get_4_momenta(data, l_flv[l_index]+'_plus')
    E_plus = p_plus[0,:]
    
    p_plus_spatial= np.array([ p_plus_x, p_plus_y, p_plus_z ])
    mod_p_plus_spatial = np.sqrt(np.power(p_plus_spatial,2).sum(axis=0))

    p_minus , p_minus_x, p_minus_y, p_minus_z = get_4_momenta(data, l_flv[l_index]+'_minus')
    E_minus = p_minus[0,:]
    
    p_minus_spatial= np.array([ p_minus_x, p_minus_y, p_minus_z ])
    mod_p_minus_spatial = np.sqrt(np.power(p_minus_spatial,2).sum(axis=0))
    
    q = p_plus + p_minus
    
    q2, q_02, modq2 = get_sp(q, q)

    lambdaD=((D_m*D_m) - (pi_m*pi_m))**2 -2*q2*(D_m*D_m +pi_m*pi_m) +np.power(q2,2)
    lambdaD_pos = np.where(lambdaD<0,0,lambdaD)
    sqrt_lambda=np.sqrt(lambdaD_pos)

    beta = -sqrt_lambda/(np.sqrt(q2)*2*D_m)

    return beta

def get_costheta_list(data, l_index, mother_index):
    
    costheta_list=[]
    
    p_D_lab, _, _, _, = get_4_momenta(data, mother_ID[mother_index])
    p_pi_lab, _, _, _,= get_4_momenta(data, 'pi')
    p_l_plus_lab, _, _, _, = get_4_momenta(data, l_flv[l_index]+'_plus')
    p_l_minus_lab, _, _, _, = get_4_momenta(data, l_flv[l_index]+'_minus')
    
    _, beta_D_vec = get_boost_D_to_lab(data, mother_index=mother_index)
    beta_q = get_boost_D_to_q(data, l_index=l_index, mother_index=l_index)
    
    for j in range(len(data[mother_ID[mother_index]+"_ConsD_M"])):
    #for j in range(2000):
        p_D_lab_r = get_TLorentzVector(p_D_lab, j)
        p_pi_lab_r = get_TLorentzVector(p_pi_lab, j)
        p_l_plus_lab_r = get_TLorentzVector(p_l_plus_lab, j)
        p_l_minus_lab_r = get_TLorentzVector(p_l_minus_lab, j)
        
        beta_D=r.TVector3(beta_D_vec[:,j][0], beta_D_vec[:,j][1], beta_D_vec[:,j][2])
        lambda_lab_to_D=r.TLorentzRotation(-beta_D)
        
        #boost it from lab to D
        p_D_D = lambda_lab_to_D.VectorMultiplication(p_D_lab_r)
        p_pi_D = lambda_lab_to_D.VectorMultiplication(p_pi_lab_r)
        p_l_plus_D = lambda_lab_to_D.VectorMultiplication(p_l_plus_lab_r)
        p_l_minus_D = lambda_lab_to_D.VectorMultiplication(p_l_minus_lab_r)
        
        beta_q_r=r.TVector3(0,0, beta_q[j])
        lambda_D_to_q=r.TLorentzRotation(beta_q_r)
        
        #boost it from D to q
        p_D_q = lambda_D_to_q.VectorMultiplication(p_D_D)
        p_pi_q = lambda_D_to_q.VectorMultiplication(p_pi_D)
        p_l_plus_q = lambda_D_to_q.VectorMultiplication(p_l_plus_D)
        p_l_minus_q = lambda_D_to_q.VectorMultiplication(p_l_minus_D)
        
        costheta=get_costheta(p_pi_q, p_l_plus_q)
        
        if not np.isnan(costheta):
            costheta_list.append(costheta)
        elif np.isnan(costheta):
            costheta_list.append(-2)
    costheta_list=np.array(costheta_list)
    
    return costheta_list