import numpy as np
from math import pi
from utils.kin_utils_edit import *

import zfit
ztf = zfit.ztf

def momenta_map(q2, cos_theta_l, 
                mB, mBbar2, mKst, m_l, mgamma, mode=None):

    if mode =='numpy':
        cos_ = np.cos
        sin_ = np.sin
        sqrt_ = np.sqrt
        square_ = np.square
        dphi2_ = dphi2_np
        zero_= np.zeros_like(q2)
        one_= np.ones_like(q2)
        pi_ = pi*np.ones_like(q2)
        conc_ = np.concatenate
        lorentz_boost_ = lorentz_boost_np
        n_events_ = q2.shape[0]
        lambda_f = lambda_function_np


    if mode =='tf':

        cos_ = tf.cos
        sin_ = tf.sin
        sqrt_ = ztf.sqrt
        square_ = ztf.square
        dphi2_ = dphi2_tf
        zero_ = tf.zeros_like(q2, dtype=tf.float64)
        one_ = tf.ones_like(q2, dtype=tf.float64)
        pi_ = pi*tf.ones_like(q2, dtype=tf.float64)
        conc_ = tf.stack
        lorentz_boost_ = lorentz_boost
        lambda_f = lambda_function
    
    """
    Decay variables

    """

    sin_theta_l = sqrt_(1-square_(cos_theta_l))

    """
    Integrated out variables
    
    """

    # phi_gamma = draw_random_angle(q2,mode)
    
    # cos_phi_gamma = cos_(phi_gamma)
    # sin_phi_gamma = sin_(phi_gamma)
    
    
    
    # cos_theta_gamma = draw_random_costheta(q2, mode)
    # sin_theta_gamma = sqrt_(1-square_(cos_theta_gamma))
    
    """
    
    Mute variables

    """
    
    
    phi_l = draw_random_angle(q2,mode)
    phi_Kst = draw_random_angle(q2,mode)
    
    
    cos_phi_l = cos_(phi_l)
    sin_phi_l = sin_(phi_l)
    
    cos_phi_Kst = cos_(phi_Kst)
    sin_phi_Kst = sin_(phi_Kst)
    
    cos_theta_Kst = draw_random_costheta(q2, mode)
    sin_theta_Kst = sqrt_(1-square_(cos_theta_Kst))
    
    """
    q-RF
    
    """
    beta = dphi2_(sqrt_(q2), m_l, m_l)
    
    p1_1 = 0.5*sqrt_(q2)*beta*sin_theta_l*cos_phi_l
    p1_2 = 0.5*sqrt_(q2)*beta*sin_theta_l*sin_phi_l
    p1_3 = 0.5*sqrt_(q2)*beta*cos_theta_l
    p1_0 = 0.5*sqrt_(q2)
    
    p1 = conc_([               p1_0,
                               p1_1,
                               p1_2,
                               p1_3], axis =1)
    
    p2 = conc_([                p1_0,
                               -p1_1,
                               -p1_2,
                               -p1_3], axis =1)
    
    """
    pBbar-RF
    
    """
    mgamma_mass=mgamma
    lambda_bar = lambda_f(sqrt_(mBbar2),mKst,sqrt_(q2))

    beta_K = sqrt_(lambda_bar)/(sqrt_(mBbar2)**2-q2+mKst**2)
    beta_q = -sqrt_(lambda_bar)/(sqrt_(mBbar2)**2+q2-mKst**2)
    
    # E_gamma = (mB**2-sqrt_(mBbar2)**2-mgamma**2)/(2*sqrt_(mBbar2))
    E_Kst = (sqrt_(mBbar2)**2-q2+mKst**2)/(2*sqrt_(mBbar2))
    E_q = (sqrt_(mBbar2)**2+q2-mKst**2)/(2*sqrt_(mBbar2))
    
    # modk_Bbar = sqrt_(-mgamma_mass**2+E_gamma**2)
    
    pKst_Bbar_0 = E_Kst
    pKst_Bbar_1 = E_Kst*beta_K*(sin_theta_Kst*cos_phi_Kst)
    pKst_Bbar_2 = E_Kst*beta_K*(sin_theta_Kst*sin_phi_Kst)
    pKst_Bbar_3 = E_Kst*beta_K*cos_theta_Kst
    
    q_Bbar_0 =  E_q
    q_Bbar_1 =  E_q*beta_q*(sin_theta_Kst*cos_phi_Kst)
    q_Bbar_2 =  E_q*beta_q*(sin_theta_Kst*sin_phi_Kst)
    q_Bbar_3 =  E_q*beta_q*cos_theta_Kst
    
    # k_Bbar_0 =  E_gamma
    # k_Bbar_1 =  modk_Bbar*sin_theta_gamma*cos_phi_gamma
    # k_Bbar_2 =  modk_Bbar*sin_theta_gamma*sin_phi_gamma
    # k_Bbar_3 =  modk_Bbar*cos_theta_gamma
    
    # pB_Bbar_0 = sqrt_(mBbar2)+k_Bbar_0
    # pB_Bbar_1 = k_Bbar_1
    # pB_Bbar_2 = k_Bbar_2
    # pB_Bbar_3 = k_Bbar_3   
    
    pKst_Bbar = conc_([         
                                pKst_Bbar_0,
                                pKst_Bbar_1,
                                pKst_Bbar_2,
                                pKst_Bbar_3], axis =1)
    
    # k_Bbar = conc_([k_Bbar_1,
    #                          k_Bbar_2,
    #                          k_Bbar_3,
    #                          k_Bbar_0], axis =1)
    
    q_Bbar = conc_([
                            q_Bbar_0,
                            q_Bbar_1,
                            q_Bbar_2,
                            q_Bbar_3
                            ], axis=1)
    
    # pB_Bbar = conc_([
    #                         pB_Bbar_1,
    #                         pB_Bbar_2,
    #                         pB_Bbar_3,
    #                         pB_Bbar_0
    #                         ], axis=1)

    
    if mode =='numpy':

        beta_Bbar = q_Bbar/q_Bbar_0[:,0:1]

    if mode =='tf':

        beta_Bbar = q_Bbar/tf.expand_dims(q_Bbar_0,axis=-1)
    
    p1_Bbar = lorentz_boost_np(p1, beta_Bbar)
    p2_Bbar = lorentz_boost_np(p2, beta_Bbar)    
    
    
    return  pKst_Bbar, q_Bbar, p1_Bbar, p2_Bbar, p1, p2#pB_Bbar,