from __future__ import annotations

from typing import Literal

import jax
import jax.numpy as jnp
import numpy as np
import scipy.special as sc
from array_api_compat import array_namespace, is_jax_namespace, is_torch_namespace

from .annotations import Array, BooleanArray, Namespace


def abs_square[T: Array](x: T) -> T:
    xp = array_namespace(x)
    return xp.square(x.real) + xp.square(x.imag)  # pyright: ignore


# TODO: Remove when it's added to theArrayAPI:
# https://github.com/data-apis/array-api/issues/242
def outer_product[T: Array](x: T, y: T) -> T:
    """Return the broadcasted outer product of a vector with itself.

    This is xp.einsum("...i,...j->...ij", x, y).
    """
    xp = array_namespace(x, y)
    xi = xp.reshape(x, (*x.shape, 1))
    yj = xp.reshape(y.conjugate(), (*y.shape[:-1], 1, y.shape[-1]))
    return xi * yj


def matrix_vector_mul[T: Array](x: T, y: T) -> T:
    """Return the matrix-vector product.

    This is xp.einsum("...ij,...j->...i", x, y).

    Note the speed difference:
    * 14.3 µs: xp.vecdot(matrix_vector_mul(m, x), x)
    * 4.44 µs: np.einsum("...i,...ij,...j->...", x, m, x)
    """
    xp = array_namespace(x, y)
    y = xp.reshape(y, (*y.shape[:-1], 1, y.shape[-1]))
    return xp.sum(x * y, axis=-1)


def matrix_dot_product[T: Array](x: T, y: T) -> T:
    """Return the "matrix dot product" of a matrix with the outer product of a vector.

    This equals:
    * 1.19 µs: xp.einsum("...ij,...ij", x, y)
    * 1.77 µs: xp.sum(x * y, axis=(-2, -1))
    # 3.87 µs: xp.linalg.trace(xp.moveaxis(x, -2, -1) @ y)
    """
    xp = array_namespace(x, y)
    return xp.sum(x * y, axis=(-2, -1))


def divide_where[T: Array](
    dividend: T, divisor: T, *, where: BooleanArray | None = None, otherwise: T | None = None
) -> T:
    """Return the quotient or a special value when a condition is false.

    Returns: `xp.where(where, dividend / divisor, otherwise)`, but without evaluating
    `dividend / divisor` when `where` is false.  This prevents both infinite cotangents, and
    some exceptions.
    """
    if where is None:
        assert otherwise is None
        xp = array_namespace(dividend, divisor)
        return xp.divide(dividend, divisor)
    assert otherwise is not None
    xp = array_namespace(dividend, divisor, where, otherwise)
    dividend = xp.where(where, dividend, 1.0)
    divisor = xp.where(where, divisor, 1.0)
    quotient: T = xp.divide(dividend, divisor)
    return xp.where(where, quotient, otherwise)


# Remove when https://github.com/scipy/scipy/pull/18605 is released.
def softplus[T: Array](x: T, /, *, xp: Namespace | None = None) -> T:
    """Softplus, which is log(1+exp(x)).

    This has asymptotic behavior exp(x) as x -> -inf and x as x -> +inf.
    """
    if xp is None:
        xp = array_namespace(x)
    return xp.logaddexp(xp.asarray(0.0), x)


def log_softplus[T: Array](x: T, /, *, xp: Namespace | None = None) -> T:
    """Log-softplus, which is log(1+log(1+exp(x))).

    This has asymptotic behavior 0 as x -> -inf and log(x) as x -> +inf.
    """
    if xp is None:
        xp = array_namespace(x)
    z = xp.asarray(0.0)
    return xp.logaddexp(z, xp.logaddexp(z, x))


def bessel_iv_ratio(v: jax.Array, x: jax.Array, /, *, iterations: int = 200) -> jax.Array:
    """Return ``I_(v + 1)(x) / I_v(x)`` for nonnegative real ``x``.

    The implementation uses a backward recurrence, which is stable for the ratio and remains fully
    differentiable with respect to both ``v`` and ``x``.
    """
    v, x = jnp.broadcast_arrays(jnp.asarray(v), jnp.asarray(x))
    q = jnp.zeros_like(v + x)

    def body(i: int, q: jax.Array) -> jax.Array:
        order = v + (iterations - i)
        return x / (2.0 * order + x * q)

    return jax.lax.fori_loop(0, iterations, body, q)


