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

#include <catch2/catch_test_macros.hpp>

import mean_field;
import test_helpers;

namespace {
    struct NotAModelSpecification final { };

    using PolytropicMassSpecifications = mean_field::models::
        SpecificationSet<mean_field::eos::Polytrope, mean_field::models::FixedTotalMass, mean_field::surface::Isobaric>;

    using PermutedPolytropicMassSpecifications = mean_field::models::
        SpecificationSet<mean_field::surface::Isobaric, mean_field::models::FixedTotalMass, mean_field::eos::Polytrope>;

    using CentralDensityPolytropicMassSpecifications = mean_field::models::SpecificationSet<
        mean_field::models::FixedCentralDensity,
        mean_field::surface::Isobaric,
        mean_field::eos::Polytrope,
        mean_field::models::FixedTotalMass>;

    using AngularMomentumPolytropicMassSpecifications = mean_field::models::SpecificationSet<
        mean_field::models::FixedAngularMomentum,
        mean_field::surface::Isobaric,
        mean_field::eos::Polytrope,
        mean_field::models::FixedTotalMass>;

    using PolytropicMassModel         = mean_field::model::StellarModel<PolytropicMassSpecifications>;
    using PermutedPolytropicMassModel = mean_field::model::StellarModel<PermutedPolytropicMassSpecifications>;
    using CentralDensityPolytropicMassModel =
        mean_field::model::StellarModel<CentralDensityPolytropicMassSpecifications>;
    using AngularMomentumPolytropicMassModel =
        mean_field::model::StellarModel<AngularMomentumPolytropicMassSpecifications>;
} // namespace

TEST_CASE(
    "Model Specifications Form Canonical Compile-Time Model Types",
    tags::model_specification_type_contract
) {
    STATIC_CHECK(mean_field::models::ModelSpecification<mean_field::eos::Polytrope>);
    STATIC_CHECK(mean_field::models::ModelSpecification<mean_field::surface::ConstantPressureSurface>);
    STATIC_CHECK(mean_field::models::ModelSpecification<mean_field::models::FixedTotalMass>);
    STATIC_CHECK(mean_field::models::ModelSpecification<mean_field::models::FixedAngularMomentum>);
    STATIC_CHECK(mean_field::models::ModelSpecification<mean_field::models::FixedCentralDensity>);
    STATIC_CHECK_FALSE(mean_field::models::ModelSpecification<NotAModelSpecification>);
    STATIC_CHECK(mean_field::models::ResolvedModelSpecification<mean_field::eos::Polytrope>);
    STATIC_CHECK(mean_field::models::ResolvedModelSpecification<mean_field::surface::ConstantPressureSurface>);
    STATIC_CHECK(mean_field::models::ResolvedModelSpecification<mean_field::models::FixedTotalMass>);
    STATIC_CHECK(mean_field::models::ResolvedModelSpecification<mean_field::models::FixedAngularMomentum>);
    STATIC_CHECK(mean_field::models::ResolvedModelSpecification<mean_field::models::FixedCentralDensity>);

    STATIC_CHECK(
        mean_field::models::ValidModelSpecificationPack<
            mean_field::eos::Polytrope, mean_field::models::FixedTotalMass, mean_field::surface::Isobaric>
    );

    STATIC_CHECK_FALSE(
        mean_field::models::ValidModelSpecificationPack<
            mean_field::eos::Polytrope, mean_field::models::FixedTotalMass, mean_field::models::FixedTotalMass,
            mean_field::surface::Isobaric>
    );

    STATIC_CHECK_FALSE(
        mean_field::models::ValidModelSpecificationPack<
            mean_field::models::FixedTotalMass, mean_field::surface::Isobaric>
    );

    STATIC_CHECK(std::same_as<PolytropicMassModel, PermutedPolytropicMassModel>);
    STATIC_CHECK_FALSE(std::same_as<PolytropicMassModel, CentralDensityPolytropicMassModel>);
    STATIC_CHECK_FALSE(std::same_as<PolytropicMassModel, AngularMomentumPolytropicMassModel>);
    STATIC_CHECK(mean_field::model::StellarModelType<PolytropicMassModel>);
    STATIC_CHECK(mean_field::model::StellarModelType<CentralDensityPolytropicMassModel>);
    STATIC_CHECK(mean_field::model::StellarModelType<AngularMomentumPolytropicMassModel>);
}

