module;
#include <array>
#include <mfem.hpp>
#include <stdexcept>
#include <tuple>
#include <type_traits>
export module mean_field:utils.blocks;

export namespace mean_field::utils::blocks {
    inline constexpr int dynamic_block_size = -1;

    struct block { };

    struct residual_block_base : block {
        static constexpr int static_block_size = dynamic_block_size;
    };

    struct value_block_base : block {
        static constexpr int static_block_size = dynamic_block_size;
    };

    struct term { };
    struct field { };

    template <typename Residual, typename... Values> struct block_row { };
    template <int index_value>
    struct residual_block final : residual_block_base {
        static constexpr int index = index_value;

        // ReSharper disable once CppNonExplicitConversionOperator
        constexpr operator int() const noexcept {
            return index;
        }
    };

    template <int index_value> struct value_block final : value_block_base {
        static constexpr int index = index_value;

        // ReSharper disable once CppNonExplicitConversionOperator
        constexpr operator int() const noexcept {
            return index;
        }
    };

    struct density final : field {
        struct mass final : term {
            struct value final : value_block_base { };
            struct residual final : residual_block_base { };
        };

        static inline constexpr mass mass_term{};
    };

    struct displacement final : field {
        struct geometry final : term {
            struct value final : value_block_base { };
            struct residual final : residual_block_base { };
        };

        static inline constexpr geometry geometry_term{};
    };

    struct gravity final : field {
        struct gradient final : term {
            struct value final : value_block_base { };
            struct residual final : residual_block_base { };
        };

        struct poisson final : term {
            struct value final : value_block_base { };
            struct residual final : residual_block_base { };
        };

        static inline constexpr gradient gradient_term{};
        static inline constexpr poisson poisson_term{};
    };

    struct enthalpy final : field {
        struct specific final : term {
            struct value final : value_block_base { };
            struct residual final : residual_block_base { };
        };

        static inline constexpr specific specific_term{};
    };

    struct barotropic_constant final : field {
        struct mass_normalization final : term {
            struct value final : value_block_base {
                static constexpr int static_block_size = 1;
            };

            struct residual final : residual_block_base {
                static constexpr int static_block_size = 1;
            };
        };

        static inline constexpr mass_normalization mass_normalization_term{};
    };

    inline constexpr density density_field{};
    inline constexpr displacement displacement_field{};
    inline constexpr gravity gravity_field{};
    inline constexpr enthalpy enthalpy_field{};
    inline constexpr barotropic_constant barotropic_constant_field{};

    template <typename... Types> struct type_list {
        static constexpr int size = sizeof...(Types);
    };

    template <typename Query, typename List> struct contains_type;

    template <typename Query>
    struct contains_type<Query, type_list<>> : std::false_type { };

    template <typename Query, typename Head, typename... Tail>
    struct contains_type<Query, type_list<Head, Tail...>>
        : std::conditional_t<
              std::is_same_v<Query, Head>,
              std::true_type,
              contains_type<Query, type_list<Tail...>>> { };

    template <typename Query, typename List>
    inline constexpr bool contains_type_v = contains_type<Query, List>::value;

    template <typename Query, typename List> struct type_count;

    template <typename Query>
    struct type_count<Query, type_list<>> : std::integral_constant<int, 0> { };

    template <typename Query, typename Head, typename... Tail>
    struct type_count<Query, type_list<Head, Tail...>>
        : std::integral_constant<
              int,
              (std::is_same_v<Query, Head> ? 1 : 0) +
                  type_count<Query, type_list<Tail...>>::value> { };

    template <typename Query, typename List>
    inline constexpr int type_count_v = type_count<Query, List>::value;

    template <typename List> struct types_are_unique;

    template <typename... Types>
    struct types_are_unique<type_list<Types...>>
        : std::bool_constant<
              ((type_count_v<Types, type_list<Types...>> == 1) && ...)> { };

    template <typename List>
    inline constexpr bool types_are_unique_v = types_are_unique<List>::value;

    template <typename Row> struct block_row_traits {
        using residual                     = void;
        using values                       = type_list<>;

        static constexpr int value_count   = 0;
        static constexpr bool is_block_row = false;
    };

    template <typename Residual, typename... Values>
    struct block_row_traits<block_row<Residual, Values...>> {
        using residual                     = Residual;
        using values                       = type_list<Values...>;

        static constexpr int value_count   = sizeof...(Values);
        static constexpr bool is_block_row = true;
    };

    template <typename Query, typename List> struct type_index;

    template <typename Query, typename... Tail>
    struct type_index<Query, type_list<Query, Tail...>> {
        static constexpr int value = 0;
    };

    template <typename Query, typename Head, typename... Tail>
    struct type_index<Query, type_list<Head, Tail...>> {
        static constexpr int value =
            1 + type_index<Query, type_list<Tail...>>::value;
    };

