#include "fesa/assembly/load_assembler.hpp"

#include "fesa/analysis/analysis_model.hpp"
#include "fesa/fem/dof_manager.hpp"
#include "fesa/model/domain.hpp"

#include <gtest/gtest.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

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

LoadFixture makeFixture(
    const std::size_t nodeCount,
    std::vector<fesa::NodeSet> nodeSets,
    std::vector<fesa::BoundaryCondition> boundaries,
    std::vector<fesa::NodalLoad> loads) {
    const std::filesystem::path source{"models/load-assembly.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:abcdef0123456789";
    for (std::size_t index = 0U; index < nodeCount; ++index) {
        const auto label = static_cast<std::int64_t>((index + 1U) * 10U);
        definition.nodes.push_back({
            {"Beam-1", label, std::to_string(label)},
            {static_cast<double>(index), 0.0, 0.0},
            {source, index + 2U}});
    }
    definition.nodeSets = std::move(nodeSets);
    definition.steps = {{
        "Step-1",
        std::move(boundaries),
        std::move(loads),
        0.1,
        1.0,
        0.01,
        1.0,
        {source, 20U}}};

    auto domainResult = fesa::Domain::create(std::move(definition));
    if (!domainResult.HasValue()) {
        throw std::runtime_error{"Load 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{"Load fixture AnalysisModel construction failed."};
    }
    auto model = std::make_unique<fesa::AnalysisModel>(
        std::move(modelResult.Value()));

    auto dofResult = fesa::DofManager::create(*model);
    if (!dofResult.HasValue()) {
        throw std::runtime_error{"Load fixture DofManager construction failed."};
    }
    auto dofs = std::make_unique<fesa::DofManager>(
        std::move(dofResult.Value()));
    return {std::move(domain), std::move(model), std::move(dofs)};
}

LoadFixture makeShellFixture(
    std::vector<fesa::BoundaryCondition> boundaries,
    std::vector<fesa::NodalLoad> loads) {
    const std::filesystem::path source{"models/shell-load-assembly.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:1122334455667788";
    definition.nodes = {
        {{"Shell-1", 10, "10"}, {0.0, 0.0, 0.0}, {source, 2U}},
        {{"Shell-1", 20, "20"}, {1.0, 0.0, 0.0}, {source, 3U}},
        {{"Shell-1", 30, "30"}, {1.0, 1.0, 0.0}, {source, 4U}},
        {{"Shell-1", 40, "40"}, {0.0, 1.0, 0.0}, {source, 5U}}};
    definition.materials = {
        {"Material", 1000.0, 0.25, {source, 6U}}};
    definition.shellSections = {
        {"ShellSection", 0.1, 0U, {source, 7U}}};
    definition.shellElements = {{
        {"Shell-1", 1, "1"},
        fesa::ShellSourceElementType::s4,
        {0U, 1U, 2U, 3U},
        0U,
        0U,
        {source, 8U}}};
    for (std::size_t node = 0U; node < definition.nodes.size(); ++node) {
        definition.shellNodeInitialFrames.push_back({
            static_cast<fesa::EntityIndex>(node),
            {0.0, 0.0, 1.0},
            {1.0, 0.0, 0.0},
            {0.0, 1.0, 0.0}});
    }
    definition.steps = {{
        "Step-1",
        std::move(boundaries),
        std::move(loads),
        0.1,
        1.0,
        0.01,
        1.0,
        {source, 20U}}};

    auto domainResult = fesa::Domain::create(std::move(definition));
    if (!domainResult.HasValue()) {
        throw std::runtime_error{"Shell load 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{"Shell load fixture AnalysisModel construction failed."};
    }
    auto model = std::make_unique<fesa::AnalysisModel>(
        std::move(modelResult.Value()));

    auto dofResult = fesa::DofManager::create(*model);
    if (!dofResult.HasValue()) {
        throw std::runtime_error{"Shell load fixture DofManager construction failed."};
    }
    auto dofs = std::make_unique<fesa::DofManager>(
        std::move(dofResult.Value()));
    return {std::move(domain), std::move(model), std::move(dofs)};
}

fesa::SparseMatrix makeDenseSparse(
    const std::size_t rows,
    const std::size_t columns,
    const std::vector<double>& values) {
    if (values.size() != rows * columns) {
        throw std::invalid_argument{"Dense sparse fixture has the wrong value count."};
    }

    fesa::SparsePattern pattern;
    std::vector<fesa::CooContribution> contributions;
    pattern.rowOffsets.reserve(rows + 1U);
    pattern.rowOffsets.push_back(0U);
    for (std::size_t row = 0U; row < rows; ++row) {
        for (std::size_t column = 0U; column < columns; ++column) {
            pattern.columnIndices.push_back(column);
            contributions.push_back({
                row,
                column,
                values[row * columns + column],
                row,
                column});
        }
        pattern.rowOffsets.push_back(pattern.columnIndices.size());
    }

    auto result = fesa::SparseMatrix::FromCoo(
        rows, columns, std::move(contributions), pattern);
    if (!result.HasValue()) {
        throw std::runtime_error{"Sparse fixture construction failed."};
    }
    return std::move(result.Value());
}

void expectFailureCode(
    const fesa::Result<fesa::Vector>& result,
    const std::string& code) {
    ASSERT_FALSE(result.HasValue());
    EXPECT_EQ(result.GetStatus().Category(), fesa::FailureCategory::kModel);
    ASSERT_EQ(result.GetStatus().Diagnostics().size(), 1U);
    EXPECT_EQ(result.GetStatus().Diagnostics()[0U].code, code);
}

} // namespace

TEST(LoadAssembly, AssemblesNodeSetAndSixComponentLoads) {
    const std::filesystem::path source{"models/load-assembly.inp"};
    auto fixture = makeFixture(
        2U,
        {{"Pair", std::nullopt, {0U, 1U}, {source, 10U}}},
        {{"10", 1, 1, 0.0, {source, 21U}}},
        {{"pair", 1, 1.0, {source, 30U}},
         {"10", 2, 2.0, {source, 31U}},
         {"20", 3, 3.0, {source, 32U}},
         {"10", 4, -4.0, {source, 33U}},
         {"PAIR", 5, 5.0, {source, 34U}},
         {"20", 6, 6.0, {source, 35U}}});

    auto result = fesa::LoadAssembler::assembleFullNodalLoad(
        *fixture.model, *fixture.dofs);
    ASSERT_TRUE(result.HasValue());
    ASSERT_EQ(result.Value().Size(), 12U);
    EXPECT_EQ(
        std::vector<double>(result.Value().Data(), result.Value().Data() + 12U),
        (std::vector<double>{
            1.0, 2.0, 0.0, -4.0, 5.0, 0.0,
            1.0, 0.0, 3.0, 0.0, 5.0, 6.0}));
    EXPECT_EQ(fixture.dofs->constrainedDofs(),
              (std::vector<std::size_t>{0U}));
    EXPECT_DOUBLE_EQ(result.Value()[0U], 1.0);
}

TEST(LoadAssembly, AccumulatesSignedLoadsInSourceOrder) {
    const std::filesystem::path source{"models/load-assembly.inp"};
    auto firstOrder = makeFixture(
        1U,
        {},
        {},
        {{"10", 1, 1.0e16, {source, 30U}},
         {"10", 1, -1.0e16, {source, 31U}},
         {"10", 1, 1.0, {source, 32U}}});
    auto secondOrder = makeFixture(
        1U,
        {},
        {},
        {{"10", 1, 1.0e16, {source, 30U}},
         {"10", 1, 1.0, {source, 31U}},
         {"10", 1, -1.0e16, {source, 32U}}});

    auto first = fesa::LoadAssembler::assembleFullNodalLoad(
        *firstOrder.model, *firstOrder.dofs);
    auto second = fesa::LoadAssembler::assembleFullNodalLoad(
        *secondOrder.model, *secondOrder.dofs);
    ASSERT_TRUE(first.HasValue());
    ASSERT_TRUE(second.HasValue());
    EXPECT_DOUBLE_EQ(first.Value()[0U], 1.0);
    EXPECT_DOUBLE_EQ(second.Value()[0U], 0.0);
}

// MITC4-LOAD-001
TEST(LoadAssembly, AggregatesAllSixGlobalShellLoadComponentsInSourceOrder) {
    const std::filesystem::path source{"models/shell-load-assembly.inp"};
    auto fixture = makeShellFixture(
        {},
        {{"10", 1, 1.0e16, {source, 30U}},
         {"10", 1, -1.0e16, {source, 31U}},
         {"10", 1, 1.0, {source, 32U}},
         {"10", 2, 2.0, {source, 33U}},
         {"10", 3, 3.0, {source, 34U}},
         {"10", 4, 1.0e16, {source, 35U}},
         {"10", 4, -1.0e16, {source, 36U}},
         {"10", 4, 4.0, {source, 37U}},
         {"10", 5, 5.0, {source, 38U}},
         {"10", 6, 6.0, {source, 39U}},
         {"10", 6, -6.0, {source, 40U}}});

    const auto result = fesa::LoadAssembler::assembleFullNodalLoad(
        *fixture.model, *fixture.dofs);

    ASSERT_TRUE(result.HasValue());
    ASSERT_EQ(result.Value().Size(), 24U);
    EXPECT_EQ(
        std::vector<double>(result.Value().Data(), result.Value().Data() + 6U),
        (std::vector<double>{1.0, 2.0, 3.0, 4.0, 5.0, 0.0}));
}

// MITC4-LOAD-002
TEST(LoadAssembly, AcceptsExactlyZeroAggregateShellMoment) {
    const std::filesystem::path source{"models/shell-load-assembly.inp"};
    auto fixture = makeShellFixture(
        {},
        {{"10", 4, 3.0, {source, 30U}},
         {"10", 4, -3.0, {source, 31U}},
         {"10", 5, 4.0, {source, 32U}},
         {"10", 5, -4.0, {source, 33U}},
         {"10", 6, 5.0, {source, 34U}},
         {"10", 6, -5.0, {source, 35U}}});

    const auto result = fesa::LoadAssembler::assembleFullNodalLoad(
        *fixture.model, *fixture.dofs);

    ASSERT_TRUE(result.HasValue());
    EXPECT_DOUBLE_EQ(result.Value()[3U], 0.0);
    EXPECT_DOUBLE_EQ(result.Value()[4U], 0.0);
    EXPECT_DOUBLE_EQ(result.Value()[5U], 0.0);
}

// MITC4-LOAD-003
TEST(LoadAssembly, EnforcesAggregateShellMomentDirectorProjectionThreshold) {
    const std::filesystem::path source{"models/shell-load-assembly.inp"};
    auto acceptedFixture = makeShellFixture(
        {},
        {{"10", 4, 1.0, {source, 30U}},
         {"10", 6, 1.0e-12, {source, 31U}}});
    auto rejectedFixture = makeShellFixture(
        {},
        {{"10", 4, 1.0, {source, 30U}},
         {"10", 6, 2.0e-12, {source, 31U}}});

    const auto accepted = fesa::LoadAssembler::assembleFullNodalLoad(
        *acceptedFixture.model, *acceptedFixture.dofs);
    const auto rejected = fesa::LoadAssembler::assembleFullNodalLoad(
        *rejectedFixture.model, *rejectedFixture.dofs);

    ASSERT_TRUE(accepted.HasValue());
    ASSERT_FALSE(rejected.HasValue());
    EXPECT_EQ(rejected.GetStatus().Category(), fesa::FailureCategory::kModel);
    ASSERT_EQ(rejected.GetStatus().Diagnostics().size(), 1U);
    EXPECT_EQ(
        rejected.GetStatus().Diagnostics()[0U].code,
        "unsupported-drilling-load");
    EXPECT_EQ(rejected.GetStatus().Diagnostics()[0U].keyword, "CLOAD");
    EXPECT_EQ(rejected.GetStatus().Diagnostics()[0U].entity_identity, "10");
}

// MITC4-LOAD-004
TEST(LoadAssembly, RejectsDrillingMomentBeforeEffectiveRhsCanBeFormed) {
    const std::filesystem::path source{"models/shell-load-assembly.inp"};
    auto fixture = makeShellFixture(
        {{"10", 1, 1, 2.0, {source, 21U}}},
        {{"10", 6, 1.0, {source, 30U}}});

    const auto rejected = fesa::LoadAssembler::assembleFullNodalLoad(
        *fixture.model, *fixture.dofs);

    ASSERT_FALSE(rejected.HasValue());
    EXPECT_EQ(rejected.GetStatus().Category(), fesa::FailureCategory::kModel);
    ASSERT_EQ(rejected.GetStatus().Diagnostics().size(), 1U);
    EXPECT_EQ(
        rejected.GetStatus().Diagnostics()[0U].code,
        "unsupported-drilling-load");
}

TEST(LoadAssembly, FormsNonzeroPrescribedEffectiveRhs) {
    const std::filesystem::path source{"models/load-assembly.inp"};
    auto fixture = makeFixture(
        1U,
        {},
        {{"10", 2, 2, 2.0, {source, 21U}},
         {"10", 5, 5, -1.0, {source, 22U}}},
        {{"10", 1, 10.0, {source, 30U}},
         {"10", 2, 900.0, {source, 31U}},
         {"10", 3, 20.0, {source, 32U}},
         {"10", 4, 30.0, {source, 33U}},
         {"10", 5, 800.0, {source, 34U}},
         {"10", 6, 40.0, {source, 35U}}});
    auto full = fesa::LoadAssembler::assembleFullNodalLoad(
        *fixture.model, *fixture.dofs);
    ASSERT_TRUE(full.HasValue());
    const auto kfc = makeDenseSparse(
        4U,
        2U,
        {1.0, 2.0,
         3.0, 4.0,
         -2.0, 5.0,
         0.5, -1.0});

    auto rhs = fesa::LoadAssembler::effectiveFreeRhs(
        full.Value(), kfc, fixture.dofs->prescribedValues(), *fixture.dofs);
    ASSERT_TRUE(rhs.HasValue());
    ASSERT_EQ(rhs.Value().Size(), 4U);
    EXPECT_EQ(
        std::vector<double>(rhs.Value().Data(), rhs.Value().Data() + 4U),
        (std::vector<double>{10.0, 18.0, 39.0, 38.0}));
}

TEST(LoadAssembly, RejectsNonfiniteOrDimensionMismatch) {
    const std::filesystem::path source{"models/load-assembly.inp"};
    const double maximum = (std::numeric_limits<double>::max)();
    auto nonfinite = makeFixture(
        1U,
        {},
        {},
        {{"10", 1, std::numeric_limits<double>::quiet_NaN(), {source, 30U}}});
    expectFailureCode(
        fesa::LoadAssembler::assembleFullNodalLoad(
            *nonfinite.model, *nonfinite.dofs),
        "nonfinite-load-value");

    auto overflow = makeFixture(
        1U,
        {},
        {},
        {{"10", 1, maximum, {source, 30U}},
         {"10", 1, maximum, {source, 31U}}});
    expectFailureCode(
        fesa::LoadAssembler::assembleFullNodalLoad(
            *overflow.model, *overflow.dofs),
        "nonfinite-load-accumulation");

    auto oneNode = makeFixture(
        1U,
        {},
        {{"10", 2, 2, 2.0, {source, 21U}},
         {"10", 5, 5, -1.0, {source, 22U}}},
        {});
    auto twoNodes = makeFixture(2U, {}, {}, {});
    expectFailureCode(
        fesa::LoadAssembler::assembleFullNodalLoad(
            *twoNodes.model, *oneNode.dofs),
        "invalid-load-dimensions");

    const auto validKfc = makeDenseSparse(4U, 2U, std::vector<double>(8U, 0.0));
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            fesa::Vector{5U},
            validKfc,
            oneNode.dofs->prescribedValues(),
            *oneNode.dofs),
        "invalid-load-dimensions");
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            fesa::Vector{6U},
            makeDenseSparse(3U, 2U, std::vector<double>(6U, 0.0)),
            oneNode.dofs->prescribedValues(),
            *oneNode.dofs),
        "invalid-load-dimensions");
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            fesa::Vector{6U},
            makeDenseSparse(4U, 1U, std::vector<double>(4U, 0.0)),
            oneNode.dofs->prescribedValues(),
            *oneNode.dofs),
        "invalid-load-dimensions");
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            fesa::Vector{6U}, validKfc, fesa::Vector{1U}, *oneNode.dofs),
        "invalid-load-dimensions");

    fesa::Vector nonfiniteFull{6U};
    nonfiniteFull[0U] = std::numeric_limits<double>::infinity();
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            nonfiniteFull,
            validKfc,
            oneNode.dofs->prescribedValues(),
            *oneNode.dofs),
        "nonfinite-load-value");

    fesa::Vector nonfinitePrescribed{2U};
    nonfinitePrescribed[0U] = std::numeric_limits<double>::quiet_NaN();
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            fesa::Vector{6U}, validKfc, nonfinitePrescribed, *oneNode.dofs),
        "nonfinite-load-value");

    const auto overflowingKfc = makeDenseSparse(
        4U,
        2U,
        {maximum, 0.0,
         0.0, 0.0,
         0.0, 0.0,
         0.0, 0.0});
    expectFailureCode(
        fesa::LoadAssembler::effectiveFreeRhs(
            fesa::Vector{6U},
            overflowingKfc,
            oneNode.dofs->prescribedValues(),
            *oneNode.dofs),
        "nonfinite-load-accumulation");
}

