/**
 * @file PyOperator.h
 * @brief Defines a pybind11 trampoline class for mfem::Operator.
 *
 * This trampoline class allows Python classes to inherit from mfem::Operator,
 * enabling Python-defined operators to be used within the C++ MFEM library.
 */
#pragma once

#include "mfem.hpp"
#include "pybind11/pybind11.h"
#include "pybind11/stl.h" // Needed for vectors, maps, sets, strings

namespace py = pybind11;

/**
 * @namespace serif::pybind
 * @brief Contains pybind11 helper classes and trampoline classes for interfacing C++ with Python.
 */
namespace serif::pybind {

    /**
     * @brief Trampoline class for mfem::Operator.
     *
     * This class allows Python classes to inherit from mfem::Operator and correctly
     * override its virtual functions. When a virtual function is called from C++,
     * the trampoline ensures the call is forwarded to the Python implementation if one exists.
     * This is crucial for integrating Python-defined linear or non-linear operators
     * into MFEM's C++-based solvers and algorithms.
     *
     * @see mfem::Operator
     *
     * @par Python Usage Example:
     * @code{.py}
     * import mfem.ser_ext as mfem
     * import numpy as np
     *
     * class MyPythonOperator(mfem.Operator):
     *     def __init__(self, size):
     *         super().__init__(size) # Call the base C++ constructor
     *         # Or super().__init__() if default constructing and setting height/width later
     *         self.matrix = np.random.rand(size, size) # Example: store a dense matrix
     *
     *     # Must override Mult
     *     def Mult(self, x, y):
     *         # x is an mfem.Vector (input)
     *         # y is an mfem.Vector (output, y = A*x)
     *         # Ensure y is correctly sized if not already
     *         if y.Size() != self.Height():
     *             y.SetSize(self.Height())
     *
     *         # Example: y = self.matrix * x
     *         # This is a conceptual illustration. For actual matrix-vector products
     *         # with numpy, you'd convert mfem.Vector to numpy array or iterate.
     *         x_np = x.GetDataArray() # Get a numpy view if configured
     *         y_np = np.dot(self.matrix, x_np)
     *         y.Assign(y_np) # Assign result back to mfem.Vector
     *         print("MyPythonOperator.Mult called")
     *
     *     # Optionally override other methods like MultTranspose, GetGradient, etc.
     *     def MultTranspose(self, x, y):
     *         if y.Size() != self.Width():
     *             y.SetSize(self.Width())
     *         # Example: y = self.matrix.T * x
     *         x_np = x.GetDataArray()
     *         y_np = np.dot(self.matrix.T, x_np)
     *         y.Assign(y_np)
     *         print("MyPythonOperator.MultTranspose called")
     *
     * # Using the Python operator
     * op_size = 5
     * py_op = MyPythonOperator(op_size)
     *
     * x_vec = mfem.Vector(op_size)
     * x_vec.Assign(np.arange(op_size, dtype=float))
     * y_vec = mfem.Vector(op_size)
     *
     * py_op.Mult(x_vec, y_vec)
     * print("Result y:", y_vec.GetDataArray())
     * # py_op can now be passed to MFEM functions expecting an mfem::Operator
     * @endcode
     */
    class PyOperator : public mfem::Operator {
    public:
        // Inherit the constructors from the base mfem::Operator class.
        // This allows Python classes to call e.g., super().__init__(size).
        using mfem::Operator::Operator;

        // --- Trampoline declarations for all overridable virtual functions ---

        /**
         * @brief Perform the operator action: y = A*x.
         *
         * This is a pure virtual function in mfem::Operator and **must** be overridden
         * by any Python class inheriting from PyOperator.
         *
         * @param x The input vector.
         * @param y The output vector (result of A*x).
         * @note This method forwards the call to the Python override.
         *       PYBIND11_OVERRIDE_PURE is used in the .cpp file to handle this.
         */
        void Mult(const mfem::Vector &x, mfem::Vector &y) const override;

        /**
         * @brief Perform the transpose operator action: y = A^T*x.
         *
         * Optional override. If not overridden, MFEM's base implementation
         * (which may raise an error or be a no-op) will be used.
         *
         * @param x The input vector.
         * @param y The output vector (result of A^T*x).
         * @note This method forwards the call to the Python override if one exists.
         *       PYBIND11_OVERRIDE is used in the .cpp file to handle this.
         */
        void MultTranspose(const mfem::Vector &x, mfem::Vector &y) const override;

        /**
         * @brief Perform the action y += a*(A*x).
         *
         * Optional override.
         *
         * @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 This method forwards the call to the Python override if one exists.
         */
        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).
         *
         * Optional override.
         *
         * @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 This method forwards the call to the Python override if one exists.
         */
        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.
         *
         * For non-linear operators, this method should return the linearization (Jacobian)
         * of the operator at the point `x`. The returned Operator is owned by this
         * Operator and should not be deleted by the caller.
         * Optional override.
         *
         * @param x The point at which to evaluate the gradient.
         * @return A reference to the gradient operator.
         * @note This method forwards the call to the Python override if one exists.
         */
        Operator& GetGradient(const mfem::Vector &x) const override;

        /**
         * @brief Assemble the diagonal of the operator.
         *
         * For discrete operators (e.g., matrices), this method should compute and store
         * the diagonal entries of the operator in the vector `diag`.
         * Optional override.
         *
         * @param diag Output vector to store the diagonal entries.
         * @note This method forwards the call to the Python override if one exists.
         */
        void AssembleDiagonal(mfem::Vector &diag) const override;

        /**
         * @brief Get the prolongation operator.
         *
         * Used in multilevel methods (e.g., AMG). Returns a pointer to the prolongation
         * operator (interpolation from a coarser level to this level).
         * The returned Operator is typically owned by this Operator.
         * Optional override.
         *
         * @return A const pointer to the prolongation operator, or nullptr if not applicable.
         * @note This method forwards the call to the Python override if one exists.
         */
        const mfem::Operator* GetProlongation() const override;

        /**
         * @brief Get the restriction operator.
         *
         * Used in multilevel methods (e.g., AMG). Returns a pointer to the restriction
         * operator (projection from this level to a coarser level).
         * Typically, this is the transpose of the prolongation operator.
         * The returned Operator is typically owned by this Operator.
         * Optional override.
         *
         * @return A const pointer to the restriction operator, or nullptr if not applicable.
         * @note This method forwards the call to the Python override if one exists.
         */
        const mfem::Operator* GetRestriction() const override;

        /**
         * @brief Recover the FEM solution.
         *
         * For operators that are part of a system solve (e.g., static condensation),
         * this method can be used to reconstruct the full finite element solution `x`
         * from a reduced solution `X` and the right-hand side `b`.
         * Optional override.
         *
         * @param X The reduced solution vector.
         * @param b The right-hand side vector.
         * @param x Output vector for the full FEM solution.
         * @note This method forwards the call to the Python override if one exists.
         */
        void RecoverFEMSolution(const mfem::Vector &X, const mfem::Vector &b, mfem::Vector &x) override;
    };

} // namespace serif::pybind