#include "fesa/constraints/essential_constraints.hpp"

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

#include <gtest/gtest.h>

#include <cstddef>
#include <filesystem>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

fesa::DofManager makeDofs(
    std::vector<fesa::BoundaryCondition> boundaries) {
    const std::filesystem::path source{"models/essential-constraints.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:1234567890abcdef";
    definition.nodes = {
        {{"Beam-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 2U}}};
    definition.steps = {{
        "Step-1",
        std::move(boundaries),
        {},
        0.1,
        1.0,
        0.01,
        1.0,
        {source, 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::DofManager makeShellSizedDofs(
    std::vector<fesa::BoundaryCondition> boundaries) {
    const std::filesystem::path source{"models/shell-essential-constraints.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:8877665544332211";
    definition.nodes = {
        {{"Shell-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 2U}},
        {{"Shell-1", 2, "2"}, {1.0, 0.0, 0.0}, {source, 3U}},
        {{"Shell-1", 3, "3"}, {1.0, 1.0, 0.0}, {source, 4U}},
        {{"Shell-1", 4, "4"}, {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}}};
    definition.steps = {{
        "Step-1",
        std::move(boundaries),
        {},
        0.1,
        1.0,
        0.01,
        1.0,
        {source, 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::SparseMatrix makeMatrix(
    const std::size_t rows,
    const std::size_t columns,
    const std::vector<double>& denseValues) {
    EXPECT_EQ(denseValues.size(), rows * columns);
    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,
                denseValues[row * columns + column],
                row,
                column});
        }
        pattern.rowOffsets.push_back(pattern.columnIndices.size());
    }

    auto matrix = fesa::SparseMatrix::fromCoo(
        rows, columns, std::move(contributions), pattern);
    EXPECT_TRUE(matrix.hasValue());
    return std::move(matrix.value());
}

std::vector<double> sequentialDense(const std::size_t size) {
    std::vector<double> values(size * size);
    for (std::size_t row = 0U; row < size; ++row) {
        for (std::size_t column = 0U; column < size; ++column) {
            values[row * size + column] =
                static_cast<double>(row * 10U + column + 1U);
        }
    }
    return values;
}

void expectShape(
    const fesa::SparseMatrix& matrix,
    const std::size_t rows,
    const std::size_t columns) {
    EXPECT_EQ(matrix.rows(), rows);
    EXPECT_EQ(matrix.columns(), columns);
    EXPECT_TRUE(matrix.validate().isOk());
}

} // namespace

TEST(EssentialConstraints, ExtractsHandComputedBlocksInStableOrder) {
    const auto dofs = makeDofs({
        {"1", 2, 2, 2.5, {{}, 12U}},
        {"1", 5, 5, -3.25, {{}, 13U}}});
    auto fullValues = sequentialDense(6U);
    fullValues[2U * 6U + 4U] = 0.0;
    const auto full = makeMatrix(6U, 6U, fullValues);

    auto result = fesa::EssentialConstraints::partition(full, dofs);
    ASSERT_TRUE(result.hasValue());
    const auto& blocks = result.value();

    EXPECT_EQ(blocks.kff.rowOffsets(), (std::vector<std::size_t>{0U, 4U, 8U, 12U, 16U}));
    EXPECT_EQ(
        blocks.kff.columnIndices(),
        (std::vector<std::size_t>{
            0U, 1U, 2U, 3U, 0U, 1U, 2U, 3U,
            0U, 1U, 2U, 3U, 0U, 1U, 2U, 3U}));
    EXPECT_EQ(
        blocks.kff.values(),
        (std::vector<double>{
            1.0, 3.0, 4.0, 6.0,
            21.0, 23.0, 24.0, 26.0,
            31.0, 33.0, 34.0, 36.0,
            51.0, 53.0, 54.0, 56.0}));
    EXPECT_EQ(blocks.kfc.rowOffsets(), (std::vector<std::size_t>{0U, 2U, 4U, 6U, 8U}));
    EXPECT_EQ(
        blocks.kfc.columnIndices(),
        (std::vector<std::size_t>{0U, 1U, 0U, 1U, 0U, 1U, 0U, 1U}));
    EXPECT_EQ(
        blocks.kfc.values(),
        (std::vector<double>{2.0, 5.0, 22.0, 0.0, 32.0, 35.0, 52.0, 55.0}));
    EXPECT_EQ(blocks.kcf.rowOffsets(), (std::vector<std::size_t>{0U, 4U, 8U}));
    EXPECT_EQ(
        blocks.kcf.columnIndices(),
        (std::vector<std::size_t>{0U, 1U, 2U, 3U, 0U, 1U, 2U, 3U}));
    EXPECT_EQ(
        blocks.kcf.values(),
        (std::vector<double>{11.0, 13.0, 14.0, 16.0, 41.0, 43.0, 44.0, 46.0}));
    EXPECT_EQ(blocks.kcc.rowOffsets(), (std::vector<std::size_t>{0U, 2U, 4U}));
    EXPECT_EQ(blocks.kcc.columnIndices(), (std::vector<std::size_t>{0U, 1U, 0U, 1U}));
    EXPECT_EQ(blocks.kcc.values(), (std::vector<double>{12.0, 15.0, 42.0, 45.0}));

    EXPECT_EQ(blocks.kfc.values()[3U], 0.0);
    EXPECT_TRUE(blocks.kff.validate().isOk());
    EXPECT_TRUE(blocks.kfc.validate().isOk());
    EXPECT_TRUE(blocks.kcf.validate().isOk());
    EXPECT_TRUE(blocks.kcc.validate().isOk());
}

TEST(EssentialConstraints, HandlesNoAllAndMixedConstraints) {
    const auto full = makeMatrix(6U, 6U, sequentialDense(6U));

    const auto noConstraints = makeDofs({});
    auto none = fesa::EssentialConstraints::partition(full, noConstraints);
    ASSERT_TRUE(none.hasValue());
    expectShape(none.value().kff, 6U, 6U);
    expectShape(none.value().kfc, 6U, 0U);
    expectShape(none.value().kcf, 0U, 6U);
    expectShape(none.value().kcc, 0U, 0U);
    EXPECT_EQ(none.value().kff.values(), full.values());

    const auto allConstraints = makeDofs({{"1", 1, 6, 1.0, {{}, 12U}}});
    auto all = fesa::EssentialConstraints::partition(full, allConstraints);
    ASSERT_TRUE(all.hasValue());
    expectShape(all.value().kff, 0U, 0U);
    expectShape(all.value().kfc, 0U, 6U);
    expectShape(all.value().kcf, 6U, 0U);
    expectShape(all.value().kcc, 6U, 6U);
    EXPECT_EQ(all.value().kcc.values(), full.values());

    const auto mixedConstraints = makeDofs({{"1", 3, 4, 0.0, {{}, 12U}}});
    auto mixed = fesa::EssentialConstraints::partition(full, mixedConstraints);
    ASSERT_TRUE(mixed.hasValue());
    expectShape(mixed.value().kff, 4U, 4U);
    expectShape(mixed.value().kfc, 4U, 2U);
    expectShape(mixed.value().kcf, 2U, 4U);
    expectShape(mixed.value().kcc, 2U, 2U);
}

// MITC4-DOF-003
TEST(EssentialConstraints, PreservesShellSizedNoMixedAndAllConstraintRoundTrips) {
    const auto full = makeMatrix(24U, 24U, sequentialDense(24U));

    const auto noConstraints = makeShellSizedDofs({});
    auto none = fesa::EssentialConstraints::partition(full, noConstraints);
    ASSERT_TRUE(none.hasValue());
    expectShape(none.value().kff, 24U, 24U);
    expectShape(none.value().kfc, 24U, 0U);
    expectShape(none.value().kcf, 0U, 24U);
    expectShape(none.value().kcc, 0U, 0U);

    const auto mixedConstraints = makeShellSizedDofs({
        {"1", 1, 6, 0.0, {{}, 12U}},
        {"4", 2, 2, 2.5, {{}, 13U}}});
    auto mixed = fesa::EssentialConstraints::partition(full, mixedConstraints);
    ASSERT_TRUE(mixed.hasValue());
    expectShape(mixed.value().kff, 17U, 17U);
    expectShape(mixed.value().kfc, 17U, 7U);
    expectShape(mixed.value().kcf, 7U, 17U);
    expectShape(mixed.value().kcc, 7U, 7U);

    fesa::Vector mixedFull{24U};
    for (std::size_t index = 0U; index < mixedFull.size(); ++index) {
        mixedFull[index] = static_cast<double>(index) + 0.5;
    }
    for (std::size_t index = 0U;
         index < mixedConstraints.constrainedDofCount();
         ++index) {
        mixedFull[mixedConstraints.constrainedDofs()[index]] =
            mixedConstraints.prescribedValues()[index];
    }
    const auto mixedFree =
        fesa::EssentialConstraints::gatherFree(mixedFull, mixedConstraints);
    const auto mixedReconstructed =
        fesa::EssentialConstraints::reconstructFull(
            mixedFree, mixedConstraints.prescribedValues(), mixedConstraints);
    ASSERT_EQ(mixedReconstructed.size(), mixedFull.size());
    for (std::size_t index = 0U; index < mixedFull.size(); ++index) {
        EXPECT_DOUBLE_EQ(mixedReconstructed[index], mixedFull[index]);
    }

    const auto allConstraints = makeShellSizedDofs({
        {"1", 1, 6, 1.0, {{}, 14U}},
        {"2", 1, 6, 2.0, {{}, 15U}},
        {"3", 1, 6, 3.0, {{}, 16U}},
        {"4", 1, 6, 4.0, {{}, 17U}}});
    auto all = fesa::EssentialConstraints::partition(full, allConstraints);
    ASSERT_TRUE(all.hasValue());
    expectShape(all.value().kff, 0U, 0U);
    expectShape(all.value().kfc, 0U, 24U);
    expectShape(all.value().kcf, 24U, 0U);
    expectShape(all.value().kcc, 24U, 24U);

    const auto allReconstructed =
        fesa::EssentialConstraints::reconstructFull(
            fesa::Vector{0U},
            allConstraints.prescribedValues(),
            allConstraints);
    ASSERT_EQ(allReconstructed.size(), 24U);
    for (std::size_t node = 0U; node < 4U; ++node) {
        for (std::size_t component = 0U; component < 6U; ++component) {
            EXPECT_DOUBLE_EQ(allReconstructed[node * 6U + component], node + 1.0);
        }
    }
}

TEST(EssentialConstraints, ReconstructsNonzeroPrescribedValues) {
    const auto dofs = makeDofs({
        {"1", 2, 2, 2.5, {{}, 12U}},
        {"1", 5, 5, -3.25, {{}, 13U}}});
    fesa::Vector full{6U};
    full[0U] = 10.0;
    full[1U] = 2.5;
    full[2U] = 20.0;
    full[3U] = 30.0;
    full[4U] = -3.25;
    full[5U] = 40.0;

    const auto free = fesa::EssentialConstraints::gatherFree(full, dofs);
    const auto constrained =
        fesa::EssentialConstraints::gatherConstrained(full, dofs);
    EXPECT_EQ(free.size(), 4U);
    EXPECT_DOUBLE_EQ(free[0U], 10.0);
    EXPECT_DOUBLE_EQ(free[1U], 20.0);
    EXPECT_DOUBLE_EQ(free[2U], 30.0);
    EXPECT_DOUBLE_EQ(free[3U], 40.0);
    EXPECT_EQ(constrained.size(), 2U);
    EXPECT_DOUBLE_EQ(constrained[0U], 2.5);
    EXPECT_DOUBLE_EQ(constrained[1U], -3.25);
    EXPECT_EQ(constrained[0U], dofs.prescribedValues()[0U]);
    EXPECT_EQ(constrained[1U], dofs.prescribedValues()[1U]);

    const auto reconstructed = fesa::EssentialConstraints::reconstructFull(
        free, dofs.prescribedValues(), dofs);
    ASSERT_EQ(reconstructed.size(), full.size());
    for (std::size_t index = 0U; index < full.size(); ++index) {
        EXPECT_DOUBLE_EQ(reconstructed[index], full[index]);
    }
}

TEST(EssentialConstraints, RejectsDimensionOrOrderMismatch) {
    const auto dofs = makeDofs({{"1", 2, 2, 1.0, {{}, 12U}}});
    const auto wrongSquare = makeMatrix(5U, 5U, sequentialDense(5U));
    auto wrongDimension =
        fesa::EssentialConstraints::partition(wrongSquare, dofs);
    ASSERT_FALSE(wrongDimension.hasValue());
    EXPECT_EQ(
        wrongDimension.status().failureCategory(),
        fesa::FailureCategory::model);
    ASSERT_EQ(wrongDimension.status().diagnostics().size(), 1U);
    EXPECT_EQ(
        wrongDimension.status().diagnostics()[0U].code,
        "invalid-constraint-dimensions");

    const auto rectangular = makeMatrix(
        6U, 5U, std::vector<double>(30U, 0.0));
    auto wrongOrder = fesa::EssentialConstraints::partition(rectangular, dofs);
    ASSERT_FALSE(wrongOrder.hasValue());
    EXPECT_EQ(
        wrongOrder.status().diagnostics()[0U].code,
        "invalid-constraint-dimensions");

    EXPECT_THROW(
        static_cast<void>(fesa::EssentialConstraints::gatherFree(
            fesa::Vector{5U}, dofs)),
        std::invalid_argument);
    EXPECT_THROW(
        static_cast<void>(fesa::EssentialConstraints::gatherConstrained(
            fesa::Vector{7U}, dofs)),
        std::invalid_argument);
    EXPECT_THROW(
        static_cast<void>(fesa::EssentialConstraints::reconstructFull(
            fesa::Vector{4U}, fesa::Vector{2U}, dofs)),
        std::invalid_argument);
}