#include "fesa/analysis/analysis_state.hpp"

#include <gtest/gtest.h>

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

namespace {

fesa::DofManager makeEmptyDofs() {
    fesa::ModelDefinition definition{};
    definition.sourcePath = "models/result-records.inp";
    definition.sourceContentIdentity = "fnv1a64:0123456789abcdef";
    definition.steps = {{
        "Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0,
        {definition.sourcePath, 10U}}};

    auto domain = fesa::Domain::create(std::move(definition));
    EXPECT_TRUE(domain.HasValue());
    auto model = fesa::AnalysisModel::create(domain.Value());
    EXPECT_TRUE(model.HasValue());
    auto dofs = fesa::DofManager::create(model.Value());
    EXPECT_TRUE(dofs.HasValue());
    return std::move(dofs.Value());
}

fesa::ShellResultRow makeShellRow(
    fesa::EntityIndex element,
    fesa::ShellMidsurfaceLocation location,
    std::array<double, 2> naturalCoordinates,
    double seed) {
    return {
        element,
        location,
        naturalCoordinates,
        {{{1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}, {0.0, 0.0, 1.0}}},
        {seed + 1.0,
         seed + 2.0,
         seed + 3.0,
         seed + 4.0,
         seed + 5.0,
         seed + 6.0,
         seed + 7.0,
         seed + 8.0},
        {seed + 11.0,
         seed + 12.0,
         seed + 13.0,
         seed + 14.0,
         seed + 15.0,
         seed + 16.0,
         seed + 17.0,
         seed + 18.0},
        {{{fesa::ShellSectionPosition::bottom,
           -1.0,
           {seed + 21.0, seed + 22.0, seed + 23.0}},
          {fesa::ShellSectionPosition::middle,
           0.0,
           {seed + 24.0, seed + 25.0, seed + 26.0}},
          {fesa::ShellSectionPosition::top,
           1.0,
           {seed + 27.0, seed + 28.0, seed + 29.0}}}}};
}

fesa::ShellStateCandidate makeShellCandidate(
    const std::vector<fesa::EntityIndex>& elements) {
    const double gauss = 1.0 / std::sqrt(3.0);
    const std::array<fesa::ShellMidsurfaceLocation, 4> locations{
        fesa::ShellMidsurfaceLocation::gp1,
        fesa::ShellMidsurfaceLocation::gp2,
        fesa::ShellMidsurfaceLocation::gp3,
        fesa::ShellMidsurfaceLocation::gp4};
    const std::array<std::array<double, 2>, 4> coordinates{
        std::array<double, 2>{-gauss, -gauss},
        std::array<double, 2>{gauss, -gauss},
        std::array<double, 2>{gauss, gauss},
        std::array<double, 2>{-gauss, gauss}};

    fesa::ShellStateCandidate candidate{};
    for (const auto element : elements) {
        for (std::size_t point = 0U; point < locations.size(); ++point) {
            candidate.rows.push_back(makeShellRow(
                element,
                locations[point],
                coordinates[point],
                100.0 * static_cast<double>(element) +
                    10.0 * static_cast<double>(point)));
        }
    }
    candidate.physicalStrainEnergy = 35.5;
    candidate.equilibrium = {1.0, -2.0, 3.0, -4.0, 5.0, -6.0};
    candidate.verificationMetrics = {1.0e-11, 2.0e-11, 3.0e-11};
    return candidate;
}

void expectShellStateEquals(
    const fesa::AnalysisState& state,
    const std::vector<fesa::ShellResultRow>& rows,
    double physicalStrainEnergy,
    const std::array<double, 6>& equilibrium,
    const std::array<double, 3>& verificationMetrics) {
    ASSERT_EQ(state.shellResults().size(), rows.size());
    for (std::size_t row = 0U; row < rows.size(); ++row) {
        EXPECT_EQ(state.shellResults()[row].element, rows[row].element);
        EXPECT_EQ(state.shellResults()[row].location, rows[row].location);
        EXPECT_EQ(
            state.shellResults()[row].naturalCoordinates,
            rows[row].naturalCoordinates);
        EXPECT_EQ(state.shellResults()[row].localFrame, rows[row].localFrame);
        EXPECT_EQ(
            state.shellResults()[row].generalizedStrain,
            rows[row].generalizedStrain);
        EXPECT_EQ(
            state.shellResults()[row].sectionResultant,
            rows[row].sectionResultant);
        for (std::size_t position = 0U;
             position < rows[row].stress.size();
             ++position) {
            EXPECT_EQ(
                state.shellResults()[row].stress[position].position,
                rows[row].stress[position].position);
            EXPECT_DOUBLE_EQ(
                state.shellResults()[row].stress[position].zeta,
                rows[row].stress[position].zeta);
            EXPECT_EQ(
                state.shellResults()[row].stress[position].components,
                rows[row].stress[position].components);
        }
    }
    EXPECT_DOUBLE_EQ(state.physicalStrainEnergy(), physicalStrainEnergy);
    EXPECT_EQ(state.equilibrium(), equilibrium);
    EXPECT_EQ(state.verificationMetrics(), verificationMetrics);
}

void expectShellCandidateRejectedWithoutMutation(
    fesa::AnalysisState& state,
    const std::vector<fesa::EntityIndex>& expectedElements,
    const fesa::ShellStateCandidate& candidate,
    const fesa::ShellStateCandidate& committed) {
    const auto status = state.commitShellResults(expectedElements, candidate);
    EXPECT_FALSE(status.IsOk());
    EXPECT_EQ(status.Category(), fesa::FailureCategory::kModel);
    expectShellStateEquals(
        state,
        committed.rows,
        committed.physicalStrainEnergy,
        committed.equilibrium,
        committed.verificationMetrics);
    ASSERT_TRUE(state.commitShellResults(expectedElements, committed).IsOk());
}

} // namespace