    template <typename Query, typename List>
    inline constexpr int type_index_v = type_index<Query, List>::value;

    template <typename ValueBlocks, typename ResidualBlocks> struct block_form {
        using value_blocks                        = ValueBlocks;
        using residual_blocks                     = ResidualBlocks;

        static constexpr int value_block_count    = ValueBlocks::size;
        static constexpr int residual_block_count = ResidualBlocks::size;
    };

    template <typename Form> struct block_form_is_valid : std::false_type { };

    template <typename... Values, typename... Residuals>
    struct block_form_is_valid<
        block_form<type_list<Values...>, type_list<Residuals...>>>
        : std::bool_constant<
              (std::is_base_of_v<value_block_base, Values> && ...) &&
              (std::is_base_of_v<residual_block_base, Residuals> && ...) &&
              types_are_unique_v<type_list<Values...>> &&
              types_are_unique_v<type_list<Residuals...>>> { };

    template <typename Form>
    inline constexpr bool block_form_is_valid_v =
        block_form_is_valid<Form>::value;

    template <typename Row, typename ValueBlocks, typename ResidualBlocks>
    struct block_row_is_valid : std::false_type { };

    template <
        typename Residual,
        typename... Values,
        typename ValueBlocks,
        typename ResidualBlocks>
    struct block_row_is_valid<
        block_row<Residual, Values...>,
        ValueBlocks,
        ResidualBlocks>
        : std::bool_constant<
              std::is_base_of_v<residual_block_base, Residual> &&
              contains_type_v<Residual, ResidualBlocks> &&
              ((std::is_base_of_v<value_block_base, Values> &&
                contains_type_v<Values, ValueBlocks>) &&
               ...) &&
              types_are_unique_v<type_list<Values...>>> { };

    template <typename Rows> struct row_residual_list;

    template <typename... Rows> struct row_residual_list<type_list<Rows...>> {
        using type = type_list<typename block_row_traits<Rows>::residual...>;
    };

    template <typename Rows>
    using row_residual_list_t = typename row_residual_list<Rows>::type;

    template <typename Form, typename JacobianForm>
    struct jacobian_form_is_valid : std::false_type { };

    template <typename... Values, typename... Residuals, typename... Rows>
    struct jacobian_form_is_valid<
        block_form<type_list<Values...>, type_list<Residuals...>>,
        type_list<Rows...>> {
        using form_type =
            block_form<type_list<Values...>, type_list<Residuals...>>;

        using value_blocks    = type_list<Values...>;
        using residual_blocks = type_list<Residuals...>;
        using rows            = type_list<Rows...>;

        static constexpr bool value =
            block_form_is_valid_v<form_type> &&
            (block_row_is_valid<Rows, value_blocks, residual_blocks>::value &&
             ...) &&
            std::is_same_v<row_residual_list_t<rows>, residual_blocks>;
    };

    template <typename Form, typename JacobianForm>
    inline constexpr bool jacobian_form_is_valid_v =
        jacobian_form_is_valid<Form, JacobianForm>::value;

    template <typename Form, typename JacobianForm>
    concept valid_jacobian_form = jacobian_form_is_valid_v<Form, JacobianForm>;

    template <typename Residual, typename Value, typename JacobianForm>
    struct has_jacobian_coupling;

    template <typename Residual, typename Value>
    struct has_jacobian_coupling<Residual, Value, type_list<>>
        : std::false_type { };

    template <
        typename Residual,
        typename Value,
        typename RowResidual,
        typename... RowValues,
        typename... RemainingRows>
    struct has_jacobian_coupling<
        Residual,
        Value,
        type_list<block_row<RowResidual, RowValues...>, RemainingRows...>>
        : std::conditional_t<
              std::is_same_v<Residual, RowResidual>,
              std::bool_constant<(std::is_same_v<Value, RowValues> || ...)>,
              has_jacobian_coupling<
                  Residual,
                  Value,
                  type_list<RemainingRows...>>> { };

    template <typename Residual, typename Value, typename JacobianForm>
    inline constexpr bool has_jacobian_coupling_v =
        has_jacobian_coupling<Residual, Value, JacobianForm>::value;

    template <
        typename Form,
        typename Term>
    consteval auto get_value_block(const Term &) {
        using value_type = typename Term::value;
        constexpr int index =
            type_index_v<value_type, typename Form::value_blocks>;
        return value_block<index>{};
    }

    template <
        typename Form,
        typename Term>
    consteval auto get_residual_block(const Term &) {
        using residual_type = typename Term::residual;
        constexpr int index =
            type_index_v<residual_type, typename Form::residual_blocks>;
        return residual_block<index>{};
    }