TEST_CASE(
    "Invariant And Phase Specifications Generate Balanced Residual And Value Types",
    tags::model_specification_type_contract
) {
    using MassSignature = PolytropicMassModel::OperatorSignature;

    STATIC_CHECK(MassSignature::generatedValueArity == 1);
    STATIC_CHECK(MassSignature::generatedResidualArity == 1);
    STATIC_CHECK(MassSignature::symbolicallySquare);

    STATIC_CHECK(
        mean_field::models::modelTypeListContains<
            mean_field::models::MultiplierFor<mean_field::models::FixedTotalMass>,
            typename MassSignature::GeneratedValues>
    );

    using AngularMomentumSignature = AngularMomentumPolytropicMassModel::OperatorSignature;

    STATIC_CHECK(AngularMomentumSignature::generatedValueArity == 2);
    STATIC_CHECK(AngularMomentumSignature::generatedResidualArity == 2);
    STATIC_CHECK(AngularMomentumSignature::symbolicallySquare);
    STATIC_CHECK(
        mean_field::models::modelTypeListContains<
            mean_field::models::PhysicalCoordinateFor<mean_field::models::FixedAngularMomentum>,
            typename AngularMomentumSignature::GeneratedValues>
    );
    STATIC_CHECK(
        mean_field::models::modelTypeListContains<
            mean_field::models::ResidualFor<mean_field::models::FixedAngularMomentum>,
            typename AngularMomentumSignature::GeneratedResiduals>
    );

    STATIC_CHECK(
        mean_field::models::modelTypeListContains<
            mean_field::models::ResidualFor<mean_field::models::FixedTotalMass>,
            typename MassSignature::GeneratedResiduals>
    );

    using CentralDensitySignature = CentralDensityPolytropicMassModel::OperatorSignature;

    STATIC_CHECK(CentralDensitySignature::generatedValueArity == 2);
    STATIC_CHECK(CentralDensitySignature::generatedResidualArity == 2);
    STATIC_CHECK(CentralDensitySignature::symbolicallySquare);

    STATIC_CHECK(
        mean_field::models::modelTypeListContains<
            mean_field::models::BorderFor<mean_field::models::FixedCentralDensity>,
            typename CentralDensitySignature::GeneratedValues>
    );

    STATIC_CHECK(
        mean_field::models::modelTypeListContains<
            mean_field::models::ResidualFor<mean_field::models::FixedCentralDensity>,
            typename CentralDensitySignature::GeneratedResiduals>
    );
}

TEST_CASE(
    "Fixed Angular Momentum Compiles A Physical Angular Velocity And Invariant Row",
    tags::model_specification_type_contract
) {
    using namespace mean_field;
    using Request = models::FixedAngularMomentumLayoutRequest;
    using Form = operators::CompiledStellarEquilibriumForm<AngularMomentumPolytropicMassModel>;
    using Jacobian = operators::CompiledStellarEquilibriumJacobianForm<AngularMomentumPolytropicMassModel>;
    using AngularValue = utils::blocks::fixed_angular_momentum::angular_velocity::value;
    using AngularResidual = utils::blocks::fixed_angular_momentum::angular_velocity::residual;

    STATIC_CHECK(models::ConstraintLayoutRequestType<Request>);
    STATIC_CHECK(models::CompiledConstraint<models::CompiledFixedAngularMomentum>);
    STATIC_CHECK(std::same_as<
                 typename Request::GeneratedValueType,
                 models::PhysicalCoordinateFor<models::FixedAngularMomentum>>);
    STATIC_CHECK(std::same_as<
                 typename AngularValue::GeneratedType,
                 models::PhysicalCoordinateFor<models::FixedAngularMomentum>>);
    STATIC_CHECK(std::same_as<
                 typename models::CompiledFixedAngularMomentum::AngularVelocityField,
                 field::AngularVelocity>);
    STATIC_CHECK(Form::value_block_count == 7);
    STATIC_CHECK(Form::residual_block_count == 7);
    STATIC_CHECK(Request::valueBlock<Form>().index == 6);
    STATIC_CHECK(Request::residualBlock<Form>().index == 6);
    STATIC_CHECK(utils::blocks::valid_jacobian_form<Form, Jacobian>);
    STATIC_CHECK(utils::blocks::has_jacobian_coupling_v<
                 AngularResidual,
                 utils::blocks::density::mass::value,
                 Jacobian>);
    STATIC_CHECK(utils::blocks::has_jacobian_coupling_v<
                 AngularResidual,
                 utils::blocks::surface_deformation::parameters::value,
                 Jacobian>);
    STATIC_CHECK(utils::blocks::has_jacobian_coupling_v<AngularResidual, AngularValue, Jacobian>);
    STATIC_CHECK(utils::blocks::has_jacobian_coupling_v<
                 utils::blocks::surface_deformation::shape_equilibrium::residual,
                 AngularValue,
                 Jacobian>);
    STATIC_CHECK(utils::blocks::has_jacobian_coupling_v<
                 utils::blocks::enthalpy::specific::residual,
                 AngularValue,
                 Jacobian>);
    STATIC_CHECK_FALSE(utils::blocks::has_jacobian_coupling_v<
                       utils::blocks::gravity::poisson::residual,
                       AngularValue,
                       Jacobian>);

    const integral::FixedAngularMomentum specification({
        .Jtotal = dimensions::AngularMomentumValue{2.75},
        .axis = {0.0, 3.0, 4.0},
        .center = {0.25, -0.5, 0.75}
    });
    const models::CompiledFixedAngularMomentum compiled = models::compileConstraint(specification);
    CHECK(compiled.targetAngularMomentum() == dimensions::AngularMomentumValue{2.75});
    CHECK(std::abs(compiled.specification().axis()[0]) < 1.0e-15);
    CHECK(std::abs(compiled.specification().axis()[1] - 0.6) < 1.0e-15);
    CHECK(std::abs(compiled.specification().axis()[2] - 0.8) < 1.0e-15);
    CHECK(compiled.specification().center() == std::array<double, 3>{0.25, -0.5, 0.75});

    const physics::RigidRotation rotation = compiled.makeRotation(1.5);
    CHECK(std::abs(rotation.angular_velocity()(0)) < 1.0e-15);
    CHECK(std::abs(rotation.angular_velocity()(1) - 0.9) < 1.0e-15);
    CHECK(std::abs(rotation.angular_velocity()(2) - 1.2) < 1.0e-15);
    CHECK(rotation.center()(0) == 0.25);
    CHECK(rotation.center()(1) == -0.5);
    CHECK(rotation.center()(2) == 0.75);
}

