module;

#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <cstddef>
#include <optional>
#include <span>
#include <stdexcept>
#include <string_view>
#include <type_traits>

#include <mfem.hpp>

export module mean_field:operators.root_manifest;

export import :model.compiled_fixed_mass;
export import :model.compiled_fixed_central_density;
export import :model.specifications;
export import :utils.blocks;

export namespace mean_field::operators {
    enum class RootBlockKind { value, residual };
    enum class RootBlockProvenance { physical_operator, model_specification };
    enum class RootRowInjection { physical_equation, append_global, replace_carrier_rows };
    enum class RootColumnPolicy { physical_state, existing_physical_multiplier, solver_border, no_column };
    enum class RootScalePolicy { unscaled, target_relative };

    struct RootBlockDescriptor final {
        std::string_view stableId;
        std::string_view symbol;
        RootBlockKind kind;
        RootBlockProvenance provenance;
        std::string_view source;
        RootRowInjection rowInjection;
        RootColumnPolicy columnPolicy;
        RootScalePolicy scalePolicy;
        int canonicalIndex;
        int offset;
        int size;
        double scale;
    };

    struct RootRowReplacementDescriptor final {
        std::string_view stableId;
        std::string_view sourceSpecification;
        models::SpecificationRole role;
        int carrierResidualBlock;
        int replacedRowCount;
    };

    struct RootConstraintDescriptor final {
        std::string_view stableId;
        models::SpecificationRole role;
        RootRowInjection rowInjection;
        RootColumnPolicy columnPolicy;
        int valueBlock;
        int residualBlock;
        int rowArity;
        int columnArity;
        double target;
        std::optional<double> carrierTarget;
        std::string_view targetUnits;
        std::string_view residualUnits;
        double residualScale;
    };

    struct CentralDensityManifestInput final {
        double targetDensity;
        double targetEnthalpy;
        int centerDofCount;
    };

    struct RootConstraintReport final {
        RootConstraintDescriptor descriptor;
        double achieved;
        double dimensionalResidual;
        double scaledResidual;
    };

    namespace detail {
        struct StaticRootBlockDescriptor final {
            std::string_view stableId;
            std::string_view symbol;
            RootBlockProvenance provenance;
            std::string_view source;
            RootRowInjection rowInjection;
            RootColumnPolicy columnPolicy;
            RootScalePolicy scalePolicy;
        };

        template <typename Block> struct RootBlockTraits;

#define MEAN_FIELD_PHYSICAL_VALUE_BLOCK(BlockType, StableId, Symbol)                                                   \
    template <> struct RootBlockTraits<BlockType> {                                                                    \
        static constexpr StaticRootBlockDescriptor descriptor{                                                         \
            StableId,                                                                                                  \
            Symbol,                                                                                                    \
            RootBlockProvenance::physical_operator,                                                                    \
            "stellar_equilibrium",                                                                                     \
            RootRowInjection::physical_equation,                                                                       \
            RootColumnPolicy::physical_state,                                                                          \
            RootScalePolicy::unscaled                                                                                  \
        };                                                                                                             \
    }

#define MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(BlockType, StableId, Symbol)                                                \
    template <> struct RootBlockTraits<BlockType> {                                                                    \
        static constexpr StaticRootBlockDescriptor descriptor{                                                         \
            StableId,                                                                                                  \
            Symbol,                                                                                                    \
            RootBlockProvenance::physical_operator,                                                                    \
            "stellar_equilibrium",                                                                                     \
            RootRowInjection::physical_equation,                                                                       \
            RootColumnPolicy::no_column,                                                                               \
            RootScalePolicy::unscaled                                                                                  \
        };                                                                                                             \
    }

        MEAN_FIELD_PHYSICAL_VALUE_BLOCK(
            utils::blocks::density::mass::value,
            "density",
            "rho"
        );
        MEAN_FIELD_PHYSICAL_VALUE_BLOCK(
            utils::blocks::displacement::geometry::value,
            "volume_displacement",
            "d"
        );
        MEAN_FIELD_PHYSICAL_VALUE_BLOCK(
            utils::blocks::surface_deformation::parameters::value,
            "surface_deformation",
            "q"
        );
        MEAN_FIELD_PHYSICAL_VALUE_BLOCK(
            utils::blocks::gravity::gradient::value,
            "gravity_gradient",
            "g"
        );
        MEAN_FIELD_PHYSICAL_VALUE_BLOCK(
            utils::blocks::gravity::poisson::value,
            "gravity_potential",
            "Phi"
        );
        MEAN_FIELD_PHYSICAL_VALUE_BLOCK(
            utils::blocks::enthalpy::specific::value,
            "specific_enthalpy",
            "h"
        );

        MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(
            utils::blocks::gravity::gradient::residual,
            "gravity_gradient_relation",
            "R_g"
        );
        MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(
            utils::blocks::gravity::poisson::residual,
            "poisson_balance",
            "R_Phi"
        );
        MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(
            utils::blocks::density::mass::residual,
            "barotropic_closure",
            "R_rho"
        );
        MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(
            utils::blocks::displacement::geometry::residual,
            "mechanical_balance",
            "R_d"
        );
        MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(
            utils::blocks::surface_deformation::shape_equilibrium::residual,
            "surface_shape_balance",
            "R_q"
        );
        MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK(
            utils::blocks::enthalpy::specific::residual,
            "hydrostatic_balance",
            "R_h"
        );

#undef MEAN_FIELD_PHYSICAL_VALUE_BLOCK
#undef MEAN_FIELD_PHYSICAL_RESIDUAL_BLOCK

        template <> struct RootBlockTraits<utils::blocks::fixed_total_mass::mass_normalization::value> {
            static constexpr StaticRootBlockDescriptor descriptor{
                "fixed_total_mass.multiplier",
                "C",
                RootBlockProvenance::model_specification,
                "FixedTotalMass",
                RootRowInjection::physical_equation,
                RootColumnPolicy::existing_physical_multiplier,
                RootScalePolicy::unscaled
            };
        };

        template <> struct RootBlockTraits<utils::blocks::fixed_total_mass::mass_normalization::residual> {
            static constexpr StaticRootBlockDescriptor descriptor{
                "fixed_total_mass.residual",
                "R_M",
                RootBlockProvenance::model_specification,
                "FixedTotalMass",
                RootRowInjection::append_global,
                RootColumnPolicy::no_column,
                RootScalePolicy::target_relative
            };
        };

        template <> struct RootBlockTraits<utils::blocks::fixed_central_density::central_value::value> {
            static constexpr StaticRootBlockDescriptor descriptor{
                "fixed_central_density.border",
                "lambda_rho_c",
                RootBlockProvenance::model_specification,
                "FixedCentralDensity",
                RootRowInjection::physical_equation,
                RootColumnPolicy::solver_border,
                RootScalePolicy::unscaled
            };
        };

        template <> struct RootBlockTraits<utils::blocks::fixed_central_density::central_value::residual> {
            static constexpr StaticRootBlockDescriptor descriptor{
                "fixed_central_density.residual",         "R_rho_c",
                RootBlockProvenance::model_specification, "FixedCentralDensity",
                RootRowInjection::append_global,          RootColumnPolicy::no_column,
                RootScalePolicy::target_relative
            };
        };

        template <typename Block>
        [[nodiscard]] constexpr double blockScale(
            const double fixedMassScale,
            const double centralDensityScale
        ) noexcept {
            if constexpr (std::same_as<Block, utils::blocks::fixed_total_mass::mass_normalization::residual>) {
                return fixedMassScale;
            } else if constexpr (std::same_as<Block, utils::blocks::fixed_central_density::central_value::residual>) {
                return centralDensityScale;
            } else {
                return 1.0;
            }
        }

        template <
            RootBlockKind Kind,
            typename... Blocks>
        [[nodiscard]] std::array<
            RootBlockDescriptor,
            sizeof...(Blocks)>
        makeBlockDescriptors(
            const mfem::Array<int> &offsets,
            const double fixedMassScale,
            const double centralDensityScale,
            utils::blocks::type_list<Blocks...>
        ) {
            std::array<RootBlockDescriptor, sizeof...(Blocks)> descriptors{};
            int index = 0;
            ((descriptors[index] =
                  {.stableId       = RootBlockTraits<Blocks>::descriptor.stableId,
                   .symbol         = RootBlockTraits<Blocks>::descriptor.symbol,
                   .kind           = Kind,
                   .provenance     = RootBlockTraits<Blocks>::descriptor.provenance,
                   .source         = RootBlockTraits<Blocks>::descriptor.source,
                   .rowInjection   = RootBlockTraits<Blocks>::descriptor.rowInjection,
                   .columnPolicy   = RootBlockTraits<Blocks>::descriptor.columnPolicy,
                   .scalePolicy    = RootBlockTraits<Blocks>::descriptor.scalePolicy,
                   .canonicalIndex = index,
                   .offset         = offsets[index],
                   .size           = offsets[index + 1] - offsets[index],
                   .scale          = blockScale<Blocks>(fixedMassScale, centralDensityScale)},
              ++index),
             ...);
            return descriptors;
        }

