#include <fesa/elements/beam/beam3d2.hpp>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <limits>

#include <gtest/gtest.h>

namespace {

constexpr double kYoung = 210.0;
constexpr double kPoisson = 0.3;
constexpr double kArea = 0.4;
constexpr double kIy = 0.03;
constexpr double kIz = 0.05;
constexpr double kTorsionJ = 0.02;
constexpr double kShearAreaY = 0.25;
constexpr double kShearAreaZ = 0.2;

fesa::Beam3D2Input make_x_axis_input(const double length) {
    return {
        {{{0.0, 0.0, 0.0}, {length, 0.0, 0.0}}},
        {fesa::NodeId{1}, fesa::NodeId{2}},
        {fesa::MaterialId{1}, "elastic", kYoung, kPoisson},
        {
            fesa::SectionId{1},
            "section",
            kArea,
            kIy,
            kIz,
            kTorsionJ,
            kShearAreaY,
            kShearAreaZ,
            fesa::ShearPropertySource::input,
            {0.0, 1.0, 0.0},
            {},
        },
    };
}

double quadratic_form(
    const fesa::Matrix12& matrix,
    const std::array<double, 12>& vector) {
    double result = 0.0;
    for (std::size_t row = 0; row < matrix.size(); ++row) {
        for (std::size_t column = 0; column < matrix[row].size(); ++column) {
            result += vector[row] * matrix[row][column] * vector[column];
        }
    }
    return result;
}

double strain_energy(
    const fesa::Matrix12& matrix,
    const std::array<double, 12>& vector) {
    return 0.5 * quadratic_form(matrix, vector);
}

double maximum_abs_entry(const fesa::Matrix12& matrix) {
    double maximum = 0.0;
    for (const auto& row : matrix) {
        for (const double value : row) {
            maximum = std::max(maximum, std::abs(value));
        }
    }
    return maximum;
}

double squared_norm(const std::array<double, 12>& vector) {
    double result = 0.0;
    for (const double value : vector) {
        result += value * value;
    }
    return result;
}

double dot(const fesa::Vec3 first, const fesa::Vec3 second) {
    return first.x * second.x + first.y * second.y + first.z * second.z;
}

fesa::Vec3 cross(const fesa::Vec3 first, const fesa::Vec3 second) {
    return {
        first.y * second.z - first.z * second.y,
        first.z * second.x - first.x * second.z,
        first.x * second.y - first.y * second.x,
    };
}

fesa::Vec3 rotate_about_axis(
    const fesa::Vec3 value,
    const fesa::Vec3 unit_axis,
    const double angle) {
    const double cosine = std::cos(angle);
    const double sine = std::sin(angle);
    const fesa::Vec3 axis_cross_value = cross(unit_axis, value);
    const double axis_projection = dot(unit_axis, value);
    return {
        value.x * cosine + axis_cross_value.x * sine +
            unit_axis.x * axis_projection * (1.0 - cosine),
        value.y * cosine + axis_cross_value.y * sine +
            unit_axis.y * axis_projection * (1.0 - cosine),
        value.z * cosine + axis_cross_value.z * sine +
            unit_axis.z * axis_projection * (1.0 - cosine),
    };
}

std::array<double, 12> rotate_dofs(
    const std::array<double, 12>& values,
    const fesa::Vec3 unit_axis,
    const double angle) {
    std::array<double, 12> rotated{};
    for (const std::size_t offset : std::array<std::size_t, 4>{0, 3, 6, 9}) {
        const fesa::Vec3 value{
            values[offset],
            values[offset + 1],
            values[offset + 2],
        };
        const fesa::Vec3 rotated_value =
            rotate_about_axis(value, unit_axis, angle);
        rotated[offset] = rotated_value.x;
        rotated[offset + 1] = rotated_value.y;
        rotated[offset + 2] = rotated_value.z;
    }
    return rotated;
}

void expect_relative_near(
    const double actual,
    const double expected,
    const double relative_tolerance =
        512.0 * std::numeric_limits<double>::epsilon()) {
    const double scale = std::max(1.0, std::abs(expected));
    EXPECT_NEAR(actual, expected, relative_tolerance * scale);
}

TEST(Beam3D2, ProducesFiniteSymmetricLocalAndGlobalStiffness) {
    const auto result = fesa::compute_beam3d2(make_x_axis_input(2.5));

    ASSERT_TRUE(result.contribution.has_value());
    EXPECT_TRUE(result.diagnostics.empty());
    const auto& contribution = *result.contribution;
    for (std::size_t row = 0; row < 12; ++row) {
        for (std::size_t column = 0; column < 12; ++column) {
            EXPECT_TRUE(std::isfinite(
                contribution.local_stiffness[row][column]));
            EXPECT_TRUE(std::isfinite(
                contribution.global_stiffness[row][column]));
            EXPECT_DOUBLE_EQ(
                contribution.local_stiffness[row][column],
                contribution.local_stiffness[column][row]);
            EXPECT_DOUBLE_EQ(
                contribution.global_stiffness[row][column],
                contribution.global_stiffness[column][row]);
        }
    }
}

TEST(RigidBody, SixIndependentModesHaveZeroStrainEnergy) {
    constexpr double length = 3.0;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    const auto& stiffness = result.contribution->local_stiffness;

    const std::array<std::array<double, 12>, 6> modes{{
        {1.0, 0.0, 0.0, 0.0, 0.0, 0.0,
         1.0, 0.0, 0.0, 0.0, 0.0, 0.0},
        {0.0, 1.0, 0.0, 0.0, 0.0, 0.0,
         0.0, 1.0, 0.0, 0.0, 0.0, 0.0},
        {0.0, 0.0, 1.0, 0.0, 0.0, 0.0,
         0.0, 0.0, 1.0, 0.0, 0.0, 0.0},
        {0.0, 0.0, 0.0, 1.0, 0.0, 0.0,
         0.0, 0.0, 0.0, 1.0, 0.0, 0.0},
        {0.0, 0.0, 0.0, 0.0, 1.0, 0.0,
         0.0, 0.0, -length, 0.0, 1.0, 0.0},
        {0.0, 0.0, 0.0, 0.0, 0.0, 1.0,
         0.0, length, 0.0, 0.0, 0.0, 1.0},
    }};

    for (std::size_t mode = 0; mode < modes.size(); ++mode) {
        const double roundoff_bound =
            4096.0 * std::numeric_limits<double>::epsilon() *
            maximum_abs_entry(stiffness) * squared_norm(modes[mode]);
        EXPECT_NEAR(
            quadratic_form(stiffness, modes[mode]),
            0.0,
            roundoff_bound)
            << "rigid mode " << mode;
    }
}

TEST(Timoshenko, ReproducesAnalyticalAxialSubmatrix) {
    constexpr double length = 2.5;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    const auto& stiffness = result.contribution->local_stiffness;
    const double expected = kYoung * kArea / length;

    expect_relative_near(stiffness[0][0], expected);
    expect_relative_near(stiffness[0][6], -expected);
    expect_relative_near(stiffness[6][0], -expected);
    expect_relative_near(stiffness[6][6], expected);
}

TEST(Timoshenko, ReproducesAnalyticalTorsionalSubmatrix) {
    constexpr double length = 2.5;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    const auto& stiffness = result.contribution->local_stiffness;
    const double shear_modulus = kYoung / (2.0 * (1.0 + kPoisson));
    const double expected = shear_modulus * kTorsionJ / length;

    expect_relative_near(stiffness[3][3], expected);
    expect_relative_near(stiffness[3][9], -expected);
    expect_relative_near(stiffness[9][3], -expected);
    expect_relative_near(stiffness[9][9], expected);
}

TEST(Timoshenko, ReproducesConstantCurvatureEnergyAboutLocalY) {
    constexpr double length = 2.5;
    constexpr double curvature = 0.4;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    std::array<double, 12> displacement{};
    displacement[4] = -0.5 * curvature * length;
    displacement[10] = 0.5 * curvature * length;

    const double expected =
        0.5 * kYoung * kIy * curvature * curvature * length;
    expect_relative_near(
        strain_energy(result.contribution->local_stiffness, displacement),
        expected);
}

TEST(Timoshenko, ReproducesConstantCurvatureEnergyAboutLocalZ) {
    constexpr double length = 2.5;
    constexpr double curvature = 0.4;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    std::array<double, 12> displacement{};
    displacement[5] = -0.5 * curvature * length;
    displacement[11] = 0.5 * curvature * length;

    const double expected =
        0.5 * kYoung * kIz * curvature * curvature * length;
    expect_relative_near(
        strain_energy(result.contribution->local_stiffness, displacement),
        expected);
}

TEST(Timoshenko, ReproducesConstantShearEnergyInLocalY) {
    constexpr double length = 2.5;
    constexpr double shear_strain = 0.04;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    std::array<double, 12> displacement{};
    displacement[1] = -0.5 * shear_strain * length;
    displacement[7] = 0.5 * shear_strain * length;

    const double shear_modulus = kYoung / (2.0 * (1.0 + kPoisson));
    const double expected =
        0.5 * shear_modulus * kShearAreaY *
        shear_strain * shear_strain * length;
    expect_relative_near(
        strain_energy(result.contribution->local_stiffness, displacement),
        expected);
}

TEST(Timoshenko, ReproducesConstantShearEnergyInLocalZ) {
    constexpr double length = 2.5;
    constexpr double shear_strain = 0.04;
    const auto result = fesa::compute_beam3d2(make_x_axis_input(length));
    ASSERT_TRUE(result.contribution.has_value());
    std::array<double, 12> displacement{};
    displacement[2] = -0.5 * shear_strain * length;
    displacement[8] = 0.5 * shear_strain * length;

    const double shear_modulus = kYoung / (2.0 * (1.0 + kPoisson));
    const double expected =
        0.5 * shear_modulus * kShearAreaZ *
        shear_strain * shear_strain * length;
    expect_relative_near(
        strain_energy(result.contribution->local_stiffness, displacement),
        expected);
}

TEST(Timoshenko, AvoidsShearLockingAcrossSlendernessSweep) {
    constexpr double length = 2.0;
    constexpr double curvature = 0.2;
    for (const double slenderness : std::array{2.0, 10.0, 100.0, 1000.0}) {
        fesa::Beam3D2Input input = make_x_axis_input(length);
        const double side = length / slenderness;
        input.section.area = side * side;
        input.section.iy = side * side * side * side / 12.0;
        input.section.iz = input.section.iy;
        input.section.shear_area_y = 5.0 * input.section.area / 6.0;
        input.section.shear_area_z = input.section.shear_area_y;

        const auto result = fesa::compute_beam3d2(input);
        ASSERT_TRUE(result.contribution.has_value());
        std::array<double, 12> displacement{};
        displacement[5] = -0.5 * curvature * length;
        displacement[11] = 0.5 * curvature * length;

        const double expected =
            0.5 * kYoung * input.section.iz *
            curvature * curvature * length;
        const double relative_tolerance =
            4096.0 * std::numeric_limits<double>::epsilon() *
            slenderness * slenderness;
        EXPECT_NEAR(
            strain_energy(
                result.contribution->local_stiffness,
                displacement),
            expected,
            relative_tolerance * expected)
            << "L/h=" << slenderness;
    }
}

TEST(Beam3D2, PreservesGlobalEnergyUnderRigidCoordinateRotation) {
    fesa::Beam3D2Input original = make_x_axis_input(2.5);
    original.coordinates = {{
        {1.0, -2.0, 0.5},
        {3.0, -0.5, 1.25},
    }};
    original.section.orientation = {-1.0, 2.0, 3.0};

    const double inverse_sqrt_fourteen = 1.0 / std::sqrt(14.0);
    const fesa::Vec3 rotation_axis{
        inverse_sqrt_fourteen,
        2.0 * inverse_sqrt_fourteen,
        3.0 * inverse_sqrt_fourteen,
    };
    constexpr double angle = 0.73;
    fesa::Beam3D2Input rotated = original;
    for (std::size_t node = 0; node < rotated.coordinates.size(); ++node) {
        rotated.coordinates[node] = rotate_about_axis(
            original.coordinates[node],
            rotation_axis,
            angle);
    }
    rotated.section.orientation = rotate_about_axis(
        original.section.orientation,
        rotation_axis,
        angle);

    const auto original_result = fesa::compute_beam3d2(original);
    const auto rotated_result = fesa::compute_beam3d2(rotated);
    ASSERT_TRUE(original_result.contribution.has_value());
    ASSERT_TRUE(rotated_result.contribution.has_value());

    const std::array<double, 12> original_displacement{
        0.1, -0.3, 0.2, 0.04, -0.02, 0.03,
        -0.2, 0.5, -0.1, -0.01, 0.06, -0.05,
    };
    const std::array<double, 12> rotated_displacement =
        rotate_dofs(original_displacement, rotation_axis, angle);
    const double original_energy = strain_energy(
        original_result.contribution->global_stiffness,
        original_displacement);
    const double rotated_energy = strain_energy(
        rotated_result.contribution->global_stiffness,
        rotated_displacement);

    expect_relative_near(
        rotated_energy,
        original_energy,
        4096.0 * std::numeric_limits<double>::epsilon());
}

TEST(BeamRecovery, RecoversPureAxialStrainAndForceAtBothEnds) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[6] = 0.4;

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    EXPECT_DOUBLE_EQ(results[0].xi, -1.0);
    EXPECT_EQ(results[0].end_node, fesa::NodeId{101});
    EXPECT_DOUBLE_EQ(results[1].xi, 1.0);
    EXPECT_EQ(results[1].end_node, fesa::NodeId{202});
    for (const auto& result : results) {
        expect_relative_near(result.section_strain[0], 0.2);
        expect_relative_near(result.section_force[0], 16.8);
        expect_relative_near(result.centroid_sigma_xx, 42.0);
        for (std::size_t component = 1; component < 6; ++component) {
            expect_relative_near(result.section_strain[component], 0.0);
            expect_relative_near(result.section_force[component], 0.0);
        }
        EXPECT_TRUE(result.sigma_xx.empty());
    }
}

