from __future__ import annotations

import importlib.resources
from collections.abc import Callable
from typing import cast

import numpy as np
from numpy.typing import NDArray
from scipy.interpolate import interp1d  # type: ignore[import-untyped]

from . import beam_data


def beam_transform_to_profile(
    B_ell: NDArray[np.float64], theta: NDArray[np.float64], normalize: bool = False
) -> NDArray[np.float64]:
    """
    Given the transform B(ell) of a beam, evaluate its real space angular profile
    at a given an angle theta in radians. This code is adapted from the pixell
    library under the utils.py module (https://pixell.readthedocs.io/en/latest/readme.html).

    Args:
        B_ell: (`float` or `NDArray[float]`): The multipole transform of the beam.
        theta: (`float` or `NDArray[float]`): The angle in radians at which to evaluate the beam profile.
        normalize: (`bool`): If True, normalize the beam profile.

    Returns:
        profile: (`float` or `NDArray[float]`): The beam profile at the given angle theta.
    """
    # check to make sure B_ell is an array
    B_ell = np.asarray(B_ell)

    # multipole moments
    ell = np.arange(B_ell.size)

    # cosine of the angle for legendre polynomials
    x = np.cos(theta)

    # c_ell coefficients
    coefficients = B_ell * (2 * ell + 1) / (4 * np.pi)

    # evaluate the Legendre polynomial at the given angle
    profile = np.polynomial.legendre.legval(x, coefficients)  # type: ignore[no-untyped-call]

    if normalize:  # pragma: no cover
        profile /= np.sum(coefficients)

    return cast(NDArray[np.float64], profile)


def generate_beam_profile(
    frequency: str, space: str = "real"
) -> Callable[[NDArray[np.float64]], NDArray[np.float64]]:
    """
    This function generates the beam profile for a given frequency. The beam profile
    is generated by transforming the beam from the multipole space to the real space
    angular profile. The beam profile returned is a function of the angle theta in radians.

    Args:
        frequency: (`str`): The frequency of the beam. Can be either 'f150' or 'f090'.
        space: (`str`): The space in which to return the beam profile. Can be either 'real' or 'harmonic'.

    Returns:
        profile_func: (`Callable`): The beam profile as a function of the angle theta in radians.
    """

    if frequency == "f150":
        beam_file = "act_planck_dr5.01_s08s18_f150_daynight_beam.txt"
    elif frequency == "f090":
        beam_file = "act_planck_dr5.01_s08s18_f090_daynight_beam.txt"
    else:  # pragma: no cover
        msg = "Frequency must be either 'f150' or 'f090'"
        raise ValueError(msg)

    beam_file = importlib.resources.files(beam_data) / beam_file  # type: ignore[attr-defined]

    # load the beam transform
    ell, B_ell = np.loadtxt(beam_file, unpack=True)

    if space == "real":
        # define theta array used for interpolation
        theta_ = np.linspace(0, 100, 1000) * np.radians(1 / 60)  # radians

        # generate profile
        profile = beam_transform_to_profile(B_ell, theta_)

        # normalize the profile
        integral = np.trapz(profile * 2 * np.pi * theta_, x=theta_)

        # interpolate the profile
        def profile_func(theta: NDArray[np.float64]) -> NDArray[np.float64]:
            return np.interp(theta, theta_, profile / integral, period=np.pi)

    elif space == "harmonic":
        profile_func = interp1d(
            ell, B_ell, kind="linear", bounds_error=False, fill_value=0.0
        )
    else:  # pragma: no cover
        msg = "Space must be either 'real' or 'harmonic'"
        raise ValueError(msg)

    return profile_func
