#include <array>
#include <concepts>
#include <stdexcept>
#include <type_traits>

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

import mean_field;
import test_helpers;

namespace {
    using CanonicalModel = mean_field::models::
        Model<mean_field::eos::Polytrope, mean_field::models::FixedTotalMass, mean_field::surface::Isobaric>;

    using PermutedModel = mean_field::models::
        Model<mean_field::models::FixedTotalMass, mean_field::surface::Isobaric, mean_field::eos::Polytrope>;

    using CentralDensityModel = mean_field::models::Model<
        mean_field::models::FixedCentralDensity,
        mean_field::surface::Isobaric,
        mean_field::eos::Polytrope,
        mean_field::models::FixedTotalMass>;

    using Form                = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form;
    using JacobianForm        = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_jacobian_form;
    using Manifest            = mean_field::operators::CompiledRootManifest<CanonicalModel, Form, JacobianForm>;
    using CentralForm         = mean_field::utils::blocks::central_density_bordered_stellar_equilibrium_form;
    using CentralJacobianForm = mean_field::utils::blocks::central_density_bordered_stellar_equilibrium_jacobian_form;
    using CentralManifest =
        mean_field::operators::CompiledRootManifest<CentralDensityModel, CentralForm, CentralJacobianForm>;

    [[nodiscard]] Manifest make_manifest() {
        const std::array<int, Form::value_block_count> valueSizes{2, 3, 4, 5, 6, 1};
        const std::array<int, Form::residual_block_count> residualSizes{4, 5, 2, 3, 6, 1};
        return {valueSizes, residualSizes, 2.5, 0.125, 3};
    }
} // namespace

TEST_CASE(
    "Model Values Are Stored In Canonical Specification Order",
    tags::model_specification_type_contract
) {
    STATIC_CHECK(std::same_as<CanonicalModel, PermutedModel>);
    STATIC_CHECK(CanonicalModel::symbolicallySquare);
    STATIC_CHECK(CanonicalModel::hasCompleteRootCompiler);
    STATIC_CHECK(CanonicalModel::compilationClass == mean_field::models::ModelCompilationClass::isolated_root);
    STATIC_CHECK(CentralDensityModel::symbolicallySquare);
    STATIC_CHECK(CentralDensityModel::hasCompleteRootCompiler);
    STATIC_CHECK(CentralDensityModel::compilationClass == mean_field::models::ModelCompilationClass::isolated_root);

    const mean_field::eos::Polytrope equationOfState{2.0, 0.75};
    const mean_field::surface::Isobaric surface{mean_field::dimensions::PressureValue{0.125}};
    const mean_field::models::FixedTotalMass mass{mean_field::dimensions::MassValue{1.75}};

    const CanonicalModel model{mass, surface, equationOfState};

    CHECK(model.specification<mean_field::eos::Polytrope>().polytropic_index() == 2.0);
    CHECK(model.specification<mean_field::surface::Isobaric>().targetPressure().value() == 0.125);
    CHECK(
        model.specification<mean_field::models::FixedTotalMass>().targetMass() ==
        mean_field::dimensions::MassValue{1.75}
    );

    const auto descriptors = model.runtimeSpecificationDescriptors();
    REQUIRE(descriptors.size() == 3);
    CHECK(descriptors[0].specification.name == "Polytrope");
    CHECK(descriptors[1].specification.name == "IsobaricSurface");
    CHECK(descriptors[2].specification.name == "FixedTotalMass");
    CHECK(descriptors[0].canonicalIndex == 0);
    CHECK(descriptors[1].canonicalIndex == 1);
    CHECK(descriptors[2].canonicalIndex == 2);
    CHECK(descriptors[0].hasRootCompiler);
    CHECK(descriptors[1].hasRootCompiler);
    CHECK(descriptors[2].hasRootCompiler);
}

TEST_CASE(
    "Central Density Root Manifest Appends A Carrier Phase Row And Solver Border",
    tags::root_manifest_type_contract
) {
    const std::array<int, CentralForm::value_block_count> valueSizes{2, 3, 4, 5, 6, 1, 1};
    const std::array<int, CentralForm::residual_block_count> residualSizes{4, 5, 2, 3, 6, 1, 1};
    const CentralManifest manifest(
        valueSizes, residualSizes, 2.5, 0.125, 3,
        mean_field::operators::CentralDensityManifestInput{
            .targetDensity = 8.0, .targetEnthalpy = 2.0, .centerDofCount = 1
        }
    );

    CHECK(manifest.layout().value_offsets().Last() == 22);
    CHECK(manifest.layout().residual_offsets().Last() == 22);
    REQUIRE(manifest.valueBlocks().size() == 7);
    REQUIRE(manifest.residualBlocks().size() == 7);
    CHECK(manifest.valueBlocks()[6].stableId == "fixed_central_density.border");
    CHECK(manifest.valueBlocks()[6].symbol == "lambda_rho_c");
    CHECK(manifest.valueBlocks()[6].columnPolicy == mean_field::operators::RootColumnPolicy::solver_border);
    CHECK(manifest.residualBlocks()[6].stableId == "fixed_central_density.residual");
    CHECK(manifest.residualBlocks()[6].symbol == "R_rho_c");
    CHECK(manifest.residualBlocks()[6].scale == 2.0);

    const auto constraints = manifest.constraints();
    REQUIRE(constraints.size() == 3);
    CHECK(constraints[2].stableId == "FixedCentralDensity");
    CHECK(constraints[2].role == mean_field::models::SpecificationRole::phase_condition);
    CHECK(constraints[2].valueBlock == 6);
    CHECK(constraints[2].residualBlock == 6);
    CHECK(constraints[2].target == 8.0);
    REQUIRE(constraints[2].carrierTarget.has_value());
    CHECK(*constraints[2].carrierTarget == 2.0);
    CHECK(constraints[2].residualScale == 2.0);
}