TEST(BeamRecovery, RecoversPureTorsionAtBothEnds) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[9] = 0.6;
    const double shear_modulus = kYoung / (2.0 * (1.0 + kPoisson));

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.section_strain[3], 0.3);
        expect_relative_near(
            result.section_force[3],
            shear_modulus * kTorsionJ * 0.3);
        for (const std::size_t component :
             std::array<std::size_t, 5>{0, 1, 2, 4, 5}) {
            expect_relative_near(result.section_strain[component], 0.0);
            expect_relative_near(result.section_force[component], 0.0);
        }
    }
}

TEST(BeamRecovery, UsesReducedIntegrationPointForShearAtBothEnds) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    std::array<double, 12> displacement{};
    displacement[7] = 1.2;
    displacement[5] = 0.2;
    displacement[11] = 0.6;
    displacement[8] = -0.4;
    displacement[4] = 0.1;
    displacement[10] = 0.5;
    const double shear_modulus = kYoung / (2.0 * (1.0 + kPoisson));

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.section_strain[1], 0.2);
        expect_relative_near(result.section_strain[2], 0.1);
        expect_relative_near(
            result.section_force[1],
            shear_modulus * kShearAreaY * 0.2);
        expect_relative_near(
            result.section_force[2],
            shear_modulus * kShearAreaZ * 0.1);
    }
}

