import numpy as np
from math import pi
import tensorflow as tf

import zfit
ztf = zfit.ztf



def metric_tensor_np():
    """Metric tensor for Lorentz space (constant)."""
    return np.array([1.,-1.,-1.,-1.])


def metric_tensor():
    """Metric tensor for Lorentz space (constant)."""
    return tf.constant([1., -1., -1., -1.], dtype=tf.float64)


def scalar_product_3(v1, v2):
    """Calculate scalar product of two 3-vectors.
    Arguments:
        vec1: First vector.
        vec2: Second vector.
    """
    return tf.reduce_sum(v1 * v2, axis=1)    

def scalar_product_4(v1, v2):
    """Calculate mass scalar for Lorentz 4-momentum.
    Arguments:
        vector: Input Lorentz momentum vector.
    """
    return tf.reduce_sum(tf.multiply(tf.multiply(v1,v2),metric_tensor()),
                                 axis=-1, keepdims=True)

def scalar_product_3_np(v1,v2):    

    
    output = (v1*v2)[:,1:4].sum(axis=1, keepdims=True)
    
    return output


def scalar_product_4_np(v1,v2):
    
    
    output = (v1*v2*metric_tensor_np()).sum(axis=1 ,keepdims=True)
    return output

def get_costheta_l(p1, p2):  
    
    num = scalar_product_3(p1, p2)

    den1= ztf.sqrt(scalar_product_3(p1, p1))
    den2= ztf.sqrt(scalar_product_3(p2, p2))
    
    costheta_l = num/(den1*den2)
    
    return costheta_l

def get_costheta_l_np(p1, p2):  
    
    num = scalar_product_3_np(p1, p2)

    den1= np.sqrt(scalar_product_3_np(p1, p1))
    den2= np.sqrt(scalar_product_3_np(p2, p2))
    
    costheta_l = num/(den1*den2)
    
    return costheta_l

def beta_l(q2, m):
    beta_l = ztf.sqrt(1-4*(ztf.square(m)/q2))
    return beta_l

def beta_l_np(q2, m):
    beta_l = np.sqrt(1-4*(np.square(m)/q2))
    return beta_l

def lambda_function(m, c1, c2):
    
    #lambd = ztf.square(ztf.square(m)- ztf.square(c1)) -2*ztf.square(c2)*(ztf.square(m) + ztf.square(c1)) +ztf.square(ztf.square(c2))
    lambd = (ztf.square(m) - ztf.square(c1 + c2))*(ztf.square(m) - ztf.square(c1 - c2))
    return lambd

def lambda_function_np(m, c1, c2):
    
    #lambd = np.square(np.square(m)- np.square(c1)) -2*np.square(c2)*(np.square(m) + np.square(c1)) +np.square(np.square(c2))
    lambd = (np.square(m) - np.square(c1 + c2))*(np.square(m) - np.square(c1 - c2))

    return lambd    

def dphi2_tf(m, c1, c2):

    lambd = lambda_function(m, c1, c2)

    out = ztf.sqrt(lambd)/(ztf.square(m))

    return out

def dphi2_np(m, c1, c2):

    lambd = lambda_function_np(m, c1, c2)

    out = np.sqrt(lambd)/(np.square(m))

    return out

def spatial_component(vector):

	return tf.gather(vector, indices=[1, 2, 3], axis=-1)

def time_component(vector):

	return tf.gather(vector, indices=[0], axis=-1)

def lorentz_vector(space, time):
    """Make a Lorentz vector from spatial and time components.
    Arguments:
        space: 3-vector of spatial components.
        time: Time component.
    """
    return tf.concat([space, time], axis=-1)

def lorentz_boost(fvector_to_boost, boost_vector):
    
   # boost_vector_time_component= time_component(boost_vector)
    #boost_vector_spatial = 
    
    beta = spatial_component(boost_vector)
    beta2 = tf.expand_dims(scalar_product_3(beta,beta), axis=-1)
    
    gamma = 1./ztf.sqrt(1.-beta2)
    gamma_ = (gamma-1.)/beta2
    
    ve = time_component(fvector_to_boost)
    vp = spatial_component(fvector_to_boost)
    
    ve2 = gamma*(ve + tf.expand_dims(scalar_product_3(vp,beta), axis=-1))
    vp2 = vp + (gamma*ve + gamma_*tf.expand_dims(scalar_product_3(vp, beta), axis=-1))*beta

    boosted_fvector=lorentz_vector(ve2,vp2)

    return boosted_fvector


def lorentz_boost_np(fvector_to_boost, boost_vector):
    
    beta = boost_vector
    beta2 = scalar_product_3_np(beta,beta)
    
    gamma = 1./np.sqrt(1.-beta2)
    gamma_ = (gamma-1.)/beta2
    
    ve = fvector_to_boost[:,0].reshape(-1,1)
    vp = fvector_to_boost

    ve2 = gamma*(ve + scalar_product_3_np(fvector_to_boost,beta))
    vp2 = vp + (gamma*ve + gamma_*scalar_product_3_np(fvector_to_boost, beta))*beta

    boosted_fvector = np.concatenate([ve2,vp2[:,1:4]],axis=1)
    return boosted_fvector



def draw_random_angle(q2, mode):

    if mode=='numpy':
        random_phi_ = np.random.uniform(size=(q2.shape[0],1), low=-pi, high=pi, )

    if mode=='tf':

        random_phi_ = tf.random.uniform(shape=tf.shape(q2), minval=-pi, maxval=pi, dtype=tf.float64)

    return random_phi_

def draw_random_costheta(q2, mode):

    if mode=='numpy':
        random_cosine_ = np.random.uniform(size=(q2.shape[0],1), low=-1.0, high=1.0, )
    if mode=='tf':
        random_cosine_ = tf.random.uniform(shape=tf.shape(q2), minval=-1., maxval=1.0, dtype=tf.float64)


    return random_cosine_

def draw_random_ones_minusones(q2, mode):
    
    
    if mode=='numpy':
        n_events=q2.shape[0]
        mask = np.random.uniform(size=(n_events,1), low=0, high=1. )>0.5
        random_one_=np.where(mask,1.,-1.)

    if mode=='tf':
        mask = tf.random.uniform(shape=tf.shape(q2), minval=0., maxval=1.0, dtype=tf.float64)
        random_one_ = tf.where(mask>0.5,tf.ones_like(q2),-tf.ones_like(q2))

    return random_one_