# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________

from abc import ABCMeta, abstractmethod
import enum
from typing import Optional, Union, Tuple
from scipy.sparse import spmatrix
import numpy as np
from pyomo.contrib.pynumero.sparse.base_block import BaseBlockVector, BaseBlockMatrix
from pyomo.contrib.pynumero.sparse import BlockVector, BlockMatrix


class LinearSolverStatus(enum.Enum):
    successful = 0
    not_enough_memory = 1
    singular = 2
    error = 3
    warning = 4
    max_iter = 5


class LinearSolverResults:
    def __init__(self, status: Optional[LinearSolverStatus] = None):
        self.status = status


class LinearSolverInterface(object, metaclass=ABCMeta):
    @abstractmethod
    def solve(
        self,
        matrix: Union[spmatrix, BlockMatrix],
        rhs: Union[np.ndarray, BlockVector],
        raise_on_error: bool = True,
    ) -> Tuple[Optional[Union[np.ndarray, BlockVector]], LinearSolverResults]:
        pass


class DirectLinearSolverInterface(LinearSolverInterface, metaclass=ABCMeta):
    @abstractmethod
    def do_symbolic_factorization(
        self, matrix: Union[spmatrix, BlockMatrix], raise_on_error: bool = True
    ) -> LinearSolverResults:
        pass

    @abstractmethod
    def do_numeric_factorization(
        self, matrix: Union[spmatrix, BlockMatrix], raise_on_error: bool = True
    ) -> LinearSolverResults:
        pass

    @abstractmethod
    def do_back_solve(
        self, rhs: Union[np.ndarray, BlockVector], raise_on_error: bool = True
    ) -> Tuple[Optional[Union[np.ndarray, BlockVector]], LinearSolverResults]:
        pass

    def solve(
        self,
        matrix: Union[spmatrix, BlockMatrix],
        rhs: Union[np.ndarray, BlockVector],
        raise_on_error: bool = True,
    ) -> Tuple[Optional[Union[np.ndarray, BlockVector]], LinearSolverResults]:
        symbolic_res = self.do_symbolic_factorization(
            matrix, raise_on_error=raise_on_error
        )
        if symbolic_res.status != LinearSolverStatus.successful:
            return None, symbolic_res
        numeric_res = self.do_numeric_factorization(
            matrix, raise_on_error=raise_on_error
        )
        if numeric_res.status != LinearSolverStatus.successful:
            return None, numeric_res
        return self.do_back_solve(rhs, raise_on_error=raise_on_error)
