module;

#include <algorithm>
#include <cmath>
#include <concepts>
#include <memory>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

export module mean_field:solver.stellar_enthalpy_strategy;
export import :solver.enthalpy_admissibility;
export import :normalization.stellar_equilibrium;
export import :utils.domain;

export namespace mean_field::solver::detail {
    template <class SurfaceConstraint>
    concept PolynomialEnthalpySurfaceConstraint = requires {
        typename SurfaceConstraint::CarrierQuantity;
        typename SurfaceConstraint::Relation;
        requires std::same_as<typename SurfaceConstraint::CarrierQuantity,dimensions::quantity::SpecificEnthalpy>;
        requires std::same_as<typename SurfaceConstraint::Relation,eos::SpecificEnthalpyFromPressure>;
    };

    // The polynomial engine fixes homogeneous trace rows. Check the physical
    // prescription before inspecting coefficients, so an unsupported surface
    // condition is never reported as an invalid numerical starting state.
    template <PolynomialEnthalpySurfaceConstraint SurfaceConstraint, class EquationOfState>
    void RequirePolynomialEnthalpySurface(const SurfaceConstraint &surface,
                                         const EquationOfState &equationOfState, MPI_Comm communicator) {
        int localUnsupported = surface.targetPressure().value() == 0.0 ? 0 : 1;
        int globalUnsupported = 0;
        if (MPI_Allreduce(&localUnsupported,&globalUnsupported,1,MPI_INT,MPI_MAX,communicator) != MPI_SUCCESS)
            throw std::runtime_error("Could not validate the polynomial enthalpy surface prescription.");
        if (globalUnsupported)
            throw std::invalid_argument("PolynomialEnthalpy supports only zero prescribed surface pressure (Psurf=0).");
        bool homogeneous = false;
        try {
            const double trace = eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
                equationOfState,surface.targetPressure()).value();
            homogeneous = std::isfinite(trace) && trace == 0.0;
        } catch (...) {
            // A rank-local EOS exception must not strand other ranks in the
            // collective compatibility check that follows.
        }
        localUnsupported = homogeneous ? 0 : 1;
        if (MPI_Allreduce(&localUnsupported,&globalUnsupported,1,MPI_INT,MPI_MAX,communicator) != MPI_SUCCESS)
            throw std::runtime_error("Could not validate the zero-pressure enthalpy reference.");
        if (globalUnsupported)
            throw std::invalid_argument("PolynomialEnthalpy requires the EOS to prescribe h(Psurf=0)=0.");
    }
}

export namespace mean_field::solver {
    /*
     * The solver owns this workspace only when a positivity policy is selected.
     * Certification uses physical h: arbitrary positive diagonal normalization
     * need not preserve the polynomial represented by its normalized entries.
     */
    template <class State>
    class PreparedStellarEnthalpyStrategy final {
        using Problem = std::remove_cvref_t<decltype(std::declval<State &>().Problem())>;
        using Form = typename Problem::FormType;
        static_assert(detail::PolynomialEnthalpySurfaceConstraint<typename Problem::CompiledSurfaceConstraintType>,
                      "PolynomialEnthalpy requires a state-independent specific-enthalpy pressure surface relation.");
        using Schema = utils::domain::CoreEnvelopeVacuumDomainSchema;
        using Domain = field::FieldDomainT<field::Enthalpy>;
        static constexpr int enthalpyIndex = utils::blocks::type_index_v<
            utils::blocks::enthalpy::specific::value,typename Form::value_blocks>;
        static_assert(enthalpyIndex >= 0, "Polynomial positivity requires the compiled specific-enthalpy block.");
        static constexpr utils::blocks::value_block<enthalpyIndex> enthalpyBlock{};
    public:
        explicit PreparedStellarEnthalpyStrategy(State &state, EnthalpyAdmissibilityOptions options = {})
            : m_state(&state),
              m_offset(state.Problem().GetManifest().layout().offset(enthalpyBlock)),
              m_size(state.Problem().GetManifest().layout().size(enthalpyBlock)) {
            detail::RequirePolynomialEnthalpySurface(state.Problem().GetCompiledSurfaceConstraint(),
                state.Problem().GetStellarModel().equationOfState(),state.Problem().GetCommunicator());
            const auto &fem = equilibrium::detail::StellarEquilibriumProblemFactory::FiniteElementModel(state.Problem());
            if (!fem.mesh || !fem.enthalpyFes)
                throw std::invalid_argument("Stellar enthalpy certification requires a complete enthalpy FE space.");
            std::vector<int> elements;
            for (int e = 0; e < fem.mesh->GetNE(); ++e)
                if (Schema::template attribute_belongs_to<Domain>(fem.mesh->GetAttribute(e))) elements.push_back(e);
            const auto &surface = state.Problem().GetPressureSurfaceRows().reduced_dofs();
            m_prepared = std::make_unique<PreparedEnthalpyAdmissibility>(
                *fem.enthalpyFes,field::make_field_dof_map<field::Enthalpy,Schema>(*fem.enthalpyFes),elements,
                std::span<const int>(surface.GetData(),surface.Size()),options);
            if (m_prepared->Size() != m_size)
                throw std::invalid_argument("Compiled enthalpy block disagrees with the supported field map.");
            ExtractEnthalpy(state.AcceptedPhysicalState(),m_anchor);
            m_anchorReport = m_prepared->Check(m_anchor);
        }

