#include "fesa/fem/dof_manager.hpp"

#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <filesystem>
#include <optional>
#include <utility>
#include <vector>

namespace {

fesa::ModelDefinition makeDefinition() {
    const std::filesystem::path source{"models/dof-manager.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:0123456789abcdef";
    definition.nodes = {
        {{"Beam-1", 10, "10"}, {0.0, 0.0, 0.0}, {source, 10U}},
        {{"Beam-1", 20, "20"}, {1.0, 0.0, 0.0}, {source, 11U}},
        {{"Beam-1", 30, "30"}, {2.0, 0.0, 0.0}, {source, 12U}}};
    definition.materials = {
        {"Material", 1000.0, 0.25, {source, 20U}}};
    definition.sections = {{
        "Section", 1.0, 1.0, 0.0, 1.0, 1.0,
        {0.0, 1.0, 0.0}, {}, {source, 30U}}};
    definition.elements = {
        {{"Beam-1", 100, "100"}, {0U, 1U}, 0U, 0U, {source, 40U}},
        {{"Beam-1", 200, "200"}, {1U, 2U}, 0U, 0U, {source, 41U}}};
    definition.nodeSets = {
        {"Root", std::optional<std::string>{"Beam-1"}, {0U}, {source, 50U}},
        {"Ends", std::optional<std::string>{"Beam-1"}, {0U, 2U}, {source, 51U}}};
    definition.steps = {{
        "Step-1",
        {{"Root", 1, 2, 0.0, {source, 60U}},
         {"ends", 3, 3, 0.25, {source, 61U}},
         {"20", 6, 6, -0.5, {source, 62U}},
         {"Root", 1, 1, 0.0, {source, 63U}}},
        {},
        0.1,
        1.0,
        0.01,
        1.0,
        {source, 59U}}};
    return definition;
}

fesa::ModelDefinition makeShellDefinition() {
    const std::filesystem::path source{"models/shell-dof-manager.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:1122334455667788";
    definition.nodes = {
        {{"Shell-1", 10, "10"}, {0.0, 0.0, 0.0}, {source, 10U}},
        {{"Shell-1", 20, "20"}, {1.0, 0.0, 0.0}, {source, 11U}},
        {{"Shell-1", 30, "30"}, {1.0, 1.0, 0.0}, {source, 12U}},
        {{"Shell-1", 40, "40"}, {0.0, 1.0, 0.0}, {source, 13U}}};
    definition.materials = {
        {"Material", 1000.0, 0.25, {source, 20U}}};
    definition.shellSections = {
        {"ShellSection", 0.1, 0U, {source, 30U}}};
    definition.shellElements = {{
        {"Shell-1", 100, "100"},
        fesa::ShellSourceElementType::s4,
        {2U, 0U, 3U, 1U},
        0U,
        0U,
        {source, 40U}}};
    definition.steps = {{
        "Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {source, 50U}}};
    return definition;
}

struct DofFixture {
    fesa::DofManager dofs;
};

DofFixture makeDofFixture(fesa::ModelDefinition definition = makeDefinition()) {
    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())};
}

std::vector<std::size_t> rowColumns(
    const fesa::SparsePattern& pattern, std::size_t row) {
    return {
        pattern.columnIndices.begin() + pattern.rowOffsets[row],
        pattern.columnIndices.begin() + pattern.rowOffsets[row + 1U]};
}

} // namespace

TEST(DofManager, NumbersSixDofsAndFreeEquationsStably) {
    const auto fixture = makeDofFixture();
    const auto& dofs = fixture.dofs;

    EXPECT_EQ(dofs.fullDofCount(), 18U);
    EXPECT_EQ(dofs.freeDofCount(), 13U);
    EXPECT_EQ(dofs.constrainedDofCount(), 5U);
    EXPECT_EQ(dofs.fullDof(0U, fesa::DofComponent::ux), 0U);
    EXPECT_EQ(dofs.fullDof(0U, fesa::DofComponent::urz), 5U);
    EXPECT_EQ(dofs.fullDof(1U, fesa::DofComponent::ux), 6U);
    EXPECT_EQ(dofs.fullDof(2U, fesa::DofComponent::urz), 17U);
    EXPECT_EQ(
        dofs.freeDofs(),
        (std::vector<std::size_t>{
            3U, 4U, 5U, 6U, 7U, 8U, 9U, 10U, 12U, 13U, 15U, 16U, 17U}));
    EXPECT_EQ(
        dofs.constrainedDofs(),
        (std::vector<std::size_t>{0U, 1U, 2U, 11U, 14U}));
    EXPECT_EQ(dofs.freeEquation(0U), std::nullopt);
    EXPECT_EQ(dofs.freeEquation(3U), std::optional<std::size_t>{0U});
    EXPECT_EQ(dofs.freeEquation(10U), std::optional<std::size_t>{7U});
    EXPECT_EQ(dofs.freeEquation(17U), std::optional<std::size_t>{12U});
}

TEST(DofManager, ExpandsAndValidatesPrescribedValues) {
    const auto fixture = makeDofFixture();
    const auto& values = fixture.dofs.prescribedValues();
    ASSERT_EQ(values.Size(), 5U);
    EXPECT_DOUBLE_EQ(values[0], 0.0);
    EXPECT_DOUBLE_EQ(values[1], 0.0);
    EXPECT_DOUBLE_EQ(values[2], 0.25);
    EXPECT_DOUBLE_EQ(values[3], -0.5);
    EXPECT_DOUBLE_EQ(values[4], 0.25);

    auto conflictingDefinition = makeDefinition();
    conflictingDefinition.steps[0].boundaries.push_back(
        {"root", 1, 1, 1.0, {conflictingDefinition.sourcePath, 77U}});
    auto domain = fesa::Domain::create(std::move(conflictingDefinition));
    ASSERT_TRUE(domain.HasValue());
    auto model = fesa::AnalysisModel::create(domain.Value());
    ASSERT_TRUE(model.HasValue());

    auto conflict = fesa::DofManager::create(model.Value());
    ASSERT_FALSE(conflict.HasValue());
    EXPECT_EQ(conflict.GetStatus().Category(), fesa::FailureCategory::kInput);
    ASSERT_EQ(conflict.GetStatus().Diagnostics().size(), 1U);
    const auto& diagnostic = conflict.GetStatus().Diagnostics()[0];
    EXPECT_EQ(diagnostic.code, "conflicting-boundary-condition");
    EXPECT_EQ(diagnostic.keyword, "BOUNDARY");
    EXPECT_EQ(diagnostic.entity_identity, "root");
    EXPECT_EQ(diagnostic.location.file, std::filesystem::path{"models/dof-manager.inp"});
    EXPECT_EQ(diagnostic.location.line, 77U);
}

TEST(DofManager, BuildsTwelveDofScatterAndSortedUniquePattern) {
    const auto fixture = makeDofFixture();
    const auto& dofs = fixture.dofs;

    EXPECT_EQ(
        dofs.elementScatter(0U),
        (std::array<std::size_t, 12>{
            0U, 1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U, 9U, 10U, 11U}));
    EXPECT_EQ(
        dofs.elementScatter(1U),
        (std::array<std::size_t, 12>{
            6U, 7U, 8U, 9U, 10U, 11U,
            12U, 13U, 14U, 15U, 16U, 17U}));

    const auto& pattern = dofs.sparsePattern();
    EXPECT_EQ(
        pattern.rowOffsets,
        (std::vector<std::size_t>{
            0U, 12U, 24U, 36U, 48U, 60U, 72U,
            90U, 108U, 126U, 144U, 162U, 180U,
            192U, 204U, 216U, 228U, 240U, 252U}));
    const std::vector<std::size_t> firstBlock{
        0U, 1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U, 9U, 10U, 11U};
    const std::vector<std::size_t> sharedBlock{
        0U, 1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U,
        9U, 10U, 11U, 12U, 13U, 14U, 15U, 16U, 17U};
    const std::vector<std::size_t> lastBlock{
        6U, 7U, 8U, 9U, 10U, 11U, 12U, 13U, 14U, 15U, 16U, 17U};
    EXPECT_EQ(rowColumns(pattern, 0U), firstBlock);
    EXPECT_EQ(rowColumns(pattern, 7U), sharedBlock);
    EXPECT_EQ(rowColumns(pattern, 17U), lastBlock);
    for (std::size_t row = 0U; row < dofs.fullDofCount(); ++row) {
        const auto columns = rowColumns(pattern, row);
        EXPECT_TRUE(std::is_sorted(columns.begin(), columns.end()));
        EXPECT_EQ(std::adjacent_find(columns.begin(), columns.end()), columns.end());
    }
}

// MITC4-DOF-001
TEST(DofManager, BuildsShellScatterInSourceNodeAndComponentOrder) {
    const auto fixture = makeDofFixture(makeShellDefinition());
    const auto& dofs = fixture.dofs;

    EXPECT_EQ(dofs.fullDofCount(), 24U);
    EXPECT_EQ(
        dofs.shellElementScatter(0U),
        (std::array<std::size_t, 24>{
            12U, 13U, 14U, 15U, 16U, 17U,
            0U, 1U, 2U, 3U, 4U, 5U,
            18U, 19U, 20U, 21U, 22U, 23U,
            6U, 7U, 8U, 9U, 10U, 11U}));
}

// MITC4-DOF-002
TEST(DofManager, IncludesShellConnectivityInSortedUniqueSparsePattern) {
    const auto fixture = makeDofFixture(makeShellDefinition());
    const auto& dofs = fixture.dofs;
    const auto& pattern = dofs.sparsePattern();

    ASSERT_EQ(pattern.rowOffsets.size(), 25U);
    ASSERT_EQ(pattern.columnIndices.size(), 24U * 24U);
    for (std::size_t row = 0U; row < dofs.fullDofCount(); ++row) {
        EXPECT_EQ(pattern.rowOffsets[row], row * 24U);
        EXPECT_EQ(pattern.rowOffsets[row + 1U], (row + 1U) * 24U);
        const auto columns = rowColumns(pattern, row);
        ASSERT_EQ(columns.size(), 24U);
        for (std::size_t column = 0U; column < columns.size(); ++column) {
            EXPECT_EQ(columns[column], column);
        }
        EXPECT_TRUE(std::binary_search(columns.begin(), columns.end(), row));
        EXPECT_EQ(std::adjacent_find(columns.begin(), columns.end()), columns.end());
    }
}

TEST(DofManager, ReconstructsFullReducedRoundTrip) {
    const auto fixture = makeDofFixture();
    const auto& dofs = fixture.dofs;
    fesa::Vector full{dofs.fullDofCount()};
    for (std::size_t index = 0U; index < full.Size(); ++index) {
        full[index] = static_cast<double>(index) + 0.5;
    }
    for (std::size_t index = 0U; index < dofs.constrainedDofCount(); ++index) {
        full[dofs.constrainedDofs()[index]] = dofs.prescribedValues()[index];
    }

    fesa::Vector reduced{dofs.freeDofCount()};
    for (std::size_t equation = 0U; equation < reduced.Size(); ++equation) {
        reduced[equation] = full[dofs.freeDofs()[equation]];
    }
    fesa::Vector reconstructed{dofs.fullDofCount()};
    for (std::size_t equation = 0U; equation < reduced.Size(); ++equation) {
        reconstructed[dofs.freeDofs()[equation]] = reduced[equation];
    }
    for (std::size_t index = 0U; index < dofs.constrainedDofCount(); ++index) {
        reconstructed[dofs.constrainedDofs()[index]] = dofs.prescribedValues()[index];
    }

    ASSERT_EQ(reconstructed.Size(), full.Size());
    for (std::size_t index = 0U; index < full.Size(); ++index) {
        EXPECT_DOUBLE_EQ(reconstructed[index], full[index]);
    }
}