import pynam.util

from typing import Callable
import numpy as np


def get_zeta_p_integrand(eps: Callable[[float], complex]) -> Callable[[float, float], complex]:
	""" Gets the integrand function zeta_p_integrand(u, y).

	Returns zeta_p_integrand(u, y), a complex valued function of two momenta in units of vacuum wavelength.

	:param eps:
	:return:
	"""
	def zeta_p_integrand(y: float, u: float) -> complex:
		"""
		Here y and u are in units of vacuum wavelength, coming from Ford-Weber / from the EWJN noise expressions.
		:param u:
		:param y:
		:return:
		"""
		u2 = u ** 2
		y2 = y ** 2
		k2 = u2 + y2
		k = np.sqrt(k2)
		eps_value = eps(k)
		term_1 = y2 / (eps_value - k2)
		term_2 = u2 / eps_value
		return (term_1 + term_2) / k2

	return zeta_p_integrand


def get_zeta_p_function(eps: Callable[[float], complex]):
	def zeta_p(u: float) -> complex:
		zeta_p_integrand = get_zeta_p_integrand(eps)

		integral_result = pynam.util.complex_quad(lambda y: zeta_p_integrand(y, u), 0, np.inf)

		print(integral_result)
		integral = integral_result[0]

		return integral * 2j

	return zeta_p