import numpy as np
from math import pi
from utils.kin_utils import *
import tensorflow as tf

import zfit
ztf = zfit.ztf



def a(q2, m, m1, m2):
    
    lambd = lambda_function(m1, m2, ztf.sqrt(q2))
    beta= dphi2(ztf.sqrt(q2), m, m)

    step1=q2*(1+ztf.square(beta)) + lambd/2 + 2*m*(ztf.square(m1)-ztf.square(m2)+q2)+4*ztf.square(m)*ztf.square(m1)
    #step2 = beta_l(q2, m)*ztf.sqrt(lambda_phase(m1, m2, q2))*step1
    
    return step1

def b(q2, m, m1, m2):

    lambd = lambda_function(m1, m2, ztf.sqrt(q2))
    beta= dphi2(ztf.sqrt(q2), m, m)
    
    step1=2*(q2*(1+ztf.square(beta)) + m*(ztf.sqrt(lambd)*beta+(ztf.square(m1)-ztf.square(m2)+q2)))
    #step2 = beta_l(q2, m)*ztf.sqrt(lambda_phase(m1, m2, q2))*step1
    
    return step1

def c(q2, m, m1, m2):
    
    lambd = lambda_function(m1, m2, ztf.sqrt(q2))
    beta= dphi2(ztf.sqrt(q2), m, m)

    step1= q2*(1+ztf.square(beta)) - lambd/2 + 2*m*ztf.sqrt(lambd)*beta
    #step2 = beta_l(q2, m)*ztf.sqrt(lambda_phase(m1, m2, q2))*step1
    
    return step1


def matrix_elt(q2, m, m1, m2, cos_theta_l):
    
    out = a(q2, m, m1, m2)+b(q2, m, m1, m2)*cos_theta_l+c(q2, m, m1, m2)*ztf.square(cos_theta_l)
    #out =a(q2, m, m1, m2)+c(q2, m, m1, m2)*ztf.square(cos_theta_l)
    
    return out