"""Validation and construction of inverse link functions."""

from __future__ import annotations

import functools
from typing import TYPE_CHECKING, Any, Callable

import jax
import jax.numpy as jnp

if TYPE_CHECKING:
    from .observation_models import Observations

from .utils import one_over_x


def _make_wrapper(func, name, description):
    """Create a wrapper function with combined docstrings."""

    @functools.wraps(func)
    def wrapper(x):
        return func(x)

    # Combine the custom description with original docstring
    original_doc = func.__doc__ or "No docstring available."
    wrapper.__doc__ = f"""{description}

Wrapper of `{func.__module__}.{func.__name__}`. Below the original jax docstring.

{original_doc}"""

    wrapper.__name__ = name
    return wrapper


exp = _make_wrapper(jnp.exp, "exp", "Exponential link function.")

softplus = _make_wrapper(jax.nn.softplus, "softplus", "Softplus link function.")

logistic = _make_wrapper(
    jax.lax.logistic, "logistic", "Logistic (sigmoid) link function."
)

norm_cdf = _make_wrapper(
    jax.scipy.stats.norm.cdf, "norm_cdf", "Norm CDF link function."
)

expit = _make_wrapper(
    jax.scipy.special.expit, "expit", "Expit (sigmoid) link function."
)

identity = _make_wrapper(lambda x: x, "identity", "Identity link function.")

log_softmax = _make_wrapper(
    jax.nn.log_softmax, "log_softmax", "Log-softmax link function."
)


LINK_NAME_TO_FUNC = {
    "exp": exp,
    "expit": expit,
    "jax._src.lax.lax.logistic": expit,
    "jax._src.nn.functions.softplus": softplus,
    "jax._src.nn.functions.log_softmax": log_softmax,
    "jax._src.numpy.ufuncs.exp": exp,
    "jax._src.scipy.special.expit": expit,
    "jax._src.scipy.stats.norm.cdf": jax.scipy.stats.norm.cdf,
    "jax.lax.logistic": logistic,
    "jax._src.scipy.special.logistic": logistic,
    "jax.nn.softplus": softplus,
    "jax.nn.log_softmax": log_softmax,
    "jax.numpy.exp": jnp.exp,
    "jax.scipy.special.expit": expit,
    "jax.scipy.stats.norm.cdf": norm_cdf,
    "logistic": logistic,
    "nemos.inverse_link_function_utils.exp": exp,
    "nemos.inverse_link_function_utils.expit": expit,
    "nemos.inverse_link_function_utils.logistic": logistic,
    "nemos.inverse_link_function_utils.norm_cdf": norm_cdf,
    "nemos.inverse_link_function_utils.one_over_x": one_over_x,
    "nemos.inverse_link_function_utils.softplus": softplus,
    "nemos.inverse_link_function_utils.log_softmax": log_softmax,
    "nemos.utils.one_over_x": one_over_x,
    "norm.cdf": norm_cdf,
    "one_over_x": one_over_x,
    "softplus": softplus,
    "log_softmax": log_softmax,
    "identity": lambda x: x,
}


def link_function_from_string(link_name: str):
    """
    Get a link function from a given name.

    Parameters
    ----------
    link_name:
        A string representation of the link function, e.g. "jax.numpy.exp".

    Returns
    -------
    :
        The link function corresponding to the provided name.

    Raises
    ------
    ValueError:
        If the provided string does not match any known link function.
    """
    if link_name in LINK_NAME_TO_FUNC:
        return LINK_NAME_TO_FUNC[link_name]
    else:
        raise ValueError(
            f"Unknown link function: {link_name}. "
            f"Link function must be one of {list(LINK_NAME_TO_FUNC.keys())}"
            f"if you want to use a custom link function, please provide it as a Callable."
            f"if you think this is a bug, please open an issue at 'https://github.com/flatironinstitute/nemos/issues'."
        )


def check_inverse_link_function(inverse_link_function: Callable):
    """
    Check if the provided inverse_link_function is usable.

    This function verifies if the inverse link function:

    1. Is callable
    2. Returns a jax.numpy.ndarray
    3. Is differentiable (via jax)

    Parameters
    ----------
    inverse_link_function :
        The function to be checked.

    Raises
    ------
    TypeError
        If the function is not callable.
    ValueError
        If the function does not return a jax.numpy.ndarray,
        or is not differentiable.
    """
    if inverse_link_function in LINK_NAME_TO_FUNC.values():
        return

    # check that it's callable
    if not callable(inverse_link_function):
        raise TypeError("The `inverse_link_function` function must be a Callable!")

    # check if the function returns a jax array for a 1D array
    array_out = inverse_link_function(jnp.array([1.0, 2.0, 3.0]))
    if not isinstance(array_out, jnp.ndarray):
        raise ValueError("The `inverse_link_function` must return a jax.numpy.ndarray!")

    # Optionally: Check for scalar input
    scalar_out = inverse_link_function(1.0)
    if not isinstance(scalar_out, (jnp.ndarray, float, int)):
        raise ValueError(
            "The `inverse_link_function` must handle scalar inputs correctly and return a scalar or a "
            "jax.numpy.ndarray!"
        )

    # check for autodiff
    try:
        gradient_fn = jax.grad(inverse_link_function)
        gradient_fn(1.0)
    except Exception as e:
        raise ValueError(
            f"The `inverse_link_function` function cannot be differentiated. Error: {e}"
        ) from e


def resolve_inverse_link_function(
    inverse_link_function: Any, observation_model: Observations
) -> Callable:
    """
    Validate and resolve an inverse link function specification.

    This helper function implements shared logic for setting inverse link functions
    in GLM-based models. It handles three cases:
    1. None: returns the observation model's default inverse link function
    2. String: parses and returns the corresponding link function
    3. Callable: validates and returns the function as-is

    Parameters
    ----------
    inverse_link_function : Callable, str, or None
        The inverse link function specification. Can be:
        - None: use the default from the observation model
        - str: name of a standard link function (e.g., "identity", "log", "logit")
        - Callable: a custom inverse link function
    observation_model : Observations
        The observation model instance that provides the default inverse link function.

    Returns
    -------
    Callable
        A validated inverse link function that maps linear predictions to the
        observation space.

    Raises
    ------
    TypeError
        If the provided function is not callable or a string.
    ValueError
        - If a callable is provided, but it does not return a jax.numpy.ndarray or scalar, or if it
        is not differentiable with respect to its inputs.
        - If a string is provided, but it cannot be parsed to a callable.

    Examples
    --------
    >>> import nemos as nmo
    >>> from nemos.inverse_link_function_utils import resolve_inverse_link_function
    >>> obs_model = nmo.observation_models.PoissonObservations()
    >>> # Use default link function
    >>> link_fn = resolve_inverse_link_function(None, obs_model)
    >>>
    >>> # Use named link function
    >>> link_fn = resolve_inverse_link_function("softplus", obs_model)
    >>>
    >>> # Use custom function
    >>> custom_fn = lambda x: x**2
    >>> link_fn = resolve_inverse_link_function(custom_fn, obs_model)
    """
    if inverse_link_function is None:
        return observation_model.default_inverse_link_function

    elif isinstance(inverse_link_function, str):
        return link_function_from_string(inverse_link_function)
    check_inverse_link_function(inverse_link_function)
    return inverse_link_function
