/**
 * @file PyMatrix.h
 * @brief Defines a pybind11 trampoline class for mfem::Matrix.
 *
 * This trampoline class allows Python classes to inherit from mfem::Matrix,
 * enabling Python-defined matrices to be used within the C++ MFEM library,
 * including overriding methods from its base class mfem::Operator.
 */
#pragma once

#include "mfem.hpp"
#include "pybind11/pybind11.h"

namespace serif::pybind {

    /**
     * @brief A trampoline class for mfem::Matrix.
     *
     * This class allows Python classes to inherit from mfem::Matrix and correctly
     * override its virtual functions, including those inherited from its base,
     * mfem::Operator. This is useful for creating custom matrix types in Python
     * that can interact seamlessly with MFEM's C++ components.
     *
     * @see mfem::Matrix
     * @see mfem::Operator
     * @see PyOperator
     *
     * @par Python Usage Example:
     * @code{.py}
     * import mfem.ser_ext as mfem
     * import numpy as np
     *
     * class MyPythonMatrix(mfem.Matrix):
     *     def __init__(self, size=0):
     *         super().__init__(size) # Call base C++ constructor
     *         if size > 0:
     *             # Example: Initialize with some data if size is provided
     *             # Note: Direct data manipulation might be complex due to
     *             # C++ ownership. This example is conceptual.
     *             # For a real-world scenario, you might manage data in Python
     *             # and use Eval or Mult to reflect its state.
     *             self._py_data = np.zeros((size, size))
     *         else:
     *             self._py_data = None
     *
     *     # --- mfem.Matrix pure virtual overrides ---
     *     def Elem(self, i, j):
     *         # This is problematic for direct override for write access
     *         # due to returning a reference.
     *         # Typically, you'd manage data and use it in Mult/GetRow etc.
     *         # For read access (const version), it's more straightforward.
     *         if self._py_data is not None:
     *             return self._py_data[i, j]
     *         raise IndexError("Matrix data not initialized or index out of bounds")
     *
     *     # const Elem version for read access
     *     # def GetElem(self, i, j): # A helper in Python if Elem is tricky
     *     #    return self._py_data[i,j]
     *
     *     def Inverse(self):
     *         # Return an mfem.MatrixInverse object
     *         # This is a simplified example; a real inverse is complex.
     *         print("MyPythonMatrix.Inverse() called, returning dummy inverse.")
     *         # For a real implementation, you'd compute and return an actual
     *         # mfem.MatrixInverse or a Python-derived version of it.
     *         # This might involve creating a new PyMatrix representing the inverse.
     *         identity_inv = mfem.DenseMatrix(self.Width())
     *         for i in range(self.Width()):
     *            identity_inv[i,i] = 1.0
     *         return mfem.MatrixInverse(identity_inv) # Example
     *
     *     # --- mfem.Matrix regular virtual overrides ---
     *     def Finalize(self, skip_zeros=1):
     *         super().Finalize(skip_zeros) # Call base
     *         print(f"MyPythonMatrix.Finalize({skip_zeros}) called.")
     *
     *     # --- mfem.Operator virtual overrides ---
     *     def Mult(self, x, y):
     *         # x is mfem.Vector (input), y is mfem.Vector (output)
     *         if self._py_data is not None:
     *             x_np = x.GetDataArray()
     *             y_np = np.dot(self._py_data, x_np)
     *             y.SetSize(self._py_data.shape[0])
     *             y.Assign(y_np)
     *         else:
     *             # Fallback to base class if no Python data
     *             # super().Mult(x,y) # This might not work as expected if base is abstract
     *             # or if the C++ mfem.Matrix itself doesn't have data.
     *             # For a sparse matrix, this would call its sparse Mult.
     *             # For a dense matrix, it would use its data.
     *             # If this PyMatrix is purely Python defined, it must implement Mult.
     *             raise NotImplementedError("Mult not implemented or data not set")
     *         print("MyPythonMatrix.Mult called.")
     *
     * # Using the Python Matrix
     * mat_size = 3
     * py_mat = MyPythonMatrix(mat_size)
     * py_mat._py_data = np.array([[1,2,3],[4,5,6],[7,8,9]], dtype=float)
     *
     * x_vec = mfem.Vector(mat_size)
     * x_vec.Assign(np.array([1,1,1], dtype=float))
     * y_vec = mfem.Vector(mat_size)
     *
     * py_mat.Mult(x_vec, y_vec)
     * print("Result y:", y_vec.GetDataArray())
     *
     * # inv_op = py_mat.Inverse() # Be cautious with dummy implementations
     * # print("Inverse type:", type(inv_op))
     *
     * # Note: Overriding Elem(i,j) to return a C++ reference from Python
     * # for write access (mfem::real_t&) is complex and often not directly feasible
     * # in a straightforward way with pybind11 for typical Python data structures.
     * # Python users would typically interact via methods like Set, Add, or by
     * # providing data that the C++ side can access (e.g., via GetData).
     * # The const version of Elem (read-only) is more manageable.
     * @endcode
     */
    class PyMatrix : public mfem::Matrix {
    public:
        // Inherit constructors from the base mfem::Matrix class.
        using mfem::Matrix::Matrix;