        [[nodiscard]] const PolynomialSignReport &AnchorReport() const noexcept { return m_anchorReport; }

        [[nodiscard]] PolynomialSignReport CheckState(const mfem::Vector &fullNormalized) const {
            m_state->normalizedOperator->DenormalizeState(fullNormalized,m_physicalWork);
            ExtractEnthalpy(m_physicalWork,m_enthalpyWork);
            return m_prepared->Check(m_enthalpyWork);
        }

        [[nodiscard]] EnthalpyDirectionRepairReport RepairDirection(State &state) {
            if (&state != m_state) throw std::invalid_argument("A prepared enthalpy strategy belongs to one context.");
            ExtractEnthalpy(state.AcceptedPhysicalState(),m_base);
            state.normalizedOperator->DenormalizeState(state.normalizedCorrection,m_physicalDirection);
            ExtractEnthalpy(m_physicalDirection,m_originalDirection);
            auto report = m_prepared->MinimumRepairFraction(m_base,m_anchor,m_originalDirection);
            if (report.status != EnthalpyDirectionRepairStatus::feasible) return report;

            const auto &surface = state.Problem().GetPressureSurfaceRows().reduced_dofs();
            bool traceUnchanged = true;
            for (int row : surface) traceUnchanged = traceUnchanged && m_originalDirection(row) == 0.0;
            int localChange = traceUnchanged ? 0 : 1, globalChange = 0;
            if (MPI_Allreduce(&localChange,&globalChange,1,MPI_INT,MPI_MAX,state.Problem().GetCommunicator()) != MPI_SUCCESS)
                throw std::runtime_error("Could not resolve collective enthalpy trace preservation.");
            if (report.minimumSafeFraction == 0.0 && globalChange == 0) {
                m_candidateState = state.acceptedNormalizedState; m_candidateState += state.normalizedCorrection;
                report.candidate = CheckState(m_candidateState);
                if (report.candidate.status != PolynomialSignStatus::certified_nonnegative) {
                    report.status = EnthalpyDirectionRepairStatus::unresolved;
                    return report;
                }
                m_chosenAction.SetSize(state.Problem().EquationSize());
                state.normalizedOperator->Mult(state.normalizedCorrection,m_chosenAction);
                report.jacobianApplications = 1;
                m_linearResidual = state.acceptedNormalizedResidual; m_linearResidual += m_chosenAction;
                report.meritDirectionalDerivative = Dot(state.acceptedNormalizedResidual,m_chosenAction);
                report.linearModelResidualNorm = std::sqrt(Dot(m_linearResidual,m_linearResidual));
                if (!std::isfinite(report.meritDirectionalDerivative) || report.meritDirectionalDerivative >= 0.0)
                    report.status = EnthalpyDirectionRepairStatus::no_descent;
                return report;
            }
            m_zeroFractionDirection = state.normalizedCorrection;
            m_physicalWork = m_physicalDirection;
            for (int row : surface) m_physicalWork(m_offset+row) = 0.0;
            state.normalizedOperator->NormalizeState(m_physicalWork,m_zeroFractionDirection);
            // Preserve unrelated blocks bit for bit. A normalize/denormalize
            // round trip is not an excuse to perturb their accepted direction.
            CopyOtherBlocks(state.normalizedCorrection,m_zeroFractionDirection);
            m_unitFractionDirection = state.normalizedCorrection;
            m_physicalWork = m_physicalDirection;
            for (int i = 0; i < m_size; ++i) m_physicalWork(m_offset+i) = m_anchor(i)-m_base(i);
            for (int row : surface) m_physicalWork(m_offset+row) = 0.0;
            state.normalizedOperator->NormalizeState(m_physicalWork,m_unitFractionDirection);
            CopyOtherBlocks(state.normalizedCorrection,m_unitFractionDirection);
            m_familyDirection = m_unitFractionDirection; m_familyDirection -= m_zeroFractionDirection;
            m_firstAction.SetSize(state.Problem().EquationSize());
            m_familyAction.SetSize(state.Problem().EquationSize());
            state.normalizedOperator->Mult(m_zeroFractionDirection,m_firstAction);
            state.normalizedOperator->Mult(m_familyDirection,m_familyAction);
            report.jacobianApplications = 2;
            m_linearResidual = state.acceptedNormalizedResidual; m_linearResidual += m_firstAction;
            const double denominator = Dot(m_familyAction,m_familyAction);
            const double numerator = Dot(m_familyAction,m_linearResidual);
            double fraction = report.minimumSafeFraction;
            if (std::isfinite(denominator) && denominator > 0.0 && std::isfinite(numerator))
                fraction = std::clamp(-numerator/denominator,report.minimumSafeFraction,1.0);

            auto candidate = m_prepared->ApplyRepairFraction(m_base,m_anchor,m_originalDirection,fraction,m_repaired);
            if (candidate.status != PolynomialSignStatus::certified_nonnegative && fraction != 1.0) {
                fraction = 1.0;
                candidate = m_prepared->ApplyRepairFraction(m_base,m_anchor,m_originalDirection,fraction,m_repaired);
            }
            if (candidate.status != PolynomialSignStatus::certified_nonnegative) {
                report.status = EnthalpyDirectionRepairStatus::unresolved;
                report.candidate = candidate;
                return report;
            }
            m_physicalWork = m_physicalDirection;
            for (int i = 0; i < m_size; ++i) m_physicalWork(m_offset+i) = m_repaired(i);
            state.normalizedOperator->NormalizeState(m_physicalWork,m_chosenDirection);
            CopyOtherBlocks(state.normalizedCorrection,m_chosenDirection);

            // State normalization and root addition introduce another rounding
            // boundary. Certify exactly what backtracking will construct.
            m_candidateState = state.acceptedNormalizedState; m_candidateState += m_chosenDirection;
            candidate = CheckState(m_candidateState);
            if (candidate.status != PolynomialSignStatus::certified_nonnegative) {
                report.status = EnthalpyDirectionRepairStatus::unresolved;
                report.candidate = candidate;
                return report;
            }
            m_chosenAction.SetSize(state.Problem().EquationSize());
            state.normalizedOperator->Mult(m_chosenDirection,m_chosenAction);
            ++report.jacobianApplications;
            m_linearResidual = state.acceptedNormalizedResidual; m_linearResidual += m_chosenAction;
            report.meritDirectionalDerivative = Dot(state.acceptedNormalizedResidual,m_chosenAction);
            report.linearModelResidualNorm = std::sqrt(Dot(m_linearResidual,m_linearResidual));
            report.chosenFraction = fraction;
            report.candidate = candidate;
            // A feasible family can lack descent. Leave the original direction
            // available for diagnostics rather than silently accepting that case.
            if (!std::isfinite(report.meritDirectionalDerivative) || report.meritDirectionalDerivative >= 0.0) {
                report.status = EnthalpyDirectionRepairStatus::no_descent;
                return report;
            }
            state.normalizedCorrection = m_chosenDirection;
            return report;
        }

