import numpy as np
import matplotlib.pyplot as plt
from scipy.special import j0, j1
from scipy.integrate import quad, quad_vec

# pip install git+https://github.com/brandonwillard/numba-scipy.git@vectorize-overloads #egg=numba-scipy
from numba import jit, njit

import time



# We store the range of n, kn and gn in arrays of length N_max
N_max = 20
n_values = np.linspace(0, N_max-1, N_max, dtype=int)
k_n_values = 2 * np.pi * n_values


def perform_integrals_quad(n, t, z, epsilon = 1e-1):
    '''
    We perform the integrals primarily using quad
    '''

    # We create the matrix of possible u(t,z) values, defined as u^2 = t^2 - z^2
    def u(t,z):
        return np.where(t <= z, 0, np.sqrt(np.maximum(t**2 - z**2, 0))) 
    
    u_value = u(t, z)

    k_n = 2*np.pi*n

    # We compute the integral from 0 to epsilon
    def integrand(u,k):
        return np.sin(k * u) / u * j1(k*u)

    epsilon_part,_ = quad(integrand, epsilon, u_value, args=(k_n), limit=100000, maxp1=100, epsabs=1e-6, epsrel=1e-3) # We use quad

    # We add everything and return it
    return epsilon_part




start_time = time.time()  # Record start time
# We compute the integral from 0 to epsilon
@njit
def integrand_numba(u,k):
    return np.sin(k * u) / u * j1(k*u)
integrand_numba(1.,2.)
end_time = time.time()    # Record end time
execution_time_one = end_time - start_time  # Calculate elapsed time
print(f"Execution time for compilation of integrand_numba: {execution_time_one} seconds")

def perform_integrals_quad_numba(integrand, n, t, z, epsilon = 1e-1):
    '''
    We perform the integrals using quad and first JITing the integrand 
    '''

    # We create the matrix of possible u(t,z) values, defined as u^2 = t^2 - z^2
    def u(t,z):
        return np.where(t <= z, 0, np.sqrt(np.maximum(t**2 - z**2, 0))) 
    
    u_value = u(t, z)

    k_n = 2*np.pi*n

    
    epsilon_part,_ = quad(integrand, epsilon, u_value, args=(k_n), limit=100000, maxp1=100, epsabs=1e-6, epsrel=1e-3) # We use quad

    return epsilon_part




start_time = time.time()  # Record start time
# We compute the integral from 0 to epsilon
@njit(parallel=True)
def integrand_vec_numba(u):
    return np.sin(k_n_values * u) / u * j1(k_n_values*u)
integrand_vec_numba(1.)
end_time = time.time()    # Record end time
execution_time_one = end_time - start_time  # Calculate elapsed time
print(f"Execution time for compilation of integrand_vec_numba: {execution_time_one} seconds")

def perform_integrals_quad_vec_numba(integrand_vec_numba, t, z, epsilon = 1e-1):
    '''
    We perform the integrals using quadvec and first JITing the integrand 
    '''

    # We create the matrix of possible u(t,z) values, defined as u^2 = t^2 - z^2
    def u(t,z):
        return np.where(t <= z, 0, np.sqrt(np.maximum(t**2 - z**2, 0))) 
    
    u_value = u(t, z)


    epsilon_part,_ = quad_vec(integrand_vec_numba, epsilon, u_value, limit=100000, epsabs=1e-6, epsrel=1e-3, workers=1) # We use quad

    return epsilon_part



def integrand_vec(u):
        return np.sin(k_n_values * u) / u * j1(k_n_values*u)

def perform_integrals_quad_vec(integrand_vec, t, z, epsilon = 1e-1):
    '''
    We perform the integrals using quad_vec
    '''

    # We create the matrix of possible u(t,z) values, defined as u^2 = t^2 - z^2
    def u(t,z):
        return np.where(t <= z, 0, np.sqrt(np.maximum(t**2 - z**2, 0))) 
    
    u_value = u(t, z)


    epsilon_part,_ = quad_vec(integrand_vec, epsilon, u_value, limit=100000, epsabs=1e-6, epsrel=1e-3, workers=1) # We use quad


    return epsilon_part



# Exact integral
def perform_integrals_exact(n, t, z):
    k = 2*np.pi*n
    return - np.cos(k*t) * j0(k*t) - np.sin(k*t) * j1(k*t)




if __name__ == "__main__":

    epsilon = 2
    t = 25
    z = 0.

    quad_result = np.empty(len(n_values))
    start_time = time.time()  # Record start time
    for n in n_values:
        quad_result[n] = perform_integrals_quad(n, t, z, epsilon)
    end_time = time.time()    # Record end time
    execution_time_one = end_time - start_time  # Calculate elapsed time
    print(f"Execution time for perform_integrals_quad: {execution_time_one} seconds")


    quad_numba_result = np.empty(len(n_values))
    start_time = time.time()  # Record start time
    for n in n_values:
        quad_numba_result[n] = perform_integrals_quad_numba(integrand_numba, n, t, z, epsilon)
    end_time = time.time()    # Record end time
    execution_time_one = end_time - start_time  # Calculate elapsed time
    print(f"Execution time for perform_integrals_quad_numba: {execution_time_one} seconds")


    quad_vec_result = np.empty(len(n_values))
    start_time = time.time()  # Record start time
    quad_vec_result = perform_integrals_quad_vec(integrand_vec, t, z, epsilon)
    end_time = time.time()    # Record end time
    execution_time_one = end_time - start_time  # Calculate elapsed time
    print(f"Execution time for perform_integrals_quad_vec: {execution_time_one} seconds")


    quad_vec_numba_result = np.empty(len(n_values))
    start_time = time.time()  # Record start time
    quad_vec_numba_result = perform_integrals_quad_vec_numba(integrand_vec_numba, t, z, epsilon)
    end_time = time.time()    # Record end time
    execution_time_one = end_time - start_time  # Calculate elapsed time
    print(f"Execution time for perform_integrals_quad_vec_numba: {execution_time_one} seconds")



    exact_result = np.empty(len(n_values))
    for n in n_values:
        exact_result[n] = perform_integrals_exact(n, t, z) - perform_integrals_exact(n, epsilon, z)



    # Create the plot
    plt.plot(n_values, quad_numba_result, label='Quad Numba', color='blue', marker='o')
    plt.plot(n_values, quad_result, label='Normal quad', color='red', marker='o') 

    plt.plot(n_values, quad_vec_numba_result, label='Quad Vec', color='black', marker='o')
    plt.plot(n_values, quad_vec_result, label='Quad Vec Numba', color='purple', marker='o') 

    plt.plot(n_values, exact_result, label='Exact', color='green') 

    # Add labels and title
    plt.xlabel('n')
    plt.ylabel('Integral')
    plt.title('Integral in function of n for two different integration methods')

    # Show legend
    plt.legend()

    # Display the plot
    plt.show()