"""
SAXS/Models/FormFactors.py

The term "form factor" is used in two different ways in the literature:
1. The form factor amplitude F(q), which is the Fourier transform of the scattering length density.
2. The form factor intensity P(q), which is the square of the amplitude.

In this module, we use "form factor intensity" to refer to the intensity P(q), which is the square of the amplitude F(q).

"""
import numpy as np
import scipy.integrate as integrate

def nquad_vec(f, limits):
    """Vectorized version of scipy.integrate.nquad.
    Integrate a function of multiple variables over given limits.
    This is borrowed from `a scipy issue <https://github.com/scipy/scipy/issues/12209>`_.

    Parameters
    ----------
    f : callable
        The integrand function. It must accept as many arguments as there are
        integration variables.
    limits : list of tuples
        The integration limits for each variable. Each tuple should contain the
        lower and upper limits for the corresponding variable.

    Returns
    -------
    result : float or array-like
        The result of the integration.
    """
    n_vars = len(limits)
    z = np.zeros(n_vars-1)

    def integrand(var, i):
        if i > 0:
            z[i-1] = var
        if i == n_vars - 1:
            return integrate.quad_vec(lambda x: f(*z, x), *limits[i])[0]
        else:
            res = integrate.quad_vec(lambda x: integrand(x, i+1), *limits[i])
            return res if i == 0 else res[0]

    return integrand(None, 0)

def _homogeneous_sphere_impl(q, R):
    """
    Calculate the form factor amplitude of a homogeneous sphere.
    (entirely generated by GitHub Copilot)

    Parameters
    ----------
    q : float or array-like
        The scattering vector magnitude.
    R : float
        The radius of the sphere.

    Returns
    -------
    F : float or array-like
        The form factor of the homogeneous sphere.
    """
    # Calculate the form factor using the formula for a homogeneous sphere
    qR = q * R
    F = (3 * (np.sin(qR) - qR * np.cos(qR))) / qR**3
    return F

def homogeneous_sphere(q, R):
    """
    Calculate the form factor intensity of a homogeneous sphere.

    Parameters
    ----------
    q : float or array-like
        The scattering vector magnitude.
    R : float
        The radius of the sphere.

    Returns
    -------
    P : float or array-like
        The form factor intensity of the homogeneous sphere.
    """
    # Calculate the form factor using the formula for a homogeneous sphere
    return _homogeneous_sphere_impl(q, R)**2

def sphere_volume(R):
    """
    Calculate the volume of a sphere.

    Parameters
    ----------
    R : float
        The radius of the sphere.

    Returns
    -------
    V : float
        The volume of the sphere.
    """
    # Calculate the volume using the formula for a sphere
    V = (4/3) * np.pi * R**3
    return V

def spherical_shell(q, R, r):
    """
    Calculate the form factor intensity of a spherical shell.

    Parameters
    ----------
    q : float or array-like
        The scattering vector magnitude.
    R : float
        The outer radius of the shell.
    r : float
        The inner radius of the shell.

    Returns
    -------
    P : float or array-like
        The form factor intensity of the spherical shell.
    """
    # Calculate the form factor using the formula for a spherical shell
    v1 = sphere_volume(R)
    v2 = sphere_volume(r)
    F = (v1*homogeneous_sphere(q, R) - v2*homogeneous_sphere(q, r)) / (v1 - v2)
    return F**2

def ellipsoid_of_revolution(q, R, epsilon):
    """
    Calculate the form factor intensity of an ellipsoid of revolution.

    Parameters
    ----------
    q : float or array-like
        The scattering vector magnitude.
    R : float
        The semi-major axis of the ellipsoid.
    epsilon : float
        The aspect ratio of the ellipsoid.

    Returns
    -------
    P : float or array-like
        The form factor intensity of the ellipsoid of revolution.
    """
    # Calculate the form factor using the formula for an ellipsoid of revolution
    def r(R, epsilon, alpha):
        return R * np.sqrt(np.sin(alpha)**2 + (epsilon**2) * np.cos(alpha)**2)

    F = integrate.quad(lambda a: _homogeneous_sphere_impl(q, r(R, epsilon, a)), 0, np.pi/2)[0]
    return F**2

def tri_axial_ellipsoid(q, a, b, c):
    """
    Calculate the form factor intensity of a tri-axial ellipsoid.

    Parameters
    ----------
    q : float or array-like
        The scattering vector magnitude.
    a : float
        The semi-major axis of the ellipsoid.
    b : float
        The semi-minor axis of the ellipsoid in the x-y plane.
    c : float
        The semi-minor axis of the ellipsoid in the z direction.

    Returns
    -------
    F : float or array-like
        The form factor intensity of the tri-axial ellipsoid.
    """
    # Calculate the form factor using the formula for a tri-axial ellipsoid
    def r(a, b, c, alpha, beta):
        return np.sqrt(((a**2 * np.sin(beta)**2 + b**2 * np.cos(beta)**2) * np.sin(alpha)**2 + (c * np.cos(alpha))**2))

    F = 2/np.pi * nquad_vec(lambda alpha, beta: _homogeneous_sphere_impl(q, r(a, b, c, alpha, beta))*np.sin(alpha), [[0, np.pi/2], [0, np.pi/2]])[0]
    return F**2