TEST(AnalysisState, PreservesStepFrameAndStableRowOrder) {
    const auto dofs = makeEmptyDofs();
    auto state = fesa::AnalysisState::create(dofs, {"Step-1", 0U});

    const fesa::EndpointResultRow firstEndpoint{
        2U,
        -1,
        {"Beam-1", 10, "010"},
        {1.0, 2.0, 3.0, 4.0, 5.0, 6.0},
        {7.0, 8.0, 9.0, 10.0}};
    const fesa::EndpointResultRow secondEndpoint{
        2U,
        1,
        {"Beam-1", 20, "020"},
        {11.0, 12.0, 13.0, 14.0, 15.0, 16.0},
        {17.0, 18.0, 19.0, 20.0}};
    const fesa::GaussResultRow firstGauss{
        2U, 1, {0.1, 0.2, 0.3, 0.4}, {1.1, 1.2, 1.3, 1.4}};
    const fesa::GaussResultRow secondGauss{
        2U, 2, {0.5, 0.6, 0.7, 0.8}, {1.5, 1.6, 1.7, 1.8}};
    const fesa::StressS11Row firstStress{
        2U, 1, 1U, -0.25, 0.5, 12.5, "input"};
    const fesa::StressS11Row secondStress{
        2U, 1, 2U, 0.25, -0.5, -7.5, "input"};

    state.endpointResults().push_back(firstEndpoint);
    state.endpointResults().push_back(secondEndpoint);
    state.gaussResults().push_back(firstGauss);
    state.gaussResults().push_back(secondGauss);
    state.stressResults().push_back(firstStress);
    state.stressResults().push_back(secondStress);

    const auto* const endpointStorage = state.endpointResults().data();
    const auto* const gaussStorage = state.gaussResults().data();
    const auto* const stressStorage = state.stressResults().data();
    const fesa::AnalysisState& constState = state;

    EXPECT_EQ(constState.identity().stepName, "Step-1");
    EXPECT_EQ(constState.identity().frameIndex, 0U);
    ASSERT_EQ(constState.endpointResults().size(), 2U);
    EXPECT_EQ(constState.endpointResults().data(), endpointStorage);
    EXPECT_EQ(constState.endpointResults()[0].element, 2U);
    EXPECT_EQ(constState.endpointResults()[0].endpoint, -1);
    EXPECT_EQ(constState.endpointResults()[0].node.instance_name, "Beam-1");
    EXPECT_EQ(constState.endpointResults()[0].node.source_label, 10);
    EXPECT_EQ(constState.endpointResults()[0].node.source_label_text, "010");
    EXPECT_EQ(
        constState.endpointResults()[0].endAction,
        (std::array<double, 6>{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}));
    EXPECT_EQ(
        constState.endpointResults()[1].sectionResultant,
        (std::array<double, 4>{17.0, 18.0, 19.0, 20.0}));

    ASSERT_EQ(constState.gaussResults().size(), 2U);
    EXPECT_EQ(constState.gaussResults().data(), gaussStorage);
    EXPECT_EQ(constState.gaussResults()[0].gaussPoint, 1);
    EXPECT_EQ(constState.gaussResults()[1].gaussPoint, 2);
    EXPECT_EQ(
        constState.gaussResults()[0].generalizedStrain,
        (std::array<double, 4>{0.1, 0.2, 0.3, 0.4}));
    EXPECT_EQ(
        constState.gaussResults()[1].generalizedResultant,
        (std::array<double, 4>{1.5, 1.6, 1.7, 1.8}));

    ASSERT_EQ(constState.stressResults().size(), 2U);
    EXPECT_EQ(constState.stressResults().data(), stressStorage);
    EXPECT_EQ(constState.stressResults()[0].sectionPoint, 1U);
    EXPECT_DOUBLE_EQ(constState.stressResults()[0].x1, -0.25);
    EXPECT_DOUBLE_EQ(constState.stressResults()[0].x2, 0.5);
    EXPECT_DOUBLE_EQ(constState.stressResults()[0].s11, 12.5);
    EXPECT_EQ(constState.stressResults()[0].source, "input");
    EXPECT_EQ(constState.stressResults()[1].sectionPoint, 2U);
}

