#include <fesa/analysis/linear_static_analysis.hpp>

#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <utility>
#include <vector>

#include <fesa/assembly/serial_assembler.hpp>
#include <fesa/fem/dof_manager.hpp>
#include <fesa/model/domain_builder.hpp>
#include <fesa/results/result_database.hpp>

#include <gtest/gtest.h>

namespace {

constexpr std::size_t dofs_per_node = 6;

std::vector<fesa::PrescribedDof> fixed_dofs(
    const fesa::NodeId node,
    const double ux = 0.0) {
    std::vector<fesa::PrescribedDof> prescribed;
    prescribed.reserve(dofs_per_node);
    for (std::uint8_t dof = 1; dof <= dofs_per_node; ++dof) {
        prescribed.push_back({node, dof, dof == 1 ? ux : 0.0});
    }
    return prescribed;
}

fesa::Domain build_axial_domain(const bool all_constrained) {
    fesa::DomainBuilder builder;
    builder.add_node({
        fesa::NodeId{20},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        fesa::Vec3{0.0, 0.0, 0.0},
    });
    builder.add_node({
        fesa::NodeId{4},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 2},
        fesa::Vec3{4.0, 0.0, 0.0},
    });
    builder.add_material({
        fesa::MaterialId{0},
        "HandCalculated",
        100.0,
        0.25,
    });
    builder.add_section({
        fesa::SectionId{0},
        "General",
        2.0,
        1.0,
        1.0,
        1.0,
        2.0,
        2.0,
        fesa::ShearPropertySource::input,
        fesa::Vec3{0.0, 1.0, 0.0},
        {{1.0, 2.0}, {-1.0, -2.0}},
    });
    builder.add_beam_element({
        fesa::ElementId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        {fesa::NodeId{20}, fesa::NodeId{4}},
        fesa::MaterialId{0},
        fesa::SectionId{0},
    });

    std::vector<fesa::PrescribedDof> prescribed =
        fixed_dofs(fesa::NodeId{20});
    std::vector<fesa::PrescribedDof> end_constraints = fixed_dofs(
        fesa::NodeId{4}, all_constrained ? 0.1 : 0.0);
    if (!all_constrained) {
        end_constraints.erase(end_constraints.begin());
    }
    prescribed.insert(
        prescribed.end(),
        end_constraints.begin(),
        end_constraints.end());

    std::vector<fesa::NodalLoad> loads;
    if (!all_constrained) {
        loads.push_back({
            fesa::NodeId{4},
            {10.0, 0.0, 0.0, 0.0, 0.0, 0.0},
        });
    }
    builder.set_step({"Load", std::move(prescribed), std::move(loads)});

    auto result = std::move(builder).build();
    if (!result.domain.has_value()) {
        throw std::runtime_error{"Test Domain failed validation."};
    }
    return std::move(*result.domain);
}

std::vector<double> full_field(
    const fesa::NodalFrame& nodal,
    const fesa::DofManager& dofs,
    const bool reaction) {
    std::vector<double> field(dofs.full_dof_count(), 0.0);
    for (std::size_t node = 0; node < nodal.node_ids.size(); ++node) {
        const auto& values = reaction
            ? nodal.reaction[node]
            : nodal.displacement[node];
        for (std::size_t component = 0; component < dofs_per_node;
             ++component) {
            field[dofs.full_dof({
                nodal.node_ids[node],
                static_cast<fesa::NodeDof>(component),
            })] = values[component];
        }
    }
    return field;
}

std::vector<double> multiply(
    const fesa::SymmetricCsr& matrix,
    const std::vector<double>& vector) {
    std::vector<double> product(matrix.order, 0.0);
    for (std::size_t row = 0; row < matrix.order; ++row) {
        const std::size_t begin =
            static_cast<std::size_t>(matrix.row_offsets[row]);
        const std::size_t end =
            static_cast<std::size_t>(matrix.row_offsets[row + 1]);
        for (std::size_t entry = begin; entry < end; ++entry) {
            const std::size_t column = static_cast<std::size_t>(
                matrix.column_indices[entry]);
            const double value = matrix.values[entry];
            product[row] += value * vector[column];
            if (column != row) {
                product[column] += value * vector[row];
            }
        }
    }
    return product;
}

TEST(LinearStaticAnalysis, SolvesHandCalculatedAxialBeamInNodeIdOrder) {
    const fesa::Domain domain = build_axial_domain(false);

    const fesa::AnalysisRunResult run =
        fesa::LinearStaticAnalysis{}.run(domain);

    ASSERT_TRUE(run.succeeded);
    ASSERT_TRUE(run.results.has_value());
    EXPECT_TRUE(run.diagnostics.empty());
    EXPECT_TRUE(fesa::validate_result_database(*run.results).succeeded);
    EXPECT_EQ(run.results->schema_version, "1.0.0");
    ASSERT_EQ(run.results->steps.size(), 1);
    EXPECT_EQ(run.results->steps[0].name, "Load");
    ASSERT_EQ(run.results->steps[0].frames.size(), 1);
    const fesa::ResultFrame& frame = run.results->steps[0].frames[0];
    EXPECT_DOUBLE_EQ(frame.step_time, 1.0);
    EXPECT_TRUE(frame.diagnostics.empty());
    ASSERT_EQ(frame.nodal.node_ids.size(), 2);
    EXPECT_EQ(
        frame.nodal.node_ids,
        (std::vector<fesa::NodeId>{
            fesa::NodeId{4}, fesa::NodeId{20}}));
    EXPECT_NEAR(frame.nodal.displacement[0][0], 0.2, 1.0e-12);
    EXPECT_NEAR(frame.nodal.displacement[1][0], 0.0, 1.0e-12);
    EXPECT_NEAR(frame.nodal.reaction[0][0], 0.0, 1.0e-12);
    EXPECT_NEAR(frame.nodal.reaction[1][0], -10.0, 1.0e-12);
}