        template <models::SpecifiedModelType Model>
        inline constexpr bool hasCentralDensity = Model::template containsSpecification<models::FixedCentralDensity>;

        template <models::SpecifiedModelType Model>
        inline constexpr std::size_t rootConstraintCount = 2 + (hasCentralDensity<Model> ? 1 : 0);

        template <
            models::SpecifiedModelType Model,
            typename Form>
        [[nodiscard]] std::array<
            RootConstraintDescriptor,
            rootConstraintCount<Model>>
        makeConstraintDescriptors(
            const double targetMass,
            const double targetSurfacePressure,
            const double fixedMassScale,
            const std::optional<CentralDensityManifestInput> centralDensity
        ) {
            std::array<RootConstraintDescriptor, rootConstraintCount<Model>> descriptors{};
            descriptors[0] = {
                .stableId      = "FixedTotalMass",
                .role          = models::SpecificationRole::invariant,
                .rowInjection  = RootRowInjection::append_global,
                .columnPolicy  = RootColumnPolicy::existing_physical_multiplier,
                .valueBlock    = models::FixedMassLayoutRequest::valueBlock<Form>().index,
                .residualBlock = models::FixedMassLayoutRequest::residualBlock<Form>().index,
                .rowArity      = 1,
                .columnArity   = 1,
                .target        = targetMass,
                .carrierTarget = targetMass,
                .targetUnits   = "mass",
                .residualUnits = "mass",
                .residualScale = fixedMassScale
            };
            descriptors[1] = {
                .stableId     = "IsobaricSurface",
                .role         = models::SpecificationRole::boundary_condition,
                .rowInjection = RootRowInjection::replace_carrier_rows,
                .columnPolicy = RootColumnPolicy::no_column,
                .valueBlock   = -1,
                .residualBlock =
                    utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term).index,
                .rowArity      = 0,
                .columnArity   = 0,
                .target        = targetSurfacePressure,
                .carrierTarget = std::nullopt,
                .targetUnits   = "pressure",
                .residualUnits = "specific_enthalpy",
                .residualScale = 1.0
            };

