#include <array>
#include <limits>
#include <utility>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>

#include "serif/discretization/blocks/fields.hpp"
#include "serif/model/state.hpp"

namespace {
    namespace blocks = serif::discretization::blocks;
    namespace dependencies = serif::utils::dependencies;
    using Density = blocks::density::mass;
    using Shape = blocks::surface_deformation::coefficients;
    using ShapeRow = blocks::surface_deformation::shape_equilibrium;
    using Constant = blocks::barotropic_constant::mass_normalization;
    using Layout = blocks::SystemLayout<dependencies::KeyList<Density, Shape, Constant>, dependencies::KeyList<ShapeRow, Constant, Density>>;
    using State = serif::model::StellarState<Layout>;
}

TEST_CASE("A stellar state owns its layout and exposes fields in physical rather than residual order", "[model][state][unit]") {
    Layout layout{.state_sizes = {2, 3, 1}, .residual_sizes = {3, 1, 2}};
    State state(layout, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0});
    layout.state_sizes[0] = 9;
    layout.residual_sizes[0] = 8;

    CHECK(state.layout().state_sizes == std::array<std::size_t, 3>{2, 3, 1});
    CHECK(state.layout().residual_sizes == std::array<std::size_t, 3>{3, 1, 2});
    CHECK(state.layout().offset<blocks::RepresentedValueKind::residual, Density>() == 4);
    CHECK(state.get<Density>().data() == state.values().data());
    CHECK(state.get<Shape>().data() == state.values().data() + 2);
    CHECK(state.get<Constant>().data() == state.values().data() + 5);
    REQUIRE(state.get<Shape>().size() == 3);
    state.get<Shape>()[1] = 7.0;
    state.get<Constant>()[0] = 8.0;
    CHECK(state.values()[3] == 7.0);
    CHECK(state.values()[5] == 8.0);

    const State& read_only = state;
    CHECK(read_only.get<Shape>()[1] == 7.0);
    CHECK(read_only.values().data() == state.values().data());
}

TEST_CASE("Copied and moved stellar candidates retain independent coefficients and valid typed access", "[model][state][ownership]") {
    const Layout layout{.state_sizes = {2, 3, 1}, .residual_sizes = {3, 1, 2}};
    State original(layout, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0});
    State candidate = original;
    CHECK(candidate.values().data() != original.values().data());
    candidate.get<Density>()[0] = 9.0;
    candidate.get<Shape>()[2] = 10.0;
    CHECK(original.get<Density>()[0] == 1.0);
    CHECK(original.get<Shape>()[2] == 5.0);

    State moved(std::move(candidate));
    CHECK(moved.layout() == layout);
    CHECK(moved.get<Density>()[0] == 9.0);
    CHECK(moved.get<Shape>()[2] == 10.0);
    CHECK(moved.get<Shape>().data() == moved.values().data() + 2);
    moved.get<Constant>()[0] = 11.0;
    CHECK(original.get<Constant>()[0] == 6.0);

    State assigned(Layout{.state_sizes = {1, 1, 1}, .residual_sizes = {1, 1, 1}});
    assigned = std::move(moved);
    CHECK(assigned.layout() == layout);
    CHECK(assigned.get<Constant>()[0] == 11.0);
    CHECK(assigned.get<Constant>().data() == assigned.values().data() + 5);
    assigned = original;
    assigned.get<Density>()[1] = 12.0;
    CHECK(original.get<Density>()[1] == 2.0);
}

TEST_CASE("Stellar states reject inconsistent storage and overflowing sizes while supporting empty coordinates", "[model][state][validation]") {
    const Layout layout{.state_sizes = {2, 3, 1}, .residual_sizes = {3, 1, 2}};
    const State zero(layout);
    REQUIRE(zero.values().size() == 6);
    for (double value : zero.values()) {
        CHECK(value == 0.0);
    }
    CHECK_THROWS_WITH(State(layout, std::vector<double>(5)), Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));
    CHECK_THROWS_WITH(State(layout, std::vector<double>(7)), Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));
    const Layout overflow{.state_sizes = {std::numeric_limits<std::size_t>::max(), 1, 0}, .residual_sizes = {0, 1, 0}};
    CHECK_THROWS_WITH(State(overflow), Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));

    const State partial(Layout{.state_sizes = {0, 2, 0}, .residual_sizes = {2, 0, 0}});
    CHECK(partial.get<Density>().empty());
    CHECK(partial.get<Shape>().size() == 2);
    CHECK(partial.get<Constant>().empty());
    const blocks::SystemLayout<dependencies::KeyList<>, dependencies::KeyList<>> empty_layout{};
    const serif::model::StellarState empty(empty_layout);
    CHECK(empty.values().empty());
}