    private:
        void ExtractEnthalpy(const mfem::Vector &physical,mfem::Vector &enthalpy) const {
            if (physical.Size() != m_state->Problem().StateSize())
                throw std::invalid_argument("Stellar positivity received an incompatible root state.");
            enthalpy.SetSize(m_size);
            for (int i = 0; i < m_size; ++i) enthalpy(i) = physical(m_offset+i);
        }
        void CopyOtherBlocks(const mfem::Vector &source,mfem::Vector &target) const {
            for (int i = 0; i < m_offset; ++i) target(i) = source(i);
            for (int i = m_offset+m_size; i < target.Size(); ++i) target(i) = source(i);
        }
        [[nodiscard]] double Dot(const mfem::Vector &left,const mfem::Vector &right) const {
            const double local = left*right;
            double global = 0.0;
            if (MPI_Allreduce(&local,&global,1,MPI_DOUBLE,MPI_SUM,m_state->Problem().GetCommunicator()) != MPI_SUCCESS)
                throw std::runtime_error("Could not evaluate the collective feasible-direction model.");
            return global;
        }
        State *m_state;
        int m_offset,m_size;
        std::unique_ptr<PreparedEnthalpyAdmissibility> m_prepared;
        mfem::Vector m_anchor,m_base,m_originalDirection,m_repaired;
        PolynomialSignReport m_anchorReport;
        mutable mfem::Vector m_physicalWork,m_enthalpyWork;
        mfem::Vector m_physicalDirection,m_zeroFractionDirection,m_unitFractionDirection,m_familyDirection;
        mfem::Vector m_firstAction,m_familyAction,m_chosenAction,m_linearResidual,m_chosenDirection,m_candidateState;
    };
}