import warnings
from typing import Any, Literal, Protocol, overload

import attrs
import numpy as np
import scipy
import scipy.special
from numpy.typing import NDArray

from jordan_form._multiplicity import (
    MultiplicityProtocol,
    _matrix_rank_from_s,
    get_tol,
)


@attrs.frozen(kw_only=True)
class CanonicalJordanChains(MultiplicityProtocol):
    """Canonical system of Jordan chains."""

    eigval: float
    chains: list[np.ndarray[tuple[int, int], np.dtype[np.number]]]
    """The Jordan chains of shape (l_chain, n)."""

    @property
    def chain_lengths(self) -> np.ndarray[tuple[int], np.dtype[np.int_]]:
        """
        The length of the Jordan chains.

        Known as Segre number.

        References
        ----------
        大竹 剛, 古賀 雅伸 and 三平 満司. 2002.
        行列のジョルダン標準形の数値計算法.
        システム制御情報学会論文誌 15, 7 (2002), 320–326.
        https://doi.org/10.5687/iscie.15.320

        """
        if len(self.chains) == 0:
            return np.empty((0,), dtype=np.int_)
        return np.stack([c.shape[0] for c in self.chains])

    @property
    def dim_ith_generalized_eigenvectors(
        self,
    ) -> np.ndarray[tuple[int], np.dtype[np.int_]]:
        """
        The dimension of the space generated by the i-th generalized eigenvectors
        for oridinary eigenvalue problem.

        The dimension of the quotient space
        (i-th generalized eigenspace) / (i-1-th generalized eigenspace).

        Known as Weyr number.

        References
        ----------
        大竹 剛, 古賀 雅伸 and 三平 満司. 2002.
        行列のジョルダン標準形の数値計算法.
        システム制御情報学会論文誌 15, 7 (2002), 320–326.
        https://doi.org/10.5687/iscie.15.320

        """
        if len(self.chains) == 0:
            return np.empty((0,), dtype=np.int_)
        chain_lengths = self.chain_lengths
        return np.count_nonzero(
            chain_lengths[:, None] >= np.arange(chain_lengths.max())[None, :] + 1,
            axis=0,
        )

    @property
    def dim_generalized_eigenspace(self) -> np.ndarray[tuple[int], np.dtype[np.int_]]:
        """
        The dimension of the generalized eigenspaces for oridinary eigenvalue problem.

        Known as Kågström number.

        In general, the element of Jordan chains is not orthogonal
        and this property should not be used.

        References
        ----------
        Bo Kågström and Axel Ruhe. 1980.
        An Algorithm for Numerical Computation
        of the Jordan Normal Form of a Complex Matrix.
        ACM Trans. Math. Softw. 6, 3 (Sept. 1980), 398–419.
        https://doi.org/10.1145/355900.355912

        大竹 剛, 古賀 雅伸 and 三平 満司. 2002.
        行列のジョルダン標準形の数値計算法.
        システム制御情報学会論文誌 15, 7 (2002), 320–326.
        https://doi.org/10.5687/iscie.15.320

        """
        if len(self.chains) == 0:
            return np.empty((0,), dtype=np.int_)
        return np.cumsum(self.dim_ith_generalized_eigenvectors)

    @property
    def num_chains(self) -> int:
        """The number of Jordan chains."""
        return len(self.chains)

    @property
    def algebraic_multiplicity(self) -> int:
        return np.sum(self.chain_lengths, dtype=np.int_)

    @property
    def eigvec_orthogonal(self) -> np.ndarray[tuple[int, int], np.dtype[np.number]]:
        if len(self.chains) == 0:
            return np.empty((0, 0), dtype=np.float64)
        return np.stack([c[0, :] for c in self.chains], axis=1)

    @property
    def geometric_multiplicity(self) -> int:
        return self.eigvec_orthogonal.shape[1]


class MatrixFuncProtocol(Protocol):
    def __call__(
        self,
        eigval: float,
        derv: int,
        /,
    ) -> np.ndarray[tuple[int, int], np.dtype[np.number]] | None:
        """
        Evaluate the matrix function at the eigenvalues.

        Parameters
        ----------
        eigval : NDArray[np.number]
            The eigenvalue.
        derv : int
            Maximum derivative.

        Returns
        -------
        NDArray[np.number]
            The (derv)-th derivative of
            matrix function of shape (n, n).

        """
        ...


