#pragma once

#include <array>
#include <cstddef>
#include <string_view>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <stroid/stroid.h>

#include "serif/discretization/domain/ids/boundary.hpp"
#include "serif/discretization/domain/ids/domain.hpp"
#include "serif/discretization/domain/ids/lists/lists.hpp"
#include "serif/discretization/domain/physical_domains.hpp"
#include "serif/discretization/domain/relation/lists/relation_list.hpp"
#include "serif/discretization/domain/relation/relations.hpp"
#include "serif/discretization/domain/schema/domain_schema.hpp"
#include "serif/discretization/domain/schema/validation/all.hpp"
#include "serif/discretization/domain/types.hpp"

namespace domain_test_utils {
    namespace domain     = serif::discretization::domain;
    namespace ids        = domain::ids;
    namespace relation   = domain::relation;
    namespace schema     = domain::schema;
    namespace validation = schema::validation;

    struct UnregisteredDomain final : public domain::Domain {
        static constexpr std::string_view name = "unregistered_domain";
    };

    struct UnregisteredBoundary final : public domain::Boundary {
        static constexpr std::string_view name = "unregistered_boundary";
    };

    struct BoundaryEdge {
        int firstVertexId{-1};
        int secondVertexId{-1};
        int attribute{0};
    };

    struct StroidCase {
        std::string_view name;
        int refinementLevels{0};
        int order{1};
        double flattening{0.0};
    };

    template <typename... DomainIDTs>
    concept CanFormDomainIDList = requires { typename ids::lists::DomainIDList<DomainIDTs...>; };

    template <typename... BoundaryIDTs>
    concept CanFormBoundaryIDList = requires { typename ids::lists::BoundaryIDList<BoundaryIDTs...>; };

    template <typename BoundaryT, typename... DomainTs>
    concept CanFormDomainBoundary = requires { typename relation::DomainBoundary<BoundaryT, DomainTs...>; };

    template <typename DomainIDsT, typename BoundaryIDsT, typename RelationsT>
    concept CanFormSchema = requires { typename schema::DomainSchema<DomainIDsT, BoundaryIDsT, RelationsT>; };

    [[nodiscard]] inline int vertex_id(const int xElementCount, const int x, const int y) {
        return y * (xElementCount + 1) + x;
    }

    [[nodiscard]] inline int cell_index(const int xElementCount, const int x, const int y) {
        return y * xElementCount + x;
    }

    [[nodiscard]]inline int cell_attribute(const std::vector<int> &attributes, const int xElementCount, const int x, const int y) {
        return attributes.at(static_cast<std::size_t>(cell_index(xElementCount, x, y)));
    }

    template <typename FirstPredicateT, typename SecondPredicateT>
    void append_interface_boundaries(
        std::vector<BoundaryEdge> &boundaries,
        const std::vector<int> &attributes,
        const int xElementCount,
        const int yElementCount,
        FirstPredicateT firstPredicate,
        SecondPredicateT secondPredicate,
        const int boundaryAttribute
    ) {
        /*
         * Vertical internal faces.
         */
        for (int y = 0; y < yElementCount; ++y) {
            for (int x = 1; x < xElementCount; ++x) {
                const int leftAttribute  = cell_attribute(attributes, xElementCount, x - 1, y);
                const int rightAttribute = cell_attribute(attributes, xElementCount, x, y);
                const bool matches       = (firstPredicate(leftAttribute) && secondPredicate(rightAttribute)) ||
                                           (secondPredicate(leftAttribute) && firstPredicate(rightAttribute));

                if (!matches) {
                    continue;
                }

                boundaries.push_back(
                    {.firstVertexId  = vertex_id(xElementCount, x, y),
                     .secondVertexId = vertex_id(xElementCount, x, y + 1),
                     .attribute      = boundaryAttribute}
                );
            }
        }

        /*
         * Horizontal internal faces.
         */
        for (int y = 1; y < yElementCount; ++y) {
            for (int x = 0; x < xElementCount; ++x) {
                const int lowerAttribute = cell_attribute(attributes, xElementCount, x, y - 1);
                const int upperAttribute = cell_attribute(attributes, xElementCount, x, y);
                const bool matches       = (firstPredicate(lowerAttribute) && secondPredicate(upperAttribute)) ||
                                           (secondPredicate(lowerAttribute) && firstPredicate(upperAttribute));

                if (!matches) {
                    continue;
                }

                boundaries.push_back(
                    {.firstVertexId  = vertex_id(xElementCount, x, y),
                     .secondVertexId = vertex_id(xElementCount, x + 1, y),
                     .attribute      = boundaryAttribute}
                );
            }
        }
    }