TEST_CASE(
    "Fixed Total Mass Compiles Its Generated Multiplier And Canonical Residual Row",
    tags::model_specification_type_contract
) {
    using namespace mean_field;
    using Request = models::FixedMassLayoutRequest;
    using Form    = utils::blocks::barotropic_equilibrium_form;

    STATIC_CHECK(models::ConstraintLayoutRequestType<Request>);
    STATIC_CHECK(models::CompiledConstraint<models::CompiledFixedMass>);
    STATIC_CHECK(std::same_as<typename Request::SpecificationType, models::FixedTotalMass>);
    STATIC_CHECK(std::same_as<typename Request::GeneratedValueType, models::MultiplierFor<models::FixedTotalMass>>);
    STATIC_CHECK(std::same_as<typename Request::GeneratedResidualType, models::ResidualFor<models::FixedTotalMass>>);
    STATIC_CHECK(
        std::same_as<typename Request::ValueBlockType::GeneratedType, models::MultiplierFor<models::FixedTotalMass>>
    );
    STATIC_CHECK(
        std::same_as<typename Request::ResidualBlockType::GeneratedType, models::ResidualFor<models::FixedTotalMass>>
    );
    STATIC_CHECK(std::same_as<typename models::CompiledFixedMass::MultiplierField, field::BarotropicConstant>);
    STATIC_CHECK(Request::rowInjection == models::ConstraintRowInjection::append);
    STATIC_CHECK(Request::valueArity == 1);
    STATIC_CHECK(Request::residualArity == 1);
    STATIC_CHECK(Request::valueBlock<Form>().index == Form::value_block_count - 1);
    STATIC_CHECK(Request::residualBlock<Form>().index == Form::residual_block_count - 1);

    const models::CompiledFixedMass compiled =
        models::compileConstraint(models::FixedTotalMass{dimensions::MassValue{1.75}});
    CHECK(compiled.targetMass() == dimensions::MassValue{1.75});
    CHECK(compiled.specification().targetMass() == dimensions::MassValue{1.75});
}