def geig_func(
    A: np.ndarray[tuple[int, int], np.dtype[np.number]],
    B: np.ndarray[tuple[int, int], np.dtype[np.number]] | None = None,
    /,
) -> MatrixFuncProtocol:
    """
    Create a generalized eigenvalue function.

    Parameters
    ----------
    A : NDArray[np.number]
        The matrix A of shape (n, n).
    B : NDArray[np.number], optional
        The matrix B of shape (n, n), by default None.
        If None, the identity matrix is used.
        (oridanary eigenvalue problem)

    Returns
    -------
    MatrixFuncProtocol
        The matrix function λ→(A-Bλ) where (A-Bλ) is of shape (n, n).

    """
    B_ = B if B is not None else np.eye(A.shape[0])

    def inner(
        eigval: float, derv: int
    ) -> np.ndarray[tuple[int, int], np.dtype[np.number]] | None:
        derv = int(derv)
        if derv == 0:
            return A - eigval * B_
        elif derv == 1:
            return -B_
        elif derv > 1:
            return np.zeros_like(A)
        raise ValueError(f"Invalid derivative {derv=}. ")

    return inner


def sympy_func(A: Any) -> MatrixFuncProtocol:
    """
    Create a sympy matrix function.

    Parameters
    ----------
    A : Any
        The sympy matrix.

    Returns
    -------
    MatrixFuncProtocol
        The matrix function λ→A(λ).

    """

    def inner(
        eigval: float, derv: int
    ) -> np.ndarray[tuple[int, int], np.dtype[np.number]] | None:
        return np.asarray(A.diff("x", derv).subs("x", eigval), dtype=float)

    return inner


def _jordan_chain_def_right_term(
    A: np.ndarray[tuple[int, int, int], np.dtype[np.number]],
    chains: np.ndarray[tuple[int, int, int], np.dtype[np.number]],
) -> np.ndarray[tuple[int, int], np.dtype[np.number]]:
    """
    Get the right term of the Jordan chain definition.

    Parameters
    ----------
    A : np.ndarray[tuple[int, int, int], np.dtype[np.number]]
        The matrix derivatives evaluated at the eigenvalue
        of shape (l_chain, n, n).
    chains : np.ndarray[tuple[int, int, int], np.dtype[np.number]]
        The Jordan chains of shape (n_chain, l_chain, n).

    Returns
    -------
    Tuple[np.ndarray[tuple[int, int], np.dtype[np.number]]]
        The right term of the Jordan chain definition
        of shape (n, n).

    """
    if (
        A.ndim == 3
        and chains.ndim == 3
        and (A.shape[0] == chains.shape[1])
        and (A.shape[1] == A.shape[2] == chains.shape[2])
    ):
        pass
    else:
        raise ValueError(
            f"Invalid shape {A.shape=}, {chains.shape=}. "
            "A should be (l_chain, n, n) and chains should be (n_chain, l_chain, n)."
        )
    return -np.sum(
        np.matmul(A[None, ...], np.flip(chains[..., None], axis=1))[..., 0]
        / scipy.special.factorial(np.arange(A.shape[0])[:, None, None]),
        axis=1,
    )


