#include "fesa/results/result_recovery.hpp"

#include "fesa/analysis/analysis_model.hpp"
#include "fesa/analysis/analysis_state.hpp"
#include "fesa/assembly/parallel_for.hpp"
#include "fesa/assembly/sparse_assembler.hpp"
#include "fesa/fem/dof_manager.hpp"
#include "fesa/model/domain.hpp"

#include <gtest/gtest.h>

#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

constexpr double kYoungsModulus = 100.0;
constexpr double kPoissonRatio = 0.25;
constexpr double kLength = 2.0;

struct RecoveryFixture {
    std::unique_ptr<fesa::Domain> domain;
    std::unique_ptr<fesa::AnalysisModel> model;
    std::unique_ptr<fesa::DofManager> dofs;
    std::unique_ptr<fesa::SparseMatrix> stiffness;
};

fesa::ModelDefinition makeDefinition(
    const bool twoElements = false,
    std::vector<std::array<double, 2>> sectionPoints = {},
    std::vector<fesa::NodalLoad> loads = {},
    const bool reverseSecond = false,
    const bool sectionJump = false,
    const bool nonzeroPrescription = true) {
    const std::filesystem::path source{"models/result-recovery.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:0123456789abcdef";
    definition.nodes = {
        {{"Beam-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 10U}},
        {{"Beam-1", 2, "2"}, {kLength, 0.0, 0.0}, {source, 11U}}};
    if (twoElements) {
        definition.nodes.push_back(
            {{"Beam-1", 3, "3"}, {2.0 * kLength, 0.0, 0.0}, {source, 12U}});
    }
    definition.materials = {
        {"Material", kYoungsModulus, kPoissonRatio, {source, 20U}}};
    definition.sections = {{
        "Section",
        2.0,
        3.0,
        0.0,
        4.0,
        5.0,
        {0.0, 1.0, 0.0},
        std::move(sectionPoints),
        {source, 30U}}};
    if (sectionJump) {
        auto secondSection = definition.sections.front();
        secondSection.name = "Section-2";
        secondSection.area = 2.5;
        secondSection.location.line = 31U;
        definition.sections.push_back(std::move(secondSection));
    }
    definition.elements = {
        {{"Beam-1", 10, "10"}, {0U, 1U}, 0U, 0U, {source, 40U}}};
    if (twoElements) {
        definition.elements.push_back({
            {"Beam-1", 20, "20"},
            reverseSecond ? std::array<fesa::EntityIndex, 2>{2U, 1U}
                          : std::array<fesa::EntityIndex, 2>{1U, 2U},
            0U,
            sectionJump ? 1U : 0U,
            {source, 41U}});
    }
    definition.steps = {{
        "Step-1",
        {{"1", 1, 1, nonzeroPrescription ? 0.1 : 0.0, {source, 50U}},
         {"1", 2, 6, 0.0, {source, 51U}}},
        std::move(loads),
        0.1,
        1.0,
        0.01,
        1.0,
        {source, 49U}}};
    return definition;
}

RecoveryFixture makeFixture(
    const bool twoElements = false,
    std::vector<std::array<double, 2>> sectionPoints = {},
    std::vector<fesa::NodalLoad> loads = {},
    const bool reverseSecond = false,
    const bool sectionJump = false,
    const bool nonzeroPrescription = true) {
    auto domainResult = fesa::Domain::create(makeDefinition(
        twoElements,
        std::move(sectionPoints),
        std::move(loads),
        reverseSecond,
        sectionJump,
        nonzeroPrescription));
    if (!domainResult.hasValue()) {
        throw std::runtime_error{"Recovery fixture Domain construction failed."};
    }
    auto domain = std::make_unique<fesa::Domain>(
        std::move(domainResult.value()));

    auto modelResult = fesa::AnalysisModel::create(*domain);
    if (!modelResult.hasValue()) {
        throw std::runtime_error{"Recovery fixture AnalysisModel construction failed."};
    }
    auto model = std::make_unique<fesa::AnalysisModel>(
        std::move(modelResult.value()));

    auto dofsResult = fesa::DofManager::create(*model);
    if (!dofsResult.hasValue()) {
        throw std::runtime_error{"Recovery fixture DofManager construction failed."};
    }
    auto dofs = std::make_unique<fesa::DofManager>(
        std::move(dofsResult.value()));

    fesa::SerialParallelFor serial;
    auto stiffnessResult = fesa::SparseAssembler::assembleStiffness(
        *model, *dofs, serial);
    if (!stiffnessResult.hasValue()) {
        throw std::runtime_error{"Recovery fixture stiffness assembly failed."};
    }
    auto stiffness = std::make_unique<fesa::SparseMatrix>(
        std::move(stiffnessResult.value()));
    return {
        std::move(domain),
        std::move(model),
        std::move(dofs),
        std::move(stiffness)};
}

fesa::AnalysisState makeAxialEquilibriumState(const RecoveryFixture& fixture) {
    auto state = fesa::AnalysisState::create(
        *fixture.dofs, {"Step-1", 0U});
    state.displacement()[0U] = 0.1;
    state.displacement()[6U] = 0.3;
    const fesa::Vector internal = fixture.stiffness->multiply(state.displacement());
    for (const std::size_t fullDof : fixture.dofs->freeDofs()) {
        state.externalForce()[fullDof] = internal[fullDof];
    }
    return state;
}

fesa::AnalysisState makePatchState(
    const RecoveryFixture& fixture,
    const double epsilon,
    const double twist,
    const double kappaY,
    const double kappaZ) {
    auto state = fesa::AnalysisState::create(
        *fixture.dofs, {"Step-1", 0U});
    state.displacement()[0U] = 0.1;
    state.displacement()[6U] = 0.1 + epsilon * kLength;
    state.displacement()[7U] = 0.5 * kappaZ * kLength * kLength;
    state.displacement()[8U] = -0.5 * kappaY * kLength * kLength;
    state.displacement()[9U] = twist * kLength;
    state.displacement()[10U] = kappaY * kLength;
    state.displacement()[11U] = kappaZ * kLength;
    state.externalForce() = fixture.stiffness->multiply(state.displacement());
    return state;
}

void expectStatusCode(const fesa::Status& status, const std::string& code) {
    ASSERT_FALSE(status.isOk());
    EXPECT_EQ(status.failureCategory(), fesa::FailureCategory::model);
    ASSERT_EQ(status.diagnostics().size(), 1U);
    EXPECT_EQ(status.diagnostics()[0U].code, code);
}

void expectScaledNear(
    const double actual,
    const double expected,
    const double relativeTolerance = 1.0e-12) {
    ASSERT_TRUE(std::isfinite(actual));
    ASSERT_TRUE(std::isfinite(expected));
    EXPECT_LE(
        std::abs(actual - expected),
        relativeTolerance * (std::max)(std::abs(expected), 1.0));
}

std::vector<fesa::EndpointResultRow> makeStationRows(
    const RecoveryFixture& fixture) {
    const auto& nodes = fixture.domain->nodes();
    const auto& elements = fixture.domain->elements();
    return {
        {0U, 0, nodes[elements[0U].nodeIndices[0U]].sourceId,
         {0.0, 0.0, 0.0, 0.0, 0.0, 0.0},
         {1.0, 2.0, 3.0, 4.0}},
        {0U, 1, nodes[elements[0U].nodeIndices[1U]].sourceId,
         {0.0, 0.0, 0.0, 0.0, 0.0, 0.0},
         {5.0, 6.0, 7.0, 8.0}},
        {1U, 0, nodes[elements[1U].nodeIndices[0U]].sourceId,
         {0.0, 0.0, 0.0, 0.0, 0.0, 0.0},
         {5.0, 6.0, 7.0, 8.0}},
        {1U, 1, nodes[elements[1U].nodeIndices[1U]].sourceId,
         {0.0, 0.0, 0.0, 0.0, 0.0, 0.0},
         {9.0, 10.0, 11.0, 12.0}}};
}

} // namespace

TEST(ResultRecovery, ComputesResidualReactionForNonzeroPrescription) {
    const auto fixture = makeFixture();
    auto state = makeAxialEquilibriumState(fixture);

    const fesa::Status status = fesa::ResultRecovery::recover(
        *fixture.model, *fixture.dofs, *fixture.stiffness, state);

    ASSERT_TRUE(status.isOk());
    EXPECT_DOUBLE_EQ(state.internalForce()[0U], -20.0);
    EXPECT_DOUBLE_EQ(state.internalForce()[6U], 20.0);
    EXPECT_DOUBLE_EQ(state.residual()[0U], -20.0);
    EXPECT_DOUBLE_EQ(state.residual()[6U], 0.0);
    EXPECT_DOUBLE_EQ(state.reaction()[0U], -20.0);
    for (const std::size_t fullDof : fixture.dofs->freeDofs()) {
        EXPECT_DOUBLE_EQ(state.reaction()[fullDof], 0.0);
    }
}

TEST(ResultRecovery, EnforcesNormalizedFreeResidual) {
    const auto fixture = makeFixture();

    auto failed = makeAxialEquilibriumState(fixture);
    failed.internalForce()[0U] = 91.0;
    failed.residual()[0U] = 92.0;
    failed.reaction()[0U] = 93.0;
    failed.reaction()[6U] = 94.0;
    failed.endpointResults().push_back({});
    failed.externalForce()[6U] += 1.0e-7;
    expectStatusCode(
        fesa::ResultRecovery::recover(
            *fixture.model, *fixture.dofs, *fixture.stiffness, failed),
        "free-residual-tolerance-failure");
    EXPECT_DOUBLE_EQ(failed.internalForce()[0U], 91.0);
    EXPECT_DOUBLE_EQ(failed.residual()[0U], 92.0);
    EXPECT_DOUBLE_EQ(failed.reaction()[0U], 93.0);
    EXPECT_DOUBLE_EQ(failed.reaction()[6U], 94.0);
    EXPECT_EQ(failed.endpointResults().size(), 1U);

    auto thresholdPass = makeAxialEquilibriumState(fixture);
    thresholdPass.externalForce()[6U] += 1.0e-10 * 20.0 * 0.5;
    const fesa::Status thresholdStatus = fesa::ResultRecovery::recover(
        *fixture.model,
        *fixture.dofs,
        *fixture.stiffness,
        thresholdPass);
    ASSERT_TRUE(thresholdStatus.isOk());
    EXPECT_NE(thresholdPass.residual()[6U], 0.0);
    EXPECT_DOUBLE_EQ(
        thresholdPass.reaction()[6U], thresholdPass.residual()[6U]);

    const auto zeroFixture = makeFixture(false, {}, {}, false, false, false);
    auto zeroEquilibrium = fesa::AnalysisState::create(
        *zeroFixture.dofs, {"Step-1", 0U});
    EXPECT_TRUE(fesa::ResultRecovery::recover(
                    *zeroFixture.model,
                    *zeroFixture.dofs,
                    *zeroFixture.stiffness,
                    zeroEquilibrium)
                    .isOk());

    auto wrongPrescription = makeAxialEquilibriumState(fixture);
    wrongPrescription.displacement()[0U] = 0.0;
    expectStatusCode(
        fesa::ResultRecovery::recover(
            *fixture.model,
            *fixture.dofs,
            *fixture.stiffness,
            wrongPrescription),
        "invalid-recovery-state");

    auto nonfinite = makeAxialEquilibriumState(fixture);
    nonfinite.displacement()[6U] = std::numeric_limits<double>::quiet_NaN();
    expectStatusCode(
        fesa::ResultRecovery::recover(
            *fixture.model, *fixture.dofs, *fixture.stiffness, nonfinite),
        "nonfinite-recovery-value");

    const auto wrongFixture = makeFixture(true);
    auto wrongState = fesa::AnalysisState::create(
        *wrongFixture.dofs, {"Step-1", 0U});
    expectStatusCode(
        fesa::ResultRecovery::recover(
            *fixture.model, *fixture.dofs, *fixture.stiffness, wrongState),
        "invalid-recovery-dimensions");
}

TEST(ResultRecovery, KeepsEndActionSectionAndGaussResultsDistinct) {
    const auto fixture = makeFixture();
    auto state = makePatchState(fixture, 0.02, 0.03, -0.04, 0.05);

    ASSERT_TRUE(fesa::ResultRecovery::recover(
                    *fixture.model, *fixture.dofs, *fixture.stiffness, state)
                    .isOk());
    ASSERT_EQ(state.endpointResults().size(), 2U);
    ASSERT_EQ(state.gaussResults().size(), 2U);
    EXPECT_EQ(state.endpointResults()[0U].endpoint, 0);
    EXPECT_EQ(state.endpointResults()[1U].endpoint, 1);
    EXPECT_EQ(state.gaussResults()[0U].gaussPoint, 1);
    EXPECT_EQ(state.gaussResults()[1U].gaussPoint, 2);
    EXPECT_DOUBLE_EQ(
        state.endpointResults()[0U].endAction[0U],
        -state.endpointResults()[0U].sectionResultant[0U]);
    EXPECT_DOUBLE_EQ(
        state.endpointResults()[1U].endAction[0U],
        state.endpointResults()[1U].sectionResultant[0U]);
    EXPECT_DOUBLE_EQ(
        state.gaussResults()[0U].generalizedResultant[0U],
        state.endpointResults()[0U].sectionResultant[0U]);
}

TEST(ResultRecovery, MatchesAxialTorsionAndTwoPlaneEndSigns) {
    const auto fixture = makeFixture();
    const double epsilon = 0.02;
    const double twist = -0.03;
    const double kappaY = 0.04;
    const double kappaZ = -0.05;
    auto state = makePatchState(fixture, epsilon, twist, kappaY, kappaZ);

    ASSERT_TRUE(fesa::ResultRecovery::recover(
                    *fixture.model, *fixture.dofs, *fixture.stiffness, state)
                    .isOk());
    const double shearModulus =
        kYoungsModulus / (2.0 * (1.0 + kPoissonRatio));
    const std::array<double, 4> expected = {
        kYoungsModulus * 2.0 * epsilon,
        shearModulus * 5.0 * twist,
        kYoungsModulus * 3.0 * kappaY,
        kYoungsModulus * 4.0 * kappaZ};
    const std::array<std::size_t, 4> endComponents = {0U, 3U, 4U, 5U};
    for (std::size_t component = 0U; component < expected.size(); ++component) {
        expectScaledNear(
            state.endpointResults()[0U].sectionResultant[component],
            expected[component]);
        expectScaledNear(
            state.endpointResults()[1U].sectionResultant[component],
            expected[component]);
        expectScaledNear(
            state.endpointResults()[0U].endAction[endComponents[component]],
            -expected[component]);
        expectScaledNear(
            state.endpointResults()[1U].endAction[endComponents[component]],
            expected[component]);
    }
}

TEST(ResultRecovery, OrdersStressPointsAndDefaultCentroid) {
    const std::vector<std::array<double, 2>> sectionPoints = {
        {0.25, -0.5}, {-0.4, 0.3}};
    const auto fixture = makeFixture(false, sectionPoints);
    auto state = makePatchState(fixture, 0.01, 0.0, 0.02, -0.03);
    ASSERT_TRUE(fesa::ResultRecovery::recover(
                    *fixture.model, *fixture.dofs, *fixture.stiffness, state)
                    .isOk());
    ASSERT_EQ(state.stressResults().size(), 4U);
    for (std::size_t gauss = 0U; gauss < 2U; ++gauss) {
        for (std::size_t point = 0U; point < sectionPoints.size(); ++point) {
            const auto& row = state.stressResults()[gauss * 2U + point];
            EXPECT_EQ(row.element, 0U);
            EXPECT_EQ(row.gaussPoint, static_cast<int>(gauss + 1U));
            EXPECT_EQ(row.sectionPoint, point + 1U);
            EXPECT_DOUBLE_EQ(row.x1, sectionPoints[point][0U]);
            EXPECT_DOUBLE_EQ(row.x2, sectionPoints[point][1U]);
            EXPECT_EQ(row.source, "input");
            expectScaledNear(
                row.s11,
                kYoungsModulus *
                    (0.01 + row.x2 * 0.02 - row.x1 * -0.03));
        }
    }

    const auto defaultFixture = makeFixture();
    auto defaultState = makePatchState(defaultFixture, 0.01, 0.0, 0.0, 0.0);
    ASSERT_TRUE(fesa::ResultRecovery::recover(
                    *defaultFixture.model,
                    *defaultFixture.dofs,
                    *defaultFixture.stiffness,
                    defaultState)
                    .isOk());
    ASSERT_EQ(defaultState.stressResults().size(), 2U);
    for (const auto& row : defaultState.stressResults()) {
        EXPECT_EQ(row.sectionPoint, 0U);
        EXPECT_DOUBLE_EQ(row.x1, 0.0);
        EXPECT_DOUBLE_EQ(row.x2, 0.0);
        EXPECT_EQ(row.source, "fesa-default");
    }
}

TEST(ResultRecovery, RequiresInteriorEndpointConsistencyWithoutAveraging) {
    const auto fixture = makeFixture(true);
    const std::array<double, 4> tolerances = {1.0e-6, 1.0e-6, 1.0e-6, 1.0e-6};
    auto rows = makeStationRows(fixture);
    rows[2U].sectionResultant[0U] += 0.5e-6;

    auto normalized =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *fixture.model, rows, tolerances);
    ASSERT_TRUE(normalized.hasValue());
    ASSERT_EQ(normalized.value().size(), 3U);
    EXPECT_EQ(normalized.value()[1U].representativeElement, 0U);
    EXPECT_DOUBLE_EQ(normalized.value()[1U].sectionResultant[0U], 5.0);

    rows[2U].sectionResultant[0U] = 5.0 + 2.0e-6;
    auto mismatch =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *fixture.model, rows, tolerances);
    ASSERT_FALSE(mismatch.hasValue());
    expectStatusCode(mismatch.status(), "node-station-tolerance-failure");

    rows = makeStationRows(fixture);
    rows[2U].sectionResultant[1U] =
        std::numeric_limits<double>::infinity();
    auto nonfinite =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *fixture.model, rows, tolerances);
    ASSERT_FALSE(nonfinite.hasValue());
    expectStatusCode(nonfinite.status(), "nonfinite-node-station-value");

    auto invalidTolerance =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *fixture.model,
            makeStationRows(fixture),
            {1.0e-6, -1.0, 1.0e-6, 1.0e-6});
    ASSERT_FALSE(invalidTolerance.hasValue());
    expectStatusCode(
        invalidTolerance.status(), "invalid-node-station-tolerance");

    const std::filesystem::path source{"models/result-recovery.inp"};
    const auto loadedFixture = makeFixture(
        true, {}, {{"2", 2, 1.0, {source, 60U}}});
    auto loaded =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *loadedFixture.model,
            makeStationRows(loadedFixture),
            tolerances);
    ASSERT_FALSE(loaded.hasValue());
    expectStatusCode(loaded.status(), "ineligible-node-station");

    const auto reversedFixture = makeFixture(true, {}, {}, true);
    auto reversed =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *reversedFixture.model,
            makeStationRows(reversedFixture),
            tolerances);
    ASSERT_FALSE(reversed.hasValue());
    expectStatusCode(reversed.status(), "ineligible-node-station");

    const auto jumpFixture = makeFixture(true, {}, {}, false, true);
    auto jumped =
        fesa::ResultRecovery::normalizeSectionResultantsToNodeStations(
            *jumpFixture.model,
            makeStationRows(jumpFixture),
            tolerances);
    ASSERT_FALSE(jumped.hasValue());
    expectStatusCode(jumped.status(), "ineligible-node-station");
}