            if constexpr (hasCentralDensity<Model>) {
                if (!centralDensity.has_value()) {
                    throw std::invalid_argument(
                        "A model containing FixedCentralDensity requires central-density manifest metadata."
                    );
                }
                descriptors[2] = {
                    .stableId      = "FixedCentralDensity",
                    .role          = models::SpecificationRole::phase_condition,
                    .rowInjection  = RootRowInjection::append_global,
                    .columnPolicy  = RootColumnPolicy::solver_border,
                    .valueBlock    = models::CentralDensityLayoutRequest::valueBlock<Form>().index,
                    .residualBlock = models::CentralDensityLayoutRequest::residualBlock<Form>().index,
                    .rowArity      = 1,
                    .columnArity   = 1,
                    .target        = centralDensity->targetDensity,
                    .carrierTarget = centralDensity->targetEnthalpy,
                    .targetUnits   = "density",
                    .residualUnits = "specific_enthalpy",
                    .residualScale = std::max(std::abs(centralDensity->targetEnthalpy), 1.0e-300)
                };
            } else if (centralDensity.has_value()) {
                throw std::invalid_argument(
                    "Central-density manifest metadata was provided to a model without FixedCentralDensity."
                );
            }
            return descriptors;
        }
    } // namespace detail

    template <typename Form> class RootStateView final {
    public:
        RootStateView(
            const mfem::Vector &state,
            const utils::blocks::form_layout<Form> &layout
        )
            : m_state(state),
              m_layout(layout) {
            if (state.Size() != layout.value_offsets().Last()) {
                throw std::invalid_argument("RootStateView received a vector with the wrong size.");
            }
        }

        template <typename Term> [[nodiscard]] mfem::Vector block(const Term &term) const {
            constexpr auto valueBlock = utils::blocks::get_value_block<Form>(term);
            return mfem::Vector(
                const_cast<mfem::real_t *>(m_state.GetData()) + m_layout.offset(valueBlock), m_layout.size(valueBlock)
            );
        }

        [[nodiscard]] const mfem::Vector &vector() const noexcept {
            return m_state;
        }

    private:
        const mfem::Vector &m_state;
        const utils::blocks::form_layout<Form> &m_layout;
    };

    template <typename Form> class ResidualView final {
    public:
        ResidualView(
            mfem::Vector &residual,
            const utils::blocks::form_layout<Form> &layout
        )
            : m_residual(residual),
              m_layout(layout) {
            if (residual.Size() != layout.residual_offsets().Last()) {
                throw std::invalid_argument("ResidualView received a vector with the wrong size.");
            }
        }

        template <typename Term> [[nodiscard]] mfem::Vector block(const Term &term) const {
            constexpr auto residualBlock = utils::blocks::get_residual_block<Form>(term);
            return mfem::Vector(m_residual.GetData() + m_layout.offset(residualBlock), m_layout.size(residualBlock));
        }

        template <typename Term>
        void assign(
            const Term &term,
            const mfem::Vector &source
        ) const {
            mfem::Vector destination = block(term);
            if (destination.Size() != source.Size()) {
                throw std::invalid_argument("ResidualView block assignment has the wrong size.");
            }
            destination = source;
        }

        [[nodiscard]] mfem::Vector &vector() const noexcept {
            return m_residual;
        }

    private:
        mfem::Vector &m_residual;
        const utils::blocks::form_layout<Form> &m_layout;
    };

    template <models::SpecifiedModelType Model, typename Form, typename JacobianForm>
        requires utils::blocks::valid_jacobian_form<Form, JacobianForm>
    class CompiledRootManifest final {
    public:
        using ModelType                                                 = Model;
        using FormType                                                  = Form;
        using JacobianType                                              = JacobianForm;
        using Layout                                                    = utils::blocks::form_layout<Form>;
        using StateView                                                 = RootStateView<Form>;
        using DirectionView                                             = RootStateView<Form>;
        using RootResidualView                                          = ResidualView<Form>;

        static constexpr models::ModelCompilationClass compilationClass = Model::compilationClass;
        static constexpr bool symbolicallySquare                        = Model::symbolicallySquare;

        CompiledRootManifest(
            const std::array<
                int,
                Form::value_block_count> &valueSizes,
            const std::array<
                int,
                Form::residual_block_count> &residualSizes,
            const double targetMass,
            const double targetSurfacePressure,
            const int replacedSurfaceRowCount,
            const std::optional<CentralDensityManifestInput> centralDensity = std::nullopt
        )
            : m_layout(
                  valueSizes,
                  residualSizes
              ),
              m_fixedMassScale(
                  std::max(
                      std::abs(targetMass),
                      1.0e-300
                  )
              ),
              m_centralDensityScale(
                  centralDensity.has_value() ? std::max(
                                                   std::abs(centralDensity->targetEnthalpy),
                                                   1.0e-300
                                               )
                                             : 1.0
              ),
              m_valueBlocks(
                  detail::makeBlockDescriptors<RootBlockKind::value>(
                      m_layout.value_offsets(),
                      m_fixedMassScale,
                      m_centralDensityScale,
                      typename Form::value_blocks{}
                  )
              ),
              m_residualBlocks(
                  detail::makeBlockDescriptors<RootBlockKind::residual>(
                      m_layout.residual_offsets(),
                      m_fixedMassScale,
                      m_centralDensityScale,
                      typename Form::residual_blocks{}
                  )
              ),
              m_replacements{RootRowReplacementDescriptor{
                  .stableId            = "isobaric_surface.replacement",
                  .sourceSpecification = "IsobaricSurface",
                  .role                = models::SpecificationRole::boundary_condition,
                  .carrierResidualBlock =
                      utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term).index,
                  .replacedRowCount = replacedSurfaceRowCount
              }},
              m_constraints(
                  detail::makeConstraintDescriptors<
                      Model,
                      Form>(
                      targetMass,
                      targetSurfacePressure,
                      m_fixedMassScale,
                      centralDensity
                  )
              ) {
            if (replacedSurfaceRowCount < 0) {
                throw std::invalid_argument(
                    "An equilibrium-system manifest cannot contain a negative replacement-row count."
                );
            }
            if (centralDensity.has_value() && centralDensity->centerDofCount < 0) {
                throw std::invalid_argument(
                    "An equilibrium-system manifest cannot contain a negative central-DOF count."
                );
            }
            if constexpr (compilationClass == models::EquilibriumSystemCompilation::complete_equilibrium_system) {
                if (m_layout.value_offsets().Last() != m_layout.residual_offsets().Last()) {
                    throw std::invalid_argument(
                        "A complete equilibrium system must have equal state and equation dimensions."
                    );
                }
            }
        }

        [[nodiscard]] const Layout &layout() const noexcept {
            return m_layout;
        }

        [[nodiscard]] StateView stateView(const mfem::Vector &state) const {
            return {state, m_layout};
        }

        [[nodiscard]] DirectionView directionView(const mfem::Vector &direction) const {
            return {direction, m_layout};
        }

        [[nodiscard]] RootResidualView residualView(mfem::Vector &residual) const {
            return {residual, m_layout};
        }

        [[nodiscard]] std::span<const RootBlockDescriptor> valueBlocks() const noexcept {
            return m_valueBlocks;
        }

        [[nodiscard]] std::span<const RootBlockDescriptor> residualBlocks() const noexcept {
            return m_residualBlocks;
        }

        [[nodiscard]] std::span<const RootRowReplacementDescriptor> rowReplacements() const noexcept {
            return m_replacements;
        }

        [[nodiscard]] std::span<const RootConstraintDescriptor> constraints() const noexcept {
            return m_constraints;
        }

        [[nodiscard]] static constexpr std::span<const models::RuntimeSpecificationDescriptor>
        specificationDescriptors() noexcept {
            return Model::runtimeSpecificationDescriptors();
        }

        [[nodiscard]] RootConstraintReport fixedMassReport(const double achievedMass) const {
            const RootConstraintDescriptor &descriptor = m_constraints[0];
            const double residual                      = achievedMass - descriptor.target;
            return {
                .descriptor          = descriptor,
                .achieved            = achievedMass,
                .dimensionalResidual = residual,
                .scaledResidual      = residual / descriptor.residualScale
            };
        }

    private:
        Layout m_layout;
        double m_fixedMassScale;
        double m_centralDensityScale;
        std::array<RootBlockDescriptor, Form::value_block_count> m_valueBlocks;
        std::array<RootBlockDescriptor, Form::residual_block_count> m_residualBlocks;
        std::array<RootRowReplacementDescriptor, 1> m_replacements;
        std::array<RootConstraintDescriptor, detail::rootConstraintCount<Model>> m_constraints;
    };

    // Physics-facing names for the public equilibrium-system boundary. The
    // root-oriented names remain available while existing solver consumers
    // migrate, but new APIs should expose these aliases.
    using EquilibriumBlockKind                             = RootBlockKind;
    using EquilibriumBlockProvenance                       = RootBlockProvenance;
    using EquilibriumEquationInjection                     = RootRowInjection;
    using EquilibriumGeneratedVariablePolicy               = RootColumnPolicy;
    using EquilibriumScalePolicy                           = RootScalePolicy;
    using EquilibriumBlockDescriptor                       = RootBlockDescriptor;
    using EquilibriumEquationReplacementDescriptor         = RootRowReplacementDescriptor;
    using EquilibriumSpecificationDescriptor               = RootConstraintDescriptor;
    using EquilibriumSpecificationReport                   = RootConstraintReport;

    template <typename Form> using EquilibriumStateView    = RootStateView<Form>;

    template <typename Form> using EquilibriumResidualView = ResidualView<Form>;

    template <models::SpecifiedModelType Model, typename Form, typename JacobianForm>
        requires utils::blocks::valid_jacobian_form<Form, JacobianForm>
    using EquilibriumSystemManifest = CompiledRootManifest<Model, Form, JacobianForm>;
} // namespace mean_field::operators