def _scipy_log_bessel_ive(v: np.ndarray, x: np.ndarray) -> np.ndarray:
    return np.log(sc.ive(v, x))


def _log_bessel_ive_callback(v: jax.Array, x: jax.Array) -> jax.Array:
    v, x = jnp.broadcast_arrays(jnp.asarray(v), jnp.asarray(x))
    shape = v.shape
    dtype = jnp.result_type(v, x, jnp.float32)

    def callback(v_np: np.ndarray, x_np: np.ndarray) -> np.ndarray:
        return _scipy_log_bessel_ive(v_np, x_np).astype(dtype, copy=False)

    return jax.pure_callback(callback, jax.ShapeDtypeStruct(shape, dtype), v, x)


@jax.custom_jvp
def log_bessel_ive(v: jax.Array, x: jax.Array, /) -> jax.Array:
    """Return ``log(ive(v, x))`` for nonnegative real ``x``.

    The primal value is delegated to SciPy because JAX does not currently expose ``ive`` for
    general real order. The custom JVP provides derivatives with respect to both ``v`` and ``x``.
    """
    return _log_bessel_ive_callback(v, x)


@log_bessel_ive.defjvp
def _log_bessel_ive_jvp(
    primals: tuple[jax.Array, jax.Array], tangents: tuple[jax.Array, jax.Array]
) -> tuple[jax.Array, jax.Array]:
    v, x = primals
    v_dot, x_dot = tangents
    primal_out = log_bessel_ive(v, x)

    safe_x = jnp.maximum(jnp.asarray(x), jnp.sqrt(jnp.finfo(jnp.result_type(x, jnp.float32)).eps))
    x_tangent = x_dot * (bessel_iv_ratio(v, safe_x) + v / safe_x - 1.0)

    step = jnp.cbrt(jnp.finfo(jnp.result_type(v, jnp.float32)).eps) * (1.0 + jnp.abs(v))
    v_tangent = (
        v_dot
        * (_log_bessel_ive_callback(v + step, x) - _log_bessel_ive_callback(v - step, x))
        / (2.0 * step)
    )
    return primal_out, v_tangent + x_tangent


def sublinear_softplus[T: Array](x: T, maximum: T, /, *, xp: Namespace | None = None) -> T:
    """Sublinear-softplus, which is softplus(x) / (1 + softplus(x) / maximum).

    This has asymptotic behavior exp(x) as x -> -inf and maximum as x -> +inf.
    """
    if xp is None:
        xp = array_namespace(x)
    o = xp.asarray(1.0)
    sp = softplus(x, xp=xp)
    return sp / (o + sp / maximum)


def inverse_softplus[T: Array](y: T, /, *, xp: Namespace | None = None) -> T:
    if xp is None:
        xp = array_namespace(y)
    return xp.where(
        y > 80.0,  # noqa: PLR2004
        y,
        xp.log(xp.expm1(y)),
    )


def normalize[T: Array](
    mode: Literal["l1", "l2", "max"], x: T, *, axis: tuple[int, ...] | int | None = None
) -> T:
    """Returns the L1-normalized copy of x, assuming that x is nonnegative."""
    xp = array_namespace(x)
    epsilon = 10 * xp.finfo(x.dtype).eps
    match mode:
        case "l1":
            sum_x = xp.sum(xp.abs(x), axis=axis, keepdims=True)
            size = x.size / sum_x.size
            return xp.where(sum_x < epsilon, xp.ones_like(x) / size, x / sum_x)
        case "l2":
            sum_x = xp.sqrt(xp.sum(xp.square(x), axis=axis, keepdims=True))
            size = x.size / sum_x.size
            return xp.where(sum_x < epsilon, xp.ones_like(x) * xp.pow(size, -0.5), x / sum_x)
        case "max":
            sum_x = xp.max(xp.abs(x), axis=axis, keepdims=True)
            return xp.where(sum_x < epsilon, xp.ones_like(x), x / sum_x)


def stop_gradient[U](x: U, *, xp: Namespace | None = None) -> U:
    if xp is None:
        xp = array_namespace(x)
    if is_jax_namespace(xp):
        from jax.lax import stop_gradient as sg  # noqa: PLC0415

        return sg(x)
    if is_torch_namespace(xp):
        from torch import Tensor  # noqa: PLC0415 # pyright: ignore # type: ignore

        assert isinstance(x, Tensor)
        return x.detach()  # pyright: ignore
    return x