TEST_CASE(
    "Model Specification Descriptors Preserve Roles And Generated Arity",
    tags::model_specification_type_contract
) {
    constexpr auto polytrope = mean_field::models::specificationDescriptor<mean_field::eos::Polytrope>();
    constexpr auto surface =
        mean_field::models::specificationDescriptor<mean_field::surface::ConstantPressureSurface>();
    constexpr auto mass = mean_field::models::specificationDescriptor<mean_field::models::FixedTotalMass>();
    constexpr auto angularMomentum =
        mean_field::models::specificationDescriptor<mean_field::models::FixedAngularMomentum>();
    constexpr auto centralDensity =
        mean_field::models::specificationDescriptor<mean_field::models::FixedCentralDensity>();

    STATIC_CHECK(polytrope.name == "Polytrope");
    STATIC_CHECK(polytrope.role == mean_field::models::SpecificationRole::constitutive_law);
    STATIC_CHECK(polytrope.generatedValueArity == 0);
    STATIC_CHECK(polytrope.generatedResidualArity == 0);

    STATIC_CHECK(surface.name == "IsobaricSurface");
    STATIC_CHECK(surface.role == mean_field::models::SpecificationRole::boundary_condition);

    STATIC_CHECK(mass.name == "FixedTotalMass");
    STATIC_CHECK(mass.role == mean_field::models::SpecificationRole::invariant);
    STATIC_CHECK(mass.generatedValueArity == 1);
    STATIC_CHECK(mass.generatedResidualArity == 1);

    STATIC_CHECK(angularMomentum.name == "FixedAngularMomentum");
    STATIC_CHECK(angularMomentum.role == mean_field::models::SpecificationRole::invariant);
    STATIC_CHECK(angularMomentum.generatedValueArity == 1);
    STATIC_CHECK(angularMomentum.generatedResidualArity == 1);

    STATIC_CHECK(centralDensity.name == "FixedCentralDensity");
    STATIC_CHECK(centralDensity.role == mean_field::models::SpecificationRole::phase_condition);
    STATIC_CHECK(centralDensity.generatedValueArity == 1);
    STATIC_CHECK(centralDensity.generatedResidualArity == 1);
}

TEST_CASE(
    "Invariant And Phase Specification Values Reject Invalid Targets",
    tags::model_specification_type_contract
) {
    const mean_field::models::FixedTotalMass mass{mean_field::dimensions::MassValue{1.25}};
    const mean_field::models::FixedCentralDensity centralDensity{mean_field::eos::DensityValue{2.5}};
    const mean_field::models::FixedAngularMomentum angularMomentum{
        mean_field::dimensions::AngularMomentumValue{0.75}
    };

    CHECK(mass.targetMass() == mean_field::dimensions::MassValue{1.25});
    CHECK(centralDensity.targetDensity() == mean_field::eos::DensityValue{2.5});
    CHECK(angularMomentum.targetAngularMomentum() == mean_field::dimensions::AngularMomentumValue{0.75});
    CHECK(angularMomentum.axis() == std::array<double, 3>{0.0, 0.0, 1.0});
    CHECK(angularMomentum.center() == std::array<double, 3>{0.0, 0.0, 0.0});

    CHECK_THROWS_AS(mean_field::models::FixedTotalMass{mean_field::dimensions::MassValue{0.0}}, std::invalid_argument);
    CHECK_THROWS_AS(mean_field::models::FixedTotalMass{mean_field::dimensions::MassValue{-1.0}}, std::invalid_argument);
    CHECK_THROWS_AS(
        mean_field::models::FixedTotalMass{mean_field::dimensions::MassValue{std::numeric_limits<double>::infinity()}},
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        mean_field::models::FixedAngularMomentum{mean_field::dimensions::AngularMomentumValue{-1.0}},
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        mean_field::models::FixedAngularMomentum{
            mean_field::dimensions::AngularMomentumValue{std::numeric_limits<double>::infinity()}
        },
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        mean_field::models::FixedAngularMomentum({
            .Jtotal = mean_field::dimensions::AngularMomentumValue{1.0},
            .axis = {0.0, 0.0, 0.0}
        }),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        mean_field::models::FixedAngularMomentum({
            .Jtotal = mean_field::dimensions::AngularMomentumValue{1.0},
            .axis = {0.0, 0.0, 1.0},
            .center = {0.0, std::numeric_limits<double>::quiet_NaN(), 0.0}
        }),
        std::invalid_argument
    );

    CHECK_THROWS_AS(mean_field::models::FixedCentralDensity{mean_field::eos::DensityValue{0.0}}, std::invalid_argument);
    CHECK_THROWS_AS(
        mean_field::models::FixedCentralDensity{mean_field::eos::DensityValue{-1.0}}, std::invalid_argument
    );
    CHECK_THROWS_AS(
        mean_field::models::FixedCentralDensity{mean_field::eos::DensityValue{std::numeric_limits<double>::infinity()}},
        std::invalid_argument
    );
}