TEST_CASE(
    "Compiled Root Manifest Centralizes Canonical Blocks Provenance And Scaling",
    tags::root_manifest_type_contract
) {
    const Manifest manifest = make_manifest();

    STATIC_CHECK(Manifest::symbolicallySquare);
    STATIC_CHECK(Manifest::compilationClass == mean_field::models::ModelCompilationClass::isolated_root);

    CHECK(manifest.layout().value_offsets().Last() == 21);
    CHECK(manifest.layout().residual_offsets().Last() == 21);

    const auto values    = manifest.valueBlocks();
    const auto residuals = manifest.residualBlocks();

    REQUIRE(values.size() == 6);
    REQUIRE(residuals.size() == 6);
    CHECK(values[0].stableId == "density");
    CHECK(values[0].symbol == "rho");
    CHECK(values[1].stableId == "surface_deformation");
    CHECK(values[5].stableId == "fixed_total_mass.multiplier");
    CHECK(values[5].symbol == "C");
    CHECK(values[5].provenance == mean_field::operators::RootBlockProvenance::model_specification);
    CHECK(values[5].source == "FixedTotalMass");
    CHECK(values[5].columnPolicy == mean_field::operators::RootColumnPolicy::existing_physical_multiplier);

    CHECK(residuals[5].stableId == "fixed_total_mass.residual");
    CHECK(residuals[5].symbol == "R_M");
    CHECK(residuals[5].rowInjection == mean_field::operators::RootRowInjection::append_global);
    CHECK(residuals[5].scalePolicy == mean_field::operators::RootScalePolicy::target_relative);
    CHECK(residuals[5].scale == 2.5);

    const auto replacements = manifest.rowReplacements();
    REQUIRE(replacements.size() == 1);
    CHECK(replacements[0].sourceSpecification == "IsobaricSurface");
    CHECK(replacements[0].replacedRowCount == 3);
    CHECK(replacements[0].carrierResidualBlock == 4);

    const auto constraints = manifest.constraints();
    REQUIRE(constraints.size() == 2);
    CHECK(constraints[0].stableId == "FixedTotalMass");
    CHECK(constraints[0].valueBlock == 5);
    CHECK(constraints[0].residualBlock == 5);
    CHECK(constraints[0].target == 2.5);
    CHECK(constraints[0].residualScale == 2.5);
    CHECK(constraints[1].stableId == "IsobaricSurface");
    CHECK(constraints[1].rowInjection == mean_field::operators::RootRowInjection::replace_carrier_rows);
    CHECK(constraints[1].target == 0.125);
    CHECK_FALSE(constraints[1].carrierTarget.has_value());

    const auto report = manifest.fixedMassReport(2.75);
    CHECK(report.achieved == 2.75);
    CHECK(report.dimensionalResidual == 0.25);
    CHECK(report.scaledResidual == 0.1);
}

TEST_CASE(
    "Typed Root Views Resolve Blocks Through The Compiled Manifest",
    tags::root_manifest_type_contract
) {
    const Manifest manifest = make_manifest();

    mfem::Vector state(manifest.layout().value_offsets().Last());
    for (int index = 0; index < state.Size(); ++index) {
        state(index) = static_cast<double>(index + 1);
    }

    const auto stateView       = manifest.stateView(state);
    const mfem::Vector density = stateView.block(mean_field::utils::blocks::density_field.mass_term);
    const mfem::Vector surface = stateView.block(mean_field::utils::blocks::surface_deformation_field.parameters_term);
    const mfem::Vector multiplier =
        stateView.block(mean_field::utils::blocks::fixed_total_mass_constraint.mass_normalization_term);

    REQUIRE(density.Size() == 2);
    REQUIRE(surface.Size() == 3);
    REQUIRE(multiplier.Size() == 1);
    CHECK(density(0) == 1.0);
    CHECK(surface(0) == 3.0);
    CHECK(multiplier(0) == 21.0);

    mfem::Vector residual(manifest.layout().residual_offsets().Last());
    residual                = 0.0;
    const auto residualView = manifest.residualView(residual);
    mfem::Vector massResidual(1);
    massResidual(0) = -0.375;
    residualView.assign(mean_field::utils::blocks::fixed_total_mass_constraint.mass_normalization_term, massResidual);
    CHECK(residual(20) == -0.375);

    mfem::Vector wrongState(state.Size() - 1);
    CHECK_THROWS_AS(manifest.stateView(wrongState), std::invalid_argument);
}