    template <typename Form> class form_layout {
    public:
        form_layout(
            const std::array<
                int,
                Form::value_block_count> &value_sizes,
            const std::array<
                int,
                Form::residual_block_count> &residual_sizes
        ) {
            build_offsets(
                m_value_offsets, value_sizes, typename Form::value_blocks{}
            );

            build_offsets(
                m_residual_offsets, residual_sizes,
                typename Form::residual_blocks{}
            );
        }

        template <int index> [[nodiscard]] int size(value_block<index>) const {
            return m_value_offsets[index + 1] - m_value_offsets[index];
        }

        template <int index>
        [[nodiscard]] int size(residual_block<index>) const {
            return m_residual_offsets[index + 1] - m_residual_offsets[index];
        }

        template <int index>
        [[nodiscard]] int offset(value_block<index>) const {
            return m_value_offsets[index];
        }

        template <int index>
        [[nodiscard]] int offset(residual_block<index>) const {
            return m_residual_offsets[index];
        }

        [[nodiscard]]
        const mfem::Array<int> &value_offsets() const noexcept {
            return m_value_offsets;
        }

        [[nodiscard]]
        const mfem::Array<int> &residual_offsets() const noexcept {
            return m_residual_offsets;
        }

    private:
        template <typename BlockType>
        [[nodiscard]]
        static int resolve_block_size(const int requested_size) {
            if constexpr (BlockType::static_block_size == dynamic_block_size) {
                return requested_size;
            } else {
                if (requested_size != BlockType::static_block_size) {
                    throw std::invalid_argument(
                        "A statically sized block was given an "
                        "incompatible runtime size."
                    );
                }

                return BlockType::static_block_size;
            }
        }

        template <typename... BlockTypes>
        static void build_offsets(
            mfem::Array<int> &offsets,
            const std::array<
                int,
                sizeof...(BlockTypes)> &requested_sizes,
            type_list<BlockTypes...>
        ) {
            offsets.SetSize(sizeof...(BlockTypes) + 1);
            offsets[0]      = 0;

            int block_index = 0;

            ((offsets[block_index + 1] =
                  offsets[block_index] +
                  resolve_block_size<BlockTypes>(requested_sizes[block_index]),
              ++block_index),
             ...);
        }

        mfem::Array<int> m_value_offsets;
        mfem::Array<int> m_residual_offsets;
    };
    using gravity_field_form = block_form<
        type_list<
            density::mass::value,
            displacement::geometry::value,
            gravity::gradient::value,
            gravity::poisson::value>,
        type_list<gravity::gradient::residual, gravity::poisson::residual>>;

    using gravity_jacobian_form = type_list<
        block_row<
            gravity::gradient::residual,
            gravity::gradient::value,
            gravity::poisson::value,
            displacement::geometry::value>,
        block_row<
            gravity::poisson::residual,
            gravity::gradient::value,
            density::mass::value,
            displacement::geometry::value>>;

    // Columns:
    // [rho, d, g, Phi, h, C]
    //
    // Rows:
    // [R_g, R_Phi, R_rho, R_d, R_h, R_M]
    using barotropic_equilibrium_form = block_form<
        type_list<
            density::mass::value,
            displacement::geometry::value,
            gravity::gradient::value,
            gravity::poisson::value,
            enthalpy::specific::value,
            barotropic_constant::mass_normalization::value>,
        type_list<
            gravity::gradient::residual,
            gravity::poisson::residual,
            density::mass::residual,
            displacement::geometry::residual,
            enthalpy::specific::residual,
            barotropic_constant::mass_normalization::residual>>;

    using barotropic_equilibrium_jacobian_form = type_list<
        // R_g(g, Phi, d)
        block_row<
            gravity::gradient::residual,
            gravity::gradient::value,
            gravity::poisson::value,
            displacement::geometry::value>,

        // R_Phi(g, rho, d)
        block_row<
            gravity::poisson::residual,
            gravity::gradient::value,
            density::mass::value,
            displacement::geometry::value>,

        // R_rho(rho, h, d)
        block_row<
            density::mass::residual,
            density::mass::value,
            enthalpy::specific::value,
            displacement::geometry::value>,

        // R_d(d, h)
        block_row<
            displacement::geometry::residual,
            displacement::geometry::value,
            enthalpy::specific::value>,

        // R_h(h, Phi, d, C)
        block_row<
            enthalpy::specific::residual,
            enthalpy::specific::value,
            gravity::poisson::value,
            displacement::geometry::value,
            barotropic_constant::mass_normalization::value>,

        // R_M(rho, d)
        block_row<
            barotropic_constant::mass_normalization::residual,
            density::mass::value,
            displacement::geometry::value>>;

    static_assert(valid_jacobian_form<
                  gravity_field_form,
                  gravity_jacobian_form>);

    static_assert(valid_jacobian_form<
                  barotropic_equilibrium_form,
                  barotropic_equilibrium_jacobian_form>);
} // namespace mean_field::utils::blocks