    template <typename PredicateT>
    void append_exterior_boundaries(
        std::vector<BoundaryEdge> &boundaries,
        const std::vector<int> &attributes,
        const int xElementCount,
        const int yElementCount,
        PredicateT predicate,
        const int boundaryAttribute
    ) {
        /*
         * Bottom.
         */
        for (int x = 0; x < xElementCount; ++x) {
            if (predicate(cell_attribute(attributes, xElementCount, x, 0))) {
                boundaries.push_back(
                    {.firstVertexId  = vertex_id(xElementCount, x, 0),
                     .secondVertexId = vertex_id(xElementCount, x + 1, 0),
                     .attribute      = boundaryAttribute}
                );
            }
        }

        /*
         * Top.
         */
        for (int x = 0; x < xElementCount; ++x) {
            if (predicate(cell_attribute(attributes, xElementCount, x, yElementCount - 1))) {
                boundaries.push_back(
                    {.firstVertexId  = vertex_id(xElementCount, x, yElementCount),
                     .secondVertexId = vertex_id(xElementCount, x + 1, yElementCount),
                     .attribute      = boundaryAttribute}
                );
            }
        }

        /*
         * Left.
         */
        for (int y = 0; y < yElementCount; ++y) {
            if (predicate(cell_attribute(attributes, xElementCount, 0, y))) {
                boundaries.push_back(
                    {.firstVertexId  = vertex_id(xElementCount, 0, y),
                     .secondVertexId = vertex_id(xElementCount, 0, y + 1),
                     .attribute      = boundaryAttribute}
                );
            }
        }

        /*
         * Right.
         */
        for (int y = 0; y < yElementCount; ++y) {
            if (predicate(cell_attribute(attributes, xElementCount, xElementCount - 1, y))) {
                boundaries.push_back(
                    {.firstVertexId  = vertex_id(xElementCount, xElementCount, y),
                     .secondVertexId = vertex_id(xElementCount, xElementCount, y + 1),
                     .attribute      = boundaryAttribute}
                );
            }
        }
    }

    [[nodiscard]] inline mfem::Mesh make_grid_mesh(
        const int xElementCount,
        const int yElementCount,
        const std::vector<int> &attributes,
        const std::vector<BoundaryEdge> &boundaryEdges
    ) {
        REQUIRE(static_cast<int>(attributes.size()) == xElementCount * yElementCount);

        mfem::Mesh mesh(
            2, (xElementCount + 1) * (yElementCount + 1), xElementCount * yElementCount,
            static_cast<int>(boundaryEdges.size()), 2
        );

        for (int y = 0; y <= yElementCount; ++y) {
            for (int x = 0; x <= xElementCount; ++x) {
                mesh.AddVertex(static_cast<double>(x), static_cast<double>(y));
            }
        }

        for (int y = 0; y < yElementCount; ++y) {
            for (int x = 0; x < xElementCount; ++x) {
                const int lowerLeft  = vertex_id(xElementCount, x, y);
                const int lowerRight = vertex_id(xElementCount, x + 1, y);
                const int upperRight = vertex_id(xElementCount, x + 1, y + 1);
                const int upperLeft  = vertex_id(xElementCount, x, y + 1);

                mesh.AddQuad(
                    lowerLeft, lowerRight, upperRight, upperLeft, cell_attribute(attributes, xElementCount, x, y)
                );
            }
        }

        for (const BoundaryEdge &boundary : boundaryEdges) {
            mesh.AddBdrSegment(boundary.firstVertexId, boundary.secondVertexId, boundary.attribute);
        }

        mesh.FinalizeTopology(false);
        mesh.Finalize(false, false);

        REQUIRE(mesh.GetNBE() == static_cast<int>(boundaryEdges.size()));

        return mesh;
    }