TEST(StaticEquilibrium, ReturnedFieldsSatisfyOriginalFullEquation) {
    const fesa::Domain domain = build_axial_domain(false);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);
    const fesa::EquationSystem original =
        fesa::assemble_serial(domain, dofs);

    const fesa::AnalysisRunResult run =
        fesa::LinearStaticAnalysis{}.run(domain);

    ASSERT_TRUE(run.succeeded);
    ASSERT_TRUE(run.results.has_value());
    const fesa::NodalFrame& nodal =
        run.results->steps[0].frames[0].nodal;
    const std::vector<double> displacement =
        full_field(nodal, dofs, false);
    const std::vector<double> reaction =
        full_field(nodal, dofs, true);
    const std::vector<double> internal =
        multiply(original.stiffness, displacement);
    for (std::size_t dof = 0; dof < internal.size(); ++dof) {
        EXPECT_NEAR(
            internal[dof] - original.force[dof] - reaction[dof],
            0.0,
            1.0e-12);
        EXPECT_TRUE(std::isfinite(displacement[dof]));
        EXPECT_TRUE(std::isfinite(reaction[dof]));
    }
}

TEST(LinearStaticAnalysis, SolvesAllConstrainedSystemWithoutPardiso) {
    const fesa::Domain domain = build_axial_domain(true);

    const fesa::AnalysisRunResult run =
        fesa::LinearStaticAnalysis{}.run(domain);

    ASSERT_TRUE(run.succeeded);
    ASSERT_TRUE(run.results.has_value());
    EXPECT_TRUE(run.diagnostics.empty());
    const fesa::NodalFrame& nodal =
        run.results->steps[0].frames[0].nodal;
    ASSERT_EQ(nodal.node_ids.size(), 2);
    ASSERT_EQ(nodal.node_ids[0], fesa::NodeId{4});
    EXPECT_NEAR(nodal.displacement[0][0], 0.1, 1.0e-12);
    EXPECT_NEAR(nodal.reaction[0][0], 5.0, 1.0e-12);
    EXPECT_NEAR(nodal.reaction[1][0], -5.0, 1.0e-12);
}

TEST(CompleteResultContract, PreservesOriginsConnectivityAndBeamRecovery) {
    const fesa::Domain domain = build_axial_domain(false);

    const fesa::AnalysisRunResult run =
        fesa::LinearStaticAnalysis{}.run(domain);

    ASSERT_TRUE(run.succeeded);
    ASSERT_TRUE(run.results.has_value());
    const fesa::ResultFrame& frame = run.results->steps[0].frames[0];
    EXPECT_EQ(
        frame.nodal.origins,
        (std::vector<fesa::EntityOrigin>{
            {"BeamPart", "Beam-1", 2},
            {"BeamPart", "Beam-1", 1},
        }));

    ASSERT_EQ(frame.element.beams.size(), 1U);
    const fesa::BeamElementFrame& beam = frame.element.beams[0];
    EXPECT_EQ(beam.element, fesa::ElementId{0});
    EXPECT_EQ(
        beam.origin,
        (fesa::EntityOrigin{"BeamPart", "Beam-1", 1}));
    EXPECT_DOUBLE_EQ(beam.local_frame.ex.x, 1.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ex.y, 0.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ex.z, 0.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ey.x, 0.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ey.y, 1.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ey.z, 0.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ez.x, 0.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ez.y, 0.0);
    EXPECT_DOUBLE_EQ(beam.local_frame.ez.z, 1.0);
    EXPECT_EQ(beam.end_results[0].end_node, fesa::NodeId{20});
    EXPECT_EQ(beam.end_results[1].end_node, fesa::NodeId{4});
    EXPECT_DOUBLE_EQ(beam.end_results[0].xi, -1.0);
    EXPECT_DOUBLE_EQ(beam.end_results[1].xi, 1.0);
    for (const fesa::BeamSectionResult& end : beam.end_results) {
        EXPECT_NEAR(end.section_strain[0], 0.05, 1.0e-12);
        EXPECT_NEAR(end.section_force[0], 10.0, 1.0e-12);
        EXPECT_NEAR(end.centroid_sigma_xx, 5.0, 1.0e-12);
        ASSERT_EQ(end.sigma_xx.size(), 2U);
        EXPECT_NEAR(end.sigma_xx[0], 5.0, 1.0e-12);
        EXPECT_NEAR(end.sigma_xx[1], 5.0, 1.0e-12);
    }
}

}  // namespace