TEST(BeamRecovery, UsesBeamFrameForGlobalDisplacements) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.coordinates = {{{0.0, 0.0, 0.0}, {0.0, 2.0, 0.0}}};
    input.section.orientation = {0.0, 0.0, 1.0};
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[7] = 0.4;
    displacement[10] = 0.6;
    const double shear_modulus = kYoung / (2.0 * (1.0 + kPoisson));

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.section_strain[0], 0.2);
        expect_relative_near(result.section_force[0], 16.8);
        expect_relative_near(result.section_strain[3], 0.3);
        expect_relative_near(
            result.section_force[3],
            shear_modulus * kTorsionJ * 0.3);
    }
}

TEST(SectionForce, RecoversPureBendingAboutLocalYAtBothEnds) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[4] = -0.4;
    displacement[10] = 0.4;

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.section_strain[4], 0.4);
        expect_relative_near(result.section_force[4], 2.52);
        for (const std::size_t component :
             std::array<std::size_t, 5>{0, 1, 2, 3, 5}) {
            expect_relative_near(result.section_strain[component], 0.0);
            expect_relative_near(result.section_force[component], 0.0);
        }
    }
}

TEST(SectionForce, RecoversPureBendingAboutLocalZAtBothEnds) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[5] = 0.25;
    displacement[11] = -0.25;

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.section_strain[5], -0.25);
        expect_relative_near(result.section_force[5], -2.625);
        for (const std::size_t component :
             std::array<std::size_t, 5>{0, 1, 2, 3, 4}) {
            expect_relative_near(result.section_strain[component], 0.0);
            expect_relative_near(result.section_force[component], 0.0);
        }
    }
}