        // --- Trampolines for new mfem::Matrix pure virtual methods ---
        /**
         * @brief Access element (i,j) for read/write.
         * Pure virtual in mfem::Matrix. Must be overridden in Python.
         * @param i Row index.
         * @param j Column index.
         * @return Reference to the matrix element.
         * @note Returning a C++ reference from Python for write access can be complex.
         *       Consider alternative data management strategies in Python.
         *       PYBIND11_OVERRIDE_PURE is used in the .cpp file.
         */
        mfem::real_t& Elem(int i, int j) override;

        /**
         * @brief Access element (i,j) for read-only.
         * Pure virtual in mfem::Matrix. Must be overridden in Python.
         * @param i Row index.
         * @param j Column index.
         * @return Const reference to the matrix element.
         * @note PYBIND11_OVERRIDE_PURE is used in the .cpp file.
         */
        const mfem::real_t& Elem(int i, int j) const override;

        /**
         * @brief Get the inverse of the matrix.
         * Pure virtual in mfem::Matrix. Must be overridden in Python.
         * The caller is responsible for deleting the returned MatrixInverse object.
         * @return Pointer to an mfem::MatrixInverse object.
         * @note PYBIND11_OVERRIDE_PURE is used in the .cpp file.
         */
        mfem::MatrixInverse* Inverse() const override;

        // --- Trampoline for new mfem::Matrix regular virtual methods ---
        /**
         * @brief Finalize matrix assembly.
         * For sparse matrices, this typically involves finalizing the sparse structure.
         * Can be overridden in Python.
         * @param skip_zeros See mfem::SparseMatrix::Finalize documentation.
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void Finalize(int skip_zeros) override;

        // --- Trampolines for inherited mfem::Operator virtual methods ---
        // These must be repeated here to allow Python classes inheriting from
        // Matrix to override methods originally from the Operator base class.

        /**
         * @brief Perform the operator action: y = A*x.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * If not overridden, mfem::Matrix's default implementation is used.
         * @param x The input vector.
         * @param y The output vector (result of A*x).
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void Mult(const mfem::Vector &x, mfem::Vector &y) const override;

        /**
         * @brief Perform the transpose operator action: y = A^T*x.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @param x The input vector.
         * @param y The output vector (result of A^T*x).
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void MultTranspose(const mfem::Vector &x, mfem::Vector &y) const override;

        /**
         * @brief Perform the action y += a*(A*x).
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @param x The input vector.
         * @param y The vector to which a*(A*x) is added.
         * @param a Scalar multiplier (defaults to 1.0).
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void AddMult(const mfem::Vector &x, mfem::Vector &y, const mfem::real_t a = 1.0) const override;

        /**
         * @brief Perform the action y += a*(A^T*x).
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @param x The input vector.
         * @param y The vector to which a*(A^T*x) is added.
         * @param a Scalar multiplier (defaults to 1.0).
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void AddMultTranspose(const mfem::Vector &x, mfem::Vector &y, const mfem::real_t a = 1.0) const override;

        /**
         * @brief Get the gradient operator (Jacobian) at a given point x.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * For a linear matrix operator, the gradient is typically the matrix itself.
         * @param x The point at which to evaluate the gradient (often unused for linear operators).
         * @return A reference to the gradient operator.
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        mfem::Operator& GetGradient(const mfem::Vector &x) const override;

        /**
         * @brief Assemble the diagonal of the operator.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @param diag Output vector to store the diagonal entries.
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void AssembleDiagonal(mfem::Vector &diag) const override;

        /**
         * @brief Get the prolongation operator.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @return A const pointer to the prolongation operator, or nullptr if not applicable.
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        const mfem::Operator* GetProlongation() const override;

        /**
         * @brief Get the restriction operator.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @return A const pointer to the restriction operator, or nullptr if not applicable.
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        const mfem::Operator* GetRestriction() const override;

        /**
         * @brief Recover the FEM solution.
         * Inherited from mfem::Operator. Can be overridden in Python.
         * @param X The reduced solution vector.
         * @param b The right-hand side vector.
         * @param x Output vector for the full FEM solution.
         * @note PYBIND11_OVERRIDE is used in the .cpp file.
         */
        void RecoverFEMSolution(const mfem::Vector &X, const mfem::Vector &b, mfem::Vector &x) override;
    };

} // namespace serif::pybind