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

#include <gtest/gtest.h>

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

namespace {

fesa::ModelDefinition makeDefinition() {
    const std::filesystem::path source{"models/sparse-assembly.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"}, {2.0, 0.0, 0.0}, {source, 11U}},
        {{"Beam-1", 3, "3"}, {5.0, 0.0, 0.0}, {source, 12U}}};
    definition.materials = {
        {"Material", 120.0, 0.25, {source, 20U}}};
    definition.sections = {{
        "Section",
        2.0,
        1.5,
        0.0,
        0.75,
        0.5,
        {0.0, 1.0, 0.0},
        {},
        {source, 30U}}};
    definition.elements = {
        {{"Beam-1", 10, "10"}, {0U, 1U}, 0U, 0U, {source, 40U}},
        {{"Beam-1", 20, "20"}, {1U, 2U}, 0U, 0U, {source, 41U}}};
    definition.steps = {{
        "Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {source, 50U}}};
    return definition;
}

fesa::ModelDefinition makeShellDefinition(
    const fesa::ShellSourceElementType sourceType,
    const bool twoElements = false) {
    const std::filesystem::path source{"models/shell-sparse-assembly.inp"};
    fesa::ModelDefinition definition{};
    definition.sourcePath = source;
    definition.sourceContentIdentity = "fnv1a64:fedcba9876543210";
    if (twoElements) {
        definition.nodes = {
            {{"Shell-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 10U}},
            {{"Shell-1", 2, "2"}, {1.0, 0.0, 0.0}, {source, 11U}},
            {{"Shell-1", 3, "3"}, {2.0, 0.0, 0.0}, {source, 12U}},
            {{"Shell-1", 4, "4"}, {0.0, 1.0, 0.0}, {source, 13U}},
            {{"Shell-1", 5, "5"}, {1.0, 1.0, 0.0}, {source, 14U}},
            {{"Shell-1", 6, "6"}, {2.0, 1.0, 0.0}, {source, 15U}}};
        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.shellElements = {
            {{"Shell-1", 10, "10"}, sourceType, {0U, 1U, 4U, 3U},
             0U, 0U, {source, 40U}},
            {{"Shell-1", 20, "20"}, sourceType, {1U, 2U, 5U, 4U},
             0U, 0U, {source, 41U}}};
    } else {
        // A YZ-plane fixture catches any accidental global-Z director assumption.
        definition.nodes = {
            {{"Shell-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 10U}},
            {{"Shell-1", 2, "2"}, {0.0, 1.0, 0.0}, {source, 11U}},
            {{"Shell-1", 3, "3"}, {0.0, 1.0, 1.0}, {source, 12U}},
            {{"Shell-1", 4, "4"}, {0.0, 0.0, 1.0}, {source, 13U}}};
        for (std::size_t node = 0U; node < definition.nodes.size(); ++node) {
            definition.shellNodeInitialFrames.push_back({
                static_cast<fesa::EntityIndex>(node),
                {1.0, 0.0, 0.0},
                {0.0, 1.0, 0.0},
                {0.0, 0.0, 1.0}});
        }
        definition.shellElements = {{
            {"Shell-1", 10, "10"}, sourceType, {0U, 1U, 2U, 3U},
            0U, 0U, {source, 40U}}};
    }
    definition.materials = {
        {"Material", 120.0, 0.25, {source, 20U}}};
    definition.shellSections = {
        {"ShellSection", 0.2, 0U, {source, 30U}}};
    definition.steps = {{
        "Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {source, 50U}}};
    return definition;
}

fesa::Result<fesa::Mitc4Stiffness> directShellStiffness(
    const fesa::Domain& domain,
    const fesa::EntityIndex elementIndex) {
    const auto& definition = domain.shellElements().at(elementIndex);
    std::array<const fesa::Node*, 4> nodes{};
    std::array<std::array<double, 3>, 4> directors{};
    for (std::size_t node = 0U; node < definition.nodeIndices.size(); ++node) {
        const fesa::EntityIndex nodeIndex = definition.nodeIndices[node];
        nodes[node] = &domain.nodes().at(nodeIndex);
        directors[node] = domain.shellNodeInitialFrames().at(nodeIndex).director;
    }
    auto shell = fesa::Mitc4Shell::create(
        nodes,
        directors,
        domain.shellSections().at(definition.sectionIndex),
        domain.materials().at(definition.materialIndex));
    if (!shell.hasValue()) {
        return fesa::Result<fesa::Mitc4Stiffness>::failure(shell.status());
    }
    return shell.value().stiffness();
}

fesa::Result<fesa::SparseMatrix> assembleShell(
    const fesa::ShellSourceElementType sourceType,
    const fesa::ParallelFor& parallelFor,
    const bool twoElements = false) {
    auto domain = fesa::Domain::create(
        makeShellDefinition(sourceType, twoElements));
    if (!domain.hasValue()) {
        return fesa::Result<fesa::SparseMatrix>::failure(domain.status());
    }
    auto model = fesa::AnalysisModel::create(domain.value());
    if (!model.hasValue()) {
        return fesa::Result<fesa::SparseMatrix>::failure(model.status());
    }
    auto dofs = fesa::DofManager::create(model.value());
    if (!dofs.hasValue()) {
        return fesa::Result<fesa::SparseMatrix>::failure(dofs.status());
    }
    return fesa::SparseAssembler::assembleStiffness(
        model.value(), dofs.value(), parallelFor);
}

template<class T>
bool byteIdentical(const std::vector<T>& left, const std::vector<T>& right) {
    return left.size() == right.size() &&
           (left.empty() ||
            std::memcmp(
                left.data(), right.data(), left.size() * sizeof(T)) == 0);
}

double entry(
    const fesa::SparseMatrix& matrix,
    const std::size_t row,
    const std::size_t column) {
    const auto begin = matrix.columnIndices().begin() + matrix.rowOffsets()[row];
    const auto end = matrix.columnIndices().begin() + matrix.rowOffsets()[row + 1U];
    const auto found = std::lower_bound(begin, end, column);
    if (found == end || *found != column) {
        return 0.0;
    }
    return matrix.values()[static_cast<std::size_t>(
        std::distance(matrix.columnIndices().begin(), found))];
}

class ReverseParallelFor final : public fesa::ParallelFor {
public:
    void execute(
        const std::size_t count,
        const std::function<void(std::size_t)>& body) const override {
        ++calls_;
        observedCount_ = count;
        for (std::size_t index = count; index > 0U; --index) {
            body(index - 1U);
        }
    }

    std::size_t calls() const noexcept {
        return calls_;
    }

    std::size_t observedCount() const noexcept {
        return observedCount_;
    }

private:
    mutable std::size_t calls_{0U};
    mutable std::size_t observedCount_{0U};
};

void expectByteIdentical(
    const fesa::SparseMatrix& actual,
    const fesa::SparseMatrix& expected) {
    EXPECT_TRUE(byteIdentical(actual.rowOffsets(), expected.rowOffsets()));
    EXPECT_TRUE(byteIdentical(actual.columnIndices(), expected.columnIndices()));
    EXPECT_TRUE(byteIdentical(actual.values(), expected.values()));
}

TEST(SparseAssembly, SerialTbbAndRepeatedRunsAreByteIdentical) {
    auto domainResult = fesa::Domain::create(makeDefinition());
    ASSERT_TRUE(domainResult.hasValue());
    auto modelResult = fesa::AnalysisModel::create(domainResult.value());
    ASSERT_TRUE(modelResult.hasValue());
    auto dofsResult = fesa::DofManager::create(modelResult.value());
    ASSERT_TRUE(dofsResult.hasValue());

    fesa::SerialParallelFor serialExecutor;
    fesa::TbbParallelFor tbbExecutor;
    ReverseParallelFor reverseExecutor;
    auto serial = fesa::SparseAssembler::assembleStiffness(
        modelResult.value(), dofsResult.value(), serialExecutor);
    auto tbb = fesa::SparseAssembler::assembleStiffness(
        modelResult.value(), dofsResult.value(), tbbExecutor);
    auto reversed = fesa::SparseAssembler::assembleStiffness(
        modelResult.value(), dofsResult.value(), reverseExecutor);
    ASSERT_TRUE(serial.hasValue());
    ASSERT_TRUE(tbb.hasValue());
    ASSERT_TRUE(reversed.hasValue());

    EXPECT_EQ(reverseExecutor.calls(), 1U);
    EXPECT_EQ(reverseExecutor.observedCount(), 2U);
    EXPECT_EQ(serial.value().rows(), 18U);
    EXPECT_EQ(serial.value().columns(), 18U);
    EXPECT_EQ(serial.value().rowOffsets(), dofsResult.value().sparsePattern().rowOffsets);
    EXPECT_EQ(
        serial.value().columnIndices(),
        dofsResult.value().sparsePattern().columnIndices);
    EXPECT_TRUE(serial.value().validate().isOk());
    expectByteIdentical(tbb.value(), serial.value());
    expectByteIdentical(reversed.value(), serial.value());

    for (std::size_t repetition = 0U; repetition < 8U; ++repetition) {
        auto repeated = fesa::SparseAssembler::assembleStiffness(
            modelResult.value(), dofsResult.value(), tbbExecutor);
        ASSERT_TRUE(repeated.hasValue());
        expectByteIdentical(repeated.value(), serial.value());
    }

    for (std::size_t row = 0U; row < serial.value().rows(); ++row) {
        for (std::size_t column = 0U;
             column < serial.value().columns();
             ++column) {
            EXPECT_DOUBLE_EQ(
                entry(serial.value(), row, column),
                entry(serial.value(), column, row));
        }
    }

    EXPECT_NEAR(entry(serial.value(), 0U, 0U), 120.0, 1.0e-12);
    EXPECT_NEAR(entry(serial.value(), 0U, 6U), -120.0, 1.0e-12);
    EXPECT_NEAR(entry(serial.value(), 6U, 6U), 200.0, 1.0e-12);
    EXPECT_NEAR(entry(serial.value(), 6U, 12U), -80.0, 1.0e-12);
    EXPECT_NEAR(entry(serial.value(), 12U, 12U), 80.0, 1.0e-12);
}

TEST(
    SparseAssembly,
    AssemblesFourNodeTwentyFourDofKernelAndPreservesDiagonalSlots) {
    auto domain = fesa::Domain::create(
        makeShellDefinition(fesa::ShellSourceElementType::s4));
    ASSERT_TRUE(domain.hasValue());
    auto model = fesa::AnalysisModel::create(domain.value());
    ASSERT_TRUE(model.hasValue());
    auto dofs = fesa::DofManager::create(model.value());
    ASSERT_TRUE(dofs.hasValue());
    fesa::SerialParallelFor serialExecutor;

    auto assembled = fesa::SparseAssembler::assembleStiffness(
        model.value(), dofs.value(), serialExecutor);
    auto expected = directShellStiffness(domain.value(), 0U);
    ASSERT_TRUE(assembled.hasValue());
    ASSERT_TRUE(expected.hasValue());

    EXPECT_EQ(assembled.value().rows(), 24U);
    EXPECT_EQ(assembled.value().columns(), 24U);
    EXPECT_EQ(assembled.value().values().size(), 24U * 24U);
    EXPECT_EQ(
        assembled.value().rowOffsets(),
        dofs.value().sparsePattern().rowOffsets);
    EXPECT_EQ(
        assembled.value().columnIndices(),
        dofs.value().sparsePattern().columnIndices);
    for (std::size_t row = 0U; row < 24U; ++row) {
        const auto begin = assembled.value().columnIndices().begin() +
                           assembled.value().rowOffsets()[row];
        const auto end = assembled.value().columnIndices().begin() +
                         assembled.value().rowOffsets()[row + 1U];
        EXPECT_NE(std::lower_bound(begin, end, row), end);
        for (std::size_t column = 0U; column < 24U; ++column) {
            EXPECT_DOUBLE_EQ(
                entry(assembled.value(), row, column),
                expected.value().stabilizedGlobal24(row, column));
        }
    }
}

TEST(SparseAssembly, ShellSerialTbbReverseAndRepeatedRunsAreByteIdentical) {
    fesa::SerialParallelFor serialExecutor;
    fesa::TbbParallelFor tbbExecutor;
    ReverseParallelFor reverseExecutor;
    auto serial = assembleShell(
        fesa::ShellSourceElementType::s4, serialExecutor, true);
    auto tbb = assembleShell(
        fesa::ShellSourceElementType::s4, tbbExecutor, true);
    auto reversed = assembleShell(
        fesa::ShellSourceElementType::s4, reverseExecutor, true);
    ASSERT_TRUE(serial.hasValue());
    ASSERT_TRUE(tbb.hasValue());
    ASSERT_TRUE(reversed.hasValue());

    EXPECT_EQ(reverseExecutor.calls(), 1U);
    EXPECT_EQ(reverseExecutor.observedCount(), 2U);
    expectByteIdentical(tbb.value(), serial.value());
    expectByteIdentical(reversed.value(), serial.value());
    for (std::size_t repetition = 0U; repetition < 8U; ++repetition) {
        auto repeated = assembleShell(
            fesa::ShellSourceElementType::s4, tbbExecutor, true);
        ASSERT_TRUE(repeated.hasValue());
        expectByteIdentical(repeated.value(), serial.value());
    }
}

TEST(SparseAssembly, S4AndS4rSemanticFixturesAssembleIdenticalStiffness) {
    fesa::SerialParallelFor serialExecutor;
    auto s4 = assembleShell(
        fesa::ShellSourceElementType::s4, serialExecutor);
    auto s4r = assembleShell(
        fesa::ShellSourceElementType::s4r, serialExecutor);
    ASSERT_TRUE(s4.hasValue());
    ASSERT_TRUE(s4r.hasValue());
    EXPECT_TRUE(std::any_of(
        s4.value().values().begin(),
        s4.value().values().end(),
        [](const double value) { return value != 0.0; }));
    expectByteIdentical(s4r.value(), s4.value());
}

} // namespace