// MITC4-STATE-001
TEST(AnalysisState, OwnsExactShellRowsInStableElementAndLocationOrder) {
    const auto dofs = makeEmptyDofs();
    auto state = fesa::AnalysisState::create(dofs, {"Step-1", 0U});
    const std::vector<fesa::EntityIndex> expectedElements{3U, 7U};
    auto candidate = makeShellCandidate(expectedElements);

    const auto status = state.commitShellResults(expectedElements, candidate);

    ASSERT_TRUE(status.IsOk());
    const fesa::AnalysisState& constState = state;
    ASSERT_EQ(constState.shellResults().size(), 8U);
    EXPECT_EQ(constState.shellResults()[0].element, 3U);
    EXPECT_EQ(
        constState.shellResults()[0].location,
        fesa::ShellMidsurfaceLocation::gp1);
    EXPECT_EQ(
        constState.shellResults()[3].location,
        fesa::ShellMidsurfaceLocation::gp4);
    EXPECT_EQ(constState.shellResults()[4].element, 7U);
    EXPECT_EQ(
        constState.shellResults()[4].location,
        fesa::ShellMidsurfaceLocation::gp1);
    EXPECT_EQ(
        constState.shellResults()[0].naturalCoordinates,
        (std::array<double, 2>{
            -1.0 / std::sqrt(3.0), -1.0 / std::sqrt(3.0)}));
    EXPECT_EQ(
        constState.shellResults()[2].generalizedStrain,
        (std::array<double, 8>{
            321.0, 322.0, 323.0, 324.0,
            325.0, 326.0, 327.0, 328.0}));
    EXPECT_EQ(
        constState.shellResults()[7].sectionResultant,
        (std::array<double, 8>{
            741.0, 742.0, 743.0, 744.0,
            745.0, 746.0, 747.0, 748.0}));
    EXPECT_EQ(
        constState.shellResults()[7].stress[0].position,
        fesa::ShellSectionPosition::bottom);
    EXPECT_DOUBLE_EQ(constState.shellResults()[7].stress[0].zeta, -1.0);
    EXPECT_EQ(
        constState.shellResults()[7].stress[2].components,
        (std::array<double, 3>{757.0, 758.0, 759.0}));
}

// MITC4-STATE-002
TEST(AnalysisState, CommitsFiniteShellGlobalEvidence) {
    const auto dofs = makeEmptyDofs();
    auto state = fesa::AnalysisState::create(dofs, {"Step-1", 0U});
    const std::vector<fesa::EntityIndex> expectedElements{5U};
    const auto candidate = makeShellCandidate(expectedElements);

    const auto status = state.commitShellResults(expectedElements, candidate);

    ASSERT_TRUE(status.IsOk());
    EXPECT_DOUBLE_EQ(state.physicalStrainEnergy(), 35.5);
    EXPECT_EQ(
        state.equilibrium(),
        (std::array<double, 6>{1.0, -2.0, 3.0, -4.0, 5.0, -6.0}));
    EXPECT_EQ(
        state.verificationMetrics(),
        (std::array<double, 3>{1.0e-11, 2.0e-11, 3.0e-11}));
}

// MITC4-STATE-003
TEST(AnalysisState, InvalidShellCandidatesLeavePriorStateUnchanged) {
    const auto dofs = makeEmptyDofs();
    auto state = fesa::AnalysisState::create(dofs, {"Step-1", 0U});
    const std::vector<fesa::EntityIndex> expectedElements{5U};
    const auto committed = makeShellCandidate(expectedElements);
    ASSERT_TRUE(state.commitShellResults(expectedElements, committed).IsOk());

    auto invalidLocation = makeShellCandidate(expectedElements);
    invalidLocation.rows[0].location = fesa::ShellMidsurfaceLocation::gp2;
    expectShellCandidateRejectedWithoutMutation(
        state, expectedElements, invalidLocation, committed);

    auto nonfinite = makeShellCandidate(expectedElements);
    nonfinite.rows[2].generalizedStrain[6] =
        (std::numeric_limits<double>::quiet_NaN)();
    expectShellCandidateRejectedWithoutMutation(
        state, expectedElements, nonfinite, committed);

    auto nonfiniteFrame = makeShellCandidate(expectedElements);
    nonfiniteFrame.rows[1].localFrame[2][0] =
        (std::numeric_limits<double>::infinity)();
    expectShellCandidateRejectedWithoutMutation(
        state, expectedElements, nonfiniteFrame, committed);

    auto invalidSectionPosition = makeShellCandidate(expectedElements);
    invalidSectionPosition.rows[3].stress[0].position =
        fesa::ShellSectionPosition::top;
    expectShellCandidateRejectedWithoutMutation(
        state, expectedElements, invalidSectionPosition, committed);

    auto nonfiniteGlobalEvidence = makeShellCandidate(expectedElements);
    nonfiniteGlobalEvidence.verificationMetrics[1] =
        (std::numeric_limits<double>::infinity)();
    expectShellCandidateRejectedWithoutMutation(
        state, expectedElements, nonfiniteGlobalEvidence, committed);

    auto incompleteInventory = makeShellCandidate(expectedElements);
    incompleteInventory.rows.pop_back();
    expectShellCandidateRejectedWithoutMutation(
        state, expectedElements, incompleteInventory, committed);
}