TEST(LoadAssembly, ZeroLoadsRemainZero) {
    const std::filesystem::path source{"models/load-assembly.inp"};
    auto freeFixture = makeFixture(
        1U,
        {},
        {},
        {{"10", 3, 0.0, {source, 30U}}});
    auto full = fesa::LoadAssembler::assembleFullNodalLoad(
        *freeFixture.model, *freeFixture.dofs);
    ASSERT_TRUE(full.HasValue());
    EXPECT_TRUE(std::all_of(
        full.Value().Data(),
        full.Value().Data() + full.Value().Size(),
        [](const double value) { return value == 0.0; }));
    const auto noConstrainedColumns = makeDenseSparse(6U, 0U, {});
    auto freeRhs = fesa::LoadAssembler::effectiveFreeRhs(
        full.Value(),
        noConstrainedColumns,
        freeFixture.dofs->prescribedValues(),
        *freeFixture.dofs);
    ASSERT_TRUE(freeRhs.HasValue());
    EXPECT_EQ(freeRhs.Value().Size(), 6U);
    EXPECT_TRUE(std::all_of(
        freeRhs.Value().Data(),
        freeRhs.Value().Data() + freeRhs.Value().Size(),
        [](const double value) { return value == 0.0; }));

    auto constrainedFixture = makeFixture(
        1U,
        {},
        {{"10", 1, 6, 0.0, {source, 21U}}},
        {});
    auto constrainedFull = fesa::LoadAssembler::assembleFullNodalLoad(
        *constrainedFixture.model, *constrainedFixture.dofs);
    ASSERT_TRUE(constrainedFull.HasValue());
    const auto noFreeRows = makeDenseSparse(0U, 6U, {});
    auto constrainedRhs = fesa::LoadAssembler::effectiveFreeRhs(
        constrainedFull.Value(),
        noFreeRows,
        constrainedFixture.dofs->prescribedValues(),
        *constrainedFixture.dofs);
    ASSERT_TRUE(constrainedRhs.HasValue());
    EXPECT_EQ(constrainedRhs.Value().Size(), 0U);
}