def unrestricted_jordan_chains(
    A_func: MatrixFuncProtocol,
    eigval: float,
    /,
    atol_rank: float | None = None,
    rtol_rank: float | None = None,
    atol_norm: float | None = None,
    rtol_norm: float | None = None,
    max_l_chain: int | None = None,
) -> list[np.ndarray[tuple[int, int, int], np.dtype[np.number]]]:
    """
    Get the unrestricted Jordan chains of the matrix of fixed length.

    Parameters
    ----------
    A_func : MatrixFuncProtocol
        The matrix function which takes (λ, d) as input
        and returns the d-th derivative of the matrix function
        evaluated at the eigenvalue λ.
    eigval : float
        The eigenvalue to evaluate the matrix function.
    atol_rank : float, optional
        Threshold below which SVD values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    rtol_rank : float, optional
        Threshold below which SVD values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    atol_norm : float, optional
        Threshold below which norm values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    rtol_norm : float, optional
        Threshold below which norm values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    max_l_chain : int, optional
        The maximum length of the Jordan chain to search for.
        If None, stop searching when all unrestricted Jordan chains
        are not Jordan chains, in other words,
        the first eigenvector's norm is almost 0,
        because in general there are infinite number of Jordan chains.

    Returns
    -------
    np.ndarray[tuple[int, int, int], np.dtype[np.number]]
        The list which contains
        unrestricted Jordan chains of shape (n_chain, l_chain, n)
        at the (l_chain-1)th element.

    """
    As_list: list[np.ndarray[tuple[int, int], np.dtype[np.number]]] = []
    chains = []
    i = 0
    while True:
        A_func_res = A_func(eigval, len(As_list))
        if A_func_res is None:
            break
        As_list.append(A_func_res)
        As = np.stack(As_list, axis=0)
        l_chain = As.shape[0]
        n = As.shape[1]
        mat = np.stack(
            [
                np.moveaxis(
                    np.concat(
                        (
                            np.flip(
                                As[: (j + 1), :, :]
                                / scipy.special.factorial(
                                    np.arange(j + 1)[:, None, None]
                                ),
                                axis=0,
                            ),
                            np.zeros(
                                (l_chain - j - 1, n, n),
                                dtype=As.dtype,
                                device=As.device,
                            ),
                        ),
                        axis=0,
                    ),
                    0,
                    1,
                ).reshape(n, l_chain * n)
                for j in range(l_chain)
            ],
            axis=0,
        ).reshape(l_chain * n, l_chain * n)
        # (m*n, n_jordan_chain)
        chain = scipy.linalg.null_space(mat, rtol_rank)
        if chain.size == 0:
            break
        # (n_jordan_chain, m*n)
        chain = np.moveaxis(chain, -1, 0)
        # (n_jordan_chain, m, n)
        chain = chain.reshape(chain.shape[0], l_chain, n)
        i += 1
        if max_l_chain is None:
            norm = np.linalg.norm(chain[:, 0, :], axis=-1)
            norm_filter = norm > get_tol(np.max(norm), rtol=rtol_norm, atol=atol_norm)
            if not np.any(norm_filter):
                break
        elif i > max_l_chain:
            break
        chains.append(chain)
    return chains


def _get_space(
    chains: list[NDArray[np.number]],
    length: int,
    /,
) -> NDArray[np.number]:
    """
    Get the cut chains.

    Parameters
    ----------
    chains : list[NDArray[np.number]]
        The Jordan chains.
    length : int
        The length to cut.

    Returns
    -------
    NDArray[np.number]
        The cut chains of shape (n_chain, l_chain, n).

    """
    res = np.concat([chains_[:, :length, :] for chains_ in chains], axis=0)
    return res


@overload
def canonoical_jordan_chains_from_unrestricted(
    unrestricted_chains: list[np.ndarray[tuple[int, int, int], np.dtype[np.number]]],
    /,
    *,
    hermitian: bool | None = ...,
    atol_rank: float | None = ...,
    rtol_rank: float | None = ...,
    atol_norm: float | None = ...,
    rtol_norm: float | None = ...,
    flatten: Literal[False] = ...,
) -> list[np.ndarray[tuple[int, int, int], np.dtype[np.number]]]: ...
@overload
def canonoical_jordan_chains_from_unrestricted(
    unrestricted_chains: list[np.ndarray[tuple[int, int, int], np.dtype[np.number]]],
    /,
    *,
    hermitian: bool | None = ...,
    atol_rank: float | None = ...,
    rtol_rank: float | None = ...,
    atol_norm: float | None = ...,
    rtol_norm: float | None = ...,
    flatten: Literal[True] = ...,
) -> list[np.ndarray[tuple[int, int], np.dtype[np.number]]]: ...
def canonoical_jordan_chains_from_unrestricted(
    unrestricted_chains: list[np.ndarray[tuple[int, int, int], np.dtype[np.number]]],
    /,
    *,
    hermitian: bool | None = None,
    atol_rank: float | None = None,
    rtol_rank: float | None = None,
    atol_norm: float | None = None,
    rtol_norm: float | None = None,
    flatten: bool = True,
) -> (
    list[np.ndarray[tuple[int, int], np.dtype[np.number]]]
    | list[np.ndarray[tuple[int, int, int], np.dtype[np.number]]]
):
    """
    Get the Jordan chains of the matrix.

    Parameters
    ----------
    unrestricted_chains : np.ndarray[tuple[int, int, int], np.dtype[np.number]]
        The list of unrestricted Jordan chains of shape (n_chain, l_chain, n).
    hermitian : bool, optional
        If True, `A` is assumed to be Hermitian (symmetric if real-valued),
        enabling a more efficient method for finding singular values.
        Defaults to False.
    atol_rank : float, optional
        Threshold below which SVD values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    rtol_rank : float, optional
        Threshold below which SVD values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    atol_norm : float, optional
        Threshold below which norm values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    rtol_norm : float, optional
        Threshold below which norm values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    flatten : bool, optional
        If True, flatten the chains. Defaults to True.

    Returns
    -------
    list[NDArray[np.number]]
        The Jordan chains.
        If flatten is True, each chain is of shape (l_chain, n).
        If flatten is False, each chain is of shape (n_chain_l, l_chain, n).

    """
    chains: list[NDArray[np.number]] = []
    for chain in reversed(unrestricted_chains):
        # filter chains which first eigenvector's norm is too small
        norm = np.linalg.norm(chain[:, 0, :], axis=-1)
        norm_filter = norm > get_tol(np.max(norm), rtol=rtol_norm, atol=atol_norm)
        chain = chain[norm_filter, :, :]

        # do not process empty chains
        if not chain.size:
            continue

        # remove chains overlapping from already found longer chains
        if chains:
            # chains are already normalized
            cut_chain = _get_space(chains, chain.shape[1])
            # [n_chain_cut, l_chain, n]
            # [n_chain, n], [1, n_chain_cut, n] -> [n_chain, n_chain_cut]
            cut_chain_ip = np.sum(
                chain[:, None, 0, :] * cut_chain[None, :, 0, :], axis=-1
            )
            # [n_chain, l_chain, n]
            chain = chain - np.sum(
                cut_chain_ip[:, :, None, None] * cut_chain[None, :, :, :], axis=1
            )

        # remove duplicated chains
        chain = np.swapaxes(chain, 0, 1)
        u, s, _ = np.linalg.svd(chain[0, :, :], hermitian=hermitian)
        rank = _matrix_rank_from_s(chain[0, :, :], s, atol=atol_rank, rtol=rtol_rank)
        chain = u.T[:rank] @ chain
        chain = np.swapaxes(chain, 0, 1)

        # normalize chains based on the first eigenvector
        norm = np.linalg.norm(chain[:, 0, :], axis=-1)
        chain = chain / norm[:, None, None]

        # append the chain
        if chain.size:
            chains.append(chain)
    if flatten:
        chains = [c for chain_ in chains for c in chain_]
    return chains


