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

mother_ID=['Dplus','Ds']
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):
       
    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_E = np.array(data[particle+"_PE"])/1000      
    p_spatial = np.array([ p_x, p_y, p_z ])
    
    p_E = p_E.reshape(1,p_x.shape[0])
    
    p = np.concatenate((p_E, p_spatial),axis=0)
    
    return p, p_spatial


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_costheta_list(data, l_index, mother_index):

    cos_theta_list=[]

    p_D_lab, _, = get_4_momenta(data, mother_ID[mother_index])
    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')
    p_pi_lab, _,  = get_4_momenta(data, 'pi')

    for j in range(len(data[mother_ID[mother_index]+"_ConsD_M"])):    
        
        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)
        p_pi_lab_r = get_TLorentzVector(p_pi_lab, j)    
        q_lab_r = p_l_plus_lab_r+p_l_minus_lab_r

        
        #Determine boost from lab to D
        lab_to_D_boost_vector = -1.*p_D_lab_r.BoostVector()
        
        #Boosting everything from lab to B
        p_D_lab_r.Boost(lab_to_D_boost_vector)
        p_pi_lab_r.Boost(lab_to_D_boost_vector)
        p_l_minus_lab_r.Boost(lab_to_D_boost_vector)
        p_l_plus_lab_r.Boost(lab_to_D_boost_vector)
        
        q_lab_r.Boost(lab_to_D_boost_vector)

        #Determine boost from B to q
        D_to_q_boost_vector = -1.*q_lab_r.BoostVector()
        
        #Boosting everything from D to q
        p_D_lab_r.Boost(D_to_q_boost_vector)
        p_pi_lab_r.Boost(D_to_q_boost_vector)
        
        p_l_minus_lab_r.Boost(D_to_q_boost_vector)
        p_l_plus_lab_r.Boost(D_to_q_boost_vector)
        
        
        vpi=p_pi_lab_r.Vect().Unit()
        vlplus=p_l_plus_lab_r.Vect().Unit()

        cos_theta_pi = vpi.Dot(vlplus)
        cos_theta_list.append(cos_theta_pi)

    cos_theta_list=np.array(cos_theta_list)
    
    return cos_theta_list