"""Extended Krylov subspace construction.

Author: Benjamin Carrel, University of Geneva, 2022-2023
"""

# %% Imports
from __future__ import annotations

from typing import Callable, Optional

import numpy as np
import scipy.linalg as la
import scipy.sparse.linalg as spsla
from numpy import ndarray
from scipy.sparse import spmatrix

from .inverted_krylov_space import InvertedKrylovSpace
from .krylov_space import KrylovSpace
from .space_structure import SpaceStructure


# %% Class definition
class ExtendedKrylovSpace(SpaceStructure):
    """Extended Krylov space.

    Constructs the extended Krylov space by combining a standard Krylov space
    and an inverted Krylov space:
        EK_m(A, X) = span{X, AX, A^2X, ..., A^(m-1)X, A^(-1)X, A^(-2)X, ..., A^(-m)X}

    This space is particularly useful for problems requiring information from both
    the range of A and its inverse.

    **Algorithm Selection:** When is_symmetric=True, both component Krylov spaces
    (standard and inverted) use the Lanczos algorithm, providing approximately 2x
    speedup over the non-symmetric case. The combined basis maintains orthogonality
    across both components.

    How to use
    ----------
    1. Create an instance: EK = ExtendedKrylovSpace(A, X, invA=None)
    2. Augment the basis as needed: EK.augment_basis()
    3. Access the basis via EK.basis or EK.Q

    Attributes
    ----------
    A : spmatrix
        Sparse matrix of shape (n, n)
    X : ndarray
        Initial vector or matrix of shape (n, r)
    invA : callable
        Function that computes A^(-1) * v for a given vector/matrix v
    krylov_space : KrylovSpace
        The standard Krylov space component
    inverted_krylov_space : InvertedKrylovSpace
        The inverted Krylov space component
    Q : ndarray
        Combined orthonormal basis of the extended Krylov space
    Q1 : ndarray
        Basis of the standard Krylov space
    Q2 : ndarray
        Basis of the inverted Krylov space
    basis : ndarray
        Pointer to Q

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.sparse import csr_matrix
    >>> from lowrank.krylov import ExtendedKrylovSpace
    >>> # Create a sparse matrix
    >>> A = csr_matrix([[4, 1], [1, 3]])
    >>> X = np.array([[1.0], [0.0]])
    >>> # Extended Krylov includes both A^k*X and A^(-k)*X
    >>> EK = ExtendedKrylovSpace(A, X)
    >>> EK.augment_basis()  # Adds both A*X and A^(-1)*X
    >>> EK.size  # Total size is sum of both components
    4
    >>> # Access individual components
    >>> EK.Q1.shape  # Standard Krylov part
    (2, 2)
    >>> EK.Q2.shape  # Inverted Krylov part
    (2, 2)
    >>> # Combined basis is orthonormal
    >>> np.allclose(EK.Q.T @ EK.Q, np.eye(EK.Q.shape[1]))
    True
    """

    def __init__(
        self, A: spmatrix, X: ndarray, invA: Optional[Callable] = None, **extra_args
    ):
        """
        Parameters
        ----------
        A : spmatrix
            The matrix A of the linear system.
        X : ndarray
            The basis of the Krylov space.
        invA: callable
            The function that computes the action of A^(-1) on a vector, or a matrix.
        """
        # Validate inputs first (before accessing any attributes)
        if not isinstance(X, ndarray):
            raise TypeError("X must be a numpy array")

        # Ensure X is 2D before calling parent constructor
        if X.ndim == 1:
            X = X.reshape(-1, 1)

        # Call parent constructor
        super().__init__(A, X, **extra_args)

        # Check specific inputs
        assert invA is None or callable(invA), "invA must be a function."
        if invA is None:
            spluA = spsla.splu(A)
            invA = spluA.solve

        # Store specific inputs
        self.invA = invA

        # Krylov space
        self.krylov_space = KrylovSpace(A, X, is_symmetric=self.is_symmetric)

        # Inverted Krylov space
        self.inverted_krylov_space = InvertedKrylovSpace(
            A, X, invA=invA, is_symmetric=self.is_symmetric
        )

        # Cache for the combined basis
        self._Q_cache: Optional[ndarray] = None
        self._cache_size = 0

    # %% Properties
    @property
    def Q1(self) -> ndarray:
        """Basis of the standard Krylov space component.

        Returns
        -------
        ndarray
            Matrix of shape (n, m*r) containing the Krylov basis vectors.
        """
        return self.krylov_space.Q

    @property
    def H1(self) -> ndarray:
        """Upper Hessenberg matrix from Arnoldi for the Krylov space.

        Returns
        -------
        ndarray
            The Hessenberg matrix (non-symmetric case only).
        """
        return self.krylov_space.H

    @property
    def Q2(self) -> ndarray:
        """Basis of the inverted Krylov space component.

        Returns
        -------
        ndarray
            Matrix of shape (n, m*r) containing the inverted Krylov basis vectors.
        """
        return self.inverted_krylov_space.Q

    @property
    def H2(self) -> ndarray:
        """Upper Hessenberg matrix from Arnoldi for the inverted Krylov space.

        Returns
        -------
        ndarray
            The Hessenberg matrix (non-symmetric case only).
        """
        return self.inverted_krylov_space.H

    @property
    def Q(self) -> ndarray:
        """Combined orthonormal basis of the extended Krylov space.

        Concatenates Q1 and Q2, then orthonormalizes the result. The result is
        cached and only recomputed when the space size changes.

        Returns
        -------
        ndarray
            Matrix of shape (n, size) containing the orthonormal basis.
        """
        # Check if cache is valid
        current_size = self.size
        if self._Q_cache is None or self._cache_size != current_size:
            # Recompute and cache
            self._Q_cache = la.qr(np.hstack((self.Q1, self.Q2)), mode="economic")[0]
            self._cache_size = current_size
        assert self._Q_cache is not None, "Q cache should be initialized"
        return self._Q_cache

    @property
    def basis(self) -> ndarray:
        """The orthonormal basis of the extended Krylov space.

        Returns
        -------
        ndarray
            Matrix of shape (n, size) (same as Q).
        """
        return self.Q

    @property
    def size(self) -> int:
        """Total size of the extended Krylov space.

        Returns
        -------
        int
            Sum of the sizes of the Krylov and inverted Krylov components.
        """
        return self.krylov_space.size + self.inverted_krylov_space.size

    # %% Methods
    def augment_basis(self):
        """Augment the basis of the extended Krylov space.

        Augments both the standard Krylov space and the inverted Krylov space
        by adding the next block of basis vectors to each component.
        """
        self.krylov_space.augment_basis()
        self.inverted_krylov_space.augment_basis()