    [[nodiscard]] inline std::vector<int> make_layered_attributes() {
        constexpr int xElementCount = 5;
        constexpr int yElementCount = 5;

        std::vector<int> attributes(xElementCount * yElementCount, 3);

        for (int y = 1; y <= 3; ++y) {
            for (int x = 1; x <= 3; ++x) {
                attributes[static_cast<std::size_t>(cell_index(xElementCount, x, y))] = 2;
            }
        }

        attributes[static_cast<std::size_t>(cell_index(xElementCount, 2, 2))] = 1;

        return attributes;
    }

    [[nodiscard]] inline mfem::Mesh make_layered_mesh(
        const bool includeStellarSurface   = true,
        const bool includeInfinitySurface  = true,
        const int stellarSurfaceAttribute  = 1,
        const int infinitySurfaceAttribute = 2
    ) {
        constexpr int xElementCount       = 5;
        constexpr int yElementCount       = 5;

        const std::vector<int> attributes = make_layered_attributes();
        std::vector<BoundaryEdge> boundaries;

        const auto isStellar = [](const int materialId) { return materialId == 1 || materialId == 2; };
        const auto isVacuum  = [](const int materialId) { return materialId == 3; };

        if (includeStellarSurface) {
            append_interface_boundaries(
                boundaries, attributes, xElementCount, yElementCount, isStellar, isVacuum, stellarSurfaceAttribute
            );
        }

        if (includeInfinitySurface) {
            append_exterior_boundaries(
                boundaries, attributes, xElementCount, yElementCount, isVacuum, infinitySurfaceAttribute
            );
        }

        return make_grid_mesh(xElementCount, yElementCount, attributes, boundaries);
    }

    template <typename SchemaT>
    void check_schema_is_valid(const mfem::Mesh &mesh) {
        const auto validation = domain::schema::validation::validate_schema<SchemaT>(mesh);

        CHECK(validation.relationResults.size() == SchemaT::relation_count);

        for (const auto &relationResult : validation.relationResults) {
            INFO("Relation index = " << relationResult.relationIndex);
            INFO("Relation name = " << relationResult.relationName);
            INFO("Failure enum = " << static_cast<int>(relationResult.result.failure));
            CHECK(relationResult.valid());
        }

        CHECK(validation.valid());
    }

    using AlternateIdSchema = schema::DomainSchema<
        ids::lists::DomainIDList<
            ids::DomainID<domain::CoreDomain, 11>,
            ids::DomainID<domain::EnvelopeDomain, 17>,
            ids::DomainID<domain::VacuumDomain, 29>>,
        ids::lists::BoundaryIDList<
            ids::BoundaryID<domain::StellarSurfaceBoundary, 101>,
            ids::BoundaryID<domain::InfinitySurfaceBoundary, 203>>,
        relation::lists::RelationList<
            relation::FullyConnected<domain::CoreDomain>,
            relation::FullyConnected<domain::EnvelopeDomain>,
            relation::FullyConnected<domain::VacuumDomain>,
            relation::Inscribed<domain::CoreDomain, domain::EnvelopeDomain>,
            relation::Inscribed<domain::StellarDomains, domain::VacuumDomain>,
            relation::DomainBoundary<
                domain::StellarSurfaceBoundary,
                domain::StellarDomains,
                domain::VacuumDomain>,
            relation::DomainBoundary<domain::InfinitySurfaceBoundary, domain::VacuumDomain>>>;

    [[nodiscard]] inline stroid::config::MeshConfig make_stroid_config(
        const int refinementLevels,
        const int order,
        const double flattening
    ) {
        stroid::config::MeshConfig config;
        config.refinement_levels       = refinementLevels;
        config.order                   = order;
        config.include_external_domain = true;
        config.r_core                  = 0.25;
        config.r_star                  = 1.0;
        config.r_infinity              = 4.0;
        config.flattening              = flattening;
        config.core_id                 = 1;
        config.envelope_id             = 2;
        config.vacuum_id               = 3;
        config.surface_bdr_id          = 1;
        config.inf_bdr_id              = 2;
        config.optimization_methods    = stroid::config::OptimizationMethods{.tmop = false, .smoothstep = true};
        return config;
    }
} // namespace domain_test_utils