def canonical_jordan_chains(
    A_func: MatrixFuncProtocol,
    eigval: float,
    /,
    atol_rank: float | None = None,
    rtol_rank: float | None = None,
    atol_norm: float | None = None,
    rtol_norm: float | None = None,
    algebraic_multiplicity: int | None = None,
) -> CanonicalJordanChains:
    """
    Get the canonical Jordan chains of the matrix function.

    Parameters
    ----------
    A_func : MatrixFuncProtocol
        The matrix function which takes (λ, d) as input
        and returns the d-th derivative of the matrix function
        evaluated at the eigenvalue λ.
    eigval : float
        The eigenvalue to evaluate the matrix function.
    atol_rank : float, optional
        Threshold below which SVD values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    rtol_rank : float, optional
        Threshold below which SVD values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    atol_norm : float, optional
        Threshold below which norm values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    rtol_norm : float, optional
        Threshold below which norm values are considered zero.
        Defaults to ``np.finfo(A.dtype).eps``.
    algebraic_multiplicity : int, optional
        The algebraic multiplicity of the eigenvalue.
        If provided, a warning is issued
        if the number of generalized eigenvectors found
        does not match the algebraic multiplicity.

    Returns
    -------
    JordanChains
        The canonical Jordan chains of the matrix function.

    """
    chains = canonoical_jordan_chains_from_unrestricted(
        unrestricted_jordan_chains(
            A_func,
            eigval,
            rtol_rank=rtol_rank,
            atol_rank=atol_rank,
            atol_norm=atol_norm,
            rtol_norm=rtol_norm,
        ),
        rtol_rank=rtol_rank,
        atol_rank=atol_rank,
        rtol_norm=rtol_norm,
        atol_norm=atol_norm,
        flatten=True,
    )
    n_generalized_eigenvectors = np.sum([c.shape[0] for c in chains])
    if (
        algebraic_multiplicity is not None
        and n_generalized_eigenvectors != algebraic_multiplicity
    ):
        warnings.warn(
            "The number of generalized eigenvectors found "
            "does not match the provided algebraic multiplicity. "
            "Consider using a larger value for `atol` or `rtol`.",
            RuntimeWarning,
            stacklevel=2,
        )
    return CanonicalJordanChains(
        eigval=eigval,
        chains=chains,
    )