TEST(SectionForce, PreservesBiaxialBendingSignsAtBothEnds) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[4] = -0.4;
    displacement[10] = 0.4;
    displacement[5] = 0.25;
    displacement[11] = -0.25;

    const auto results = fesa::recover_beam3d2(input, displacement, {});

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.section_force[4], 2.52);
        expect_relative_near(result.section_force[5], -2.625);
    }
}

TEST(CentroidStress, UsesAxialStressAndPreservesRecoveryPointOrder) {
    fesa::Beam3D2Input input = make_x_axis_input(2.0);
    input.node_ids = {fesa::NodeId{101}, fesa::NodeId{202}};
    std::array<double, 12> displacement{};
    displacement[6] = 0.2;
    displacement[4] = -0.2;
    displacement[10] = 0.2;
    displacement[5] = 0.3;
    displacement[11] = -0.3;
    const std::array<std::array<double, 2>, 3> recovery_points{{
        {2.0, 3.0},
        {-1.0, 4.0},
        {5.0, -2.0},
    }};

    const auto results = fesa::recover_beam3d2(
        input,
        displacement,
        recovery_points);

    ASSERT_EQ(results.size(), 2U);
    for (const auto& result : results) {
        expect_relative_near(result.centroid_sigma_xx, 21.0);
        ASSERT_EQ(result.sigma_xx.size(), recovery_points.size());
        expect_relative_near(result.sigma_xx[0], 273.0);
        expect_relative_near(result.sigma_xx[1], 126.0);
        expect_relative_near(result.sigma_xx[2], 252.0);
    }
}

}  // namespace
