from typing import Callable, Tuple

import scipy.linalg

from jax import jit, lax
import jax.numpy as jnp
from jax._src.numpy.linalg import norm
from jax._src.numpy.util import _wraps
from jax._src.scipy.linalg import rsf2csf, schur
from jax._src.typing import ArrayLike, Array


@jit
def _algorithm_11_1_1(F: Array, T: Array) -> Tuple[Array, Array]:
  # Algorithm 11.1.1 from Golub and Van Loan "Matrix Computations"
  N = T.shape[0]
  minden = jnp.abs(T[0, 0])

  def _outer_loop(p, F_minden):
    _, F, minden = lax.fori_loop(1, N-p+1, _inner_loop, (p, *F_minden))
    return F, minden

  def _inner_loop(i, p_F_minden):
    p, F, minden = p_F_minden
    j = i+p
    s = T[i-1, j-1] * (F[j-1, j-1] - F[i-1, i-1])
    T_row, T_col = T[i-1], T[:, j-1]
    F_row, F_col = F[i-1], F[:, j-1]
    ind = (jnp.arange(N) >= i) & (jnp.arange(N) < j-1)
    val = (jnp.where(ind, T_row, 0) @ jnp.where(ind, F_col, 0) -
            jnp.where(ind, F_row, 0) @ jnp.where(ind, T_col, 0))
    s = s + val
    den = T[j-1, j-1] - T[i-1, i-1]
    s = jnp.where(den != 0, s / den, s)
    F = F.at[i-1, j-1].set(s)
    minden = jnp.minimum(minden, jnp.abs(den))
    return p, F, minden

  return lax.fori_loop(1, N, _outer_loop, (F, minden))

_FUNM_LAX_DESCRIPTION = """\
The array returned by :py:func:`jax.scipy.linalg.funm` may differ in dtype
from the array returned by py:func:`scipy.linalg.funm`. Specifically, in cases
where all imaginary parts of the array values are close to zero, the SciPy
function may return a real-valued array, whereas the JAX implementation will
return a complex-valued array.

Additionally, unlike the SciPy implementation, when ``disp=True`` no warning
will be printed if the error in the array output is estimated to be large.
"""

@_wraps(scipy.linalg.funm, lax_description=_FUNM_LAX_DESCRIPTION)
def funm(A: ArrayLike, func: Callable[[Array], Array], disp: bool = True) -> Tuple[Array, Array]:
  A = jnp.asarray(A)
  if A.ndim != 2 or A.shape[0] != A.shape[1]:
    raise ValueError('expected square array_like input')

  T, Z = schur(A)
  T, Z = rsf2csf(T, Z)

  F = jnp.diag(func(jnp.diag(T)))
  F = F.astype(T.dtype.char)

  F, minden = _algorithm_11_1_1(F, T)
  F = Z @ F @ Z.conj().T

  if disp:
    return F

  if F.dtype.char.lower() == 'e':
    tol = jnp.finfo(jnp.float16).eps
  if F.dtype.char.lower() == 'f':
    tol = jnp.finfo(jnp.float32).eps
  else:
    tol = jnp.finfo(jnp.float64).eps

  minden = jnp.where(minden == 0.0, tol, minden)
  err = jnp.where(jnp.any(jnp.isinf(F)), jnp.inf, jnp.minimum(1, jnp.maximum(
          tol, (tol / minden) * norm(jnp.triu(T, 1), 1))))

  return F, err