#include <fesa/fem/beam_frame.hpp>

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

#include <gtest/gtest.h>

namespace {

constexpr double kUnitTolerance =
    128.0 * std::numeric_limits<double>::epsilon();

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),
    };
}

void expect_vec_near(
    const fesa::Vec3 actual,
    const fesa::Vec3 expected,
    const double tolerance = kUnitTolerance) {
    EXPECT_NEAR(actual.x, expected.x, tolerance);
    EXPECT_NEAR(actual.y, expected.y, tolerance);
    EXPECT_NEAR(actual.z, expected.z, tolerance);
}

void expect_error(
    const fesa::BeamFrameResult& result,
    const std::string_view code) {
    EXPECT_FALSE(result.frame.has_value());
    ASSERT_EQ(result.diagnostics.size(), 1);
    EXPECT_EQ(result.diagnostics[0].stage, fesa::DiagnosticStage::model);
    EXPECT_EQ(result.diagnostics[0].severity, fesa::Severity::error);
    EXPECT_EQ(result.diagnostics[0].code, code);
}

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

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

TEST(BeamFrame, AlignsLocalAxesForGlobalXAxisBeam) {
    const auto result = fesa::make_beam_frame(
        {0.0, 0.0, 0.0},
        {2.0, 0.0, 0.0},
        {0.0, 3.0, 0.0});

    ASSERT_TRUE(result.frame.has_value());
    EXPECT_TRUE(result.diagnostics.empty());
    expect_vec_near(result.frame->ex, {1.0, 0.0, 0.0});
    expect_vec_near(result.frame->ey, {0.0, 1.0, 0.0});
    expect_vec_near(result.frame->ez, {0.0, 0.0, 1.0});
}

TEST(BeamFrame, ProducesRightHandedOrthonormalBasisForArbitraryAxis) {
    const auto result = fesa::make_beam_frame(
        {1.0, 2.0, 3.0},
        {3.0, 4.0, 4.0},
        {0.0, 1.0, 0.0});

    ASSERT_TRUE(result.frame.has_value());
    const fesa::BeamFrame& frame = *result.frame;
    const double inverse_sqrt_five = 1.0 / std::sqrt(5.0);
    expect_vec_near(frame.ex, {2.0 / 3.0, 2.0 / 3.0, 1.0 / 3.0});
    expect_vec_near(
        frame.ey,
        {
            -4.0 * inverse_sqrt_five / 3.0,
            5.0 * inverse_sqrt_five / 3.0,
            -2.0 * inverse_sqrt_five / 3.0,
        });
    expect_vec_near(
        frame.ez,
        {-inverse_sqrt_five, 0.0, 2.0 * inverse_sqrt_five});

    EXPECT_NEAR(dot(frame.ex, frame.ex), 1.0, kUnitTolerance);
    EXPECT_NEAR(dot(frame.ey, frame.ey), 1.0, kUnitTolerance);
    EXPECT_NEAR(dot(frame.ez, frame.ez), 1.0, kUnitTolerance);
    EXPECT_NEAR(dot(frame.ex, frame.ey), 0.0, kUnitTolerance);
    EXPECT_NEAR(dot(frame.ey, frame.ez), 0.0, kUnitTolerance);
    EXPECT_NEAR(dot(frame.ez, frame.ex), 0.0, kUnitTolerance);
    EXPECT_NEAR(dot(frame.ex, cross(frame.ey, frame.ez)), 1.0, kUnitTolerance);
}

TEST(BeamFrame, RotatesCovariantlyWithBeamGeometry) {
    constexpr fesa::Vec3 first{1.0, -2.0, 0.5};
    constexpr fesa::Vec3 second{4.0, 1.0, 2.5};
    constexpr fesa::Vec3 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;

    const auto original =
        fesa::make_beam_frame(first, second, orientation);
    const auto rotated = fesa::make_beam_frame(
        rotate_about_axis(first, rotation_axis, angle),
        rotate_about_axis(second, rotation_axis, angle),
        rotate_about_axis(orientation, rotation_axis, angle));

    ASSERT_TRUE(original.frame.has_value());
    ASSERT_TRUE(rotated.frame.has_value());
    expect_vec_near(
        rotated.frame->ex,
        rotate_about_axis(original.frame->ex, rotation_axis, angle),
        256.0 * std::numeric_limits<double>::epsilon());
    expect_vec_near(
        rotated.frame->ey,
        rotate_about_axis(original.frame->ey, rotation_axis, angle),
        256.0 * std::numeric_limits<double>::epsilon());
    expect_vec_near(
        rotated.frame->ez,
        rotate_about_axis(original.frame->ez, rotation_axis, angle),
        256.0 * std::numeric_limits<double>::epsilon());
}

TEST(BeamFrame, IsInvariantToFiniteOrientationScale) {
    const double maximum = std::numeric_limits<double>::max();
    const auto unit_scale = fesa::make_beam_frame(
        {0.0, 0.0, 0.0},
        {1.0, 1.0, 1.0},
        {1.0, -1.0, 0.0});
    const auto maximum_scale = fesa::make_beam_frame(
        {0.0, 0.0, 0.0},
        {1.0, 1.0, 1.0},
        {maximum, -maximum, 0.0});

    ASSERT_TRUE(unit_scale.frame.has_value());
    ASSERT_TRUE(maximum_scale.frame.has_value());
    expect_vec_near(maximum_scale.frame->ex, unit_scale.frame->ex);
    expect_vec_near(maximum_scale.frame->ey, unit_scale.frame->ey);
    expect_vec_near(maximum_scale.frame->ez, unit_scale.frame->ez);
}

TEST(BeamFrame, RejectsZeroLengthAxis) {
    expect_error(
        fesa::make_beam_frame(
            {1.0, 2.0, 3.0},
            {1.0, 2.0, 3.0},
            {0.0, 1.0, 0.0}),
        "model.zero_length_element");
}

TEST(BeamFrame, RejectsZeroOrientation) {
    expect_error(
        fesa::make_beam_frame(
            {0.0, 0.0, 0.0},
            {1.0, 0.0, 0.0},
            {0.0, 0.0, 0.0}),
        "model.invalid_orientation");
}

TEST(BeamFrame, RejectsOrientationParallelToElementAxis) {
    expect_error(
        fesa::make_beam_frame(
            {0.0, 0.0, 0.0},
            {2.0, 0.0, 0.0},
            {5.0, 0.0, 0.0}),
        "model.invalid_orientation");
}

TEST(BeamFrame, RejectsNumericallyParallelOrientationAtLargeScale) {
    constexpr double transverse_ratio =
        32.0 * std::numeric_limits<double>::epsilon();
    expect_error(
        fesa::make_beam_frame(
            {0.0, 0.0, 0.0},
            {2.0, 0.0, 0.0},
            {1.0e200, transverse_ratio * 1.0e200, 0.0}),
        "model.invalid_orientation");
}

TEST(BeamTransformation, IsIdentityForGlobalAlignedFrame) {
    const fesa::Matrix12 transformation = fesa::beam_transformation({
        {1.0, 0.0, 0.0},
        {0.0, 1.0, 0.0},
        {0.0, 0.0, 1.0},
    });

    for (std::size_t row = 0; row < transformation.size(); ++row) {
        for (std::size_t column = 0;
             column < transformation[row].size();
             ++column) {
            EXPECT_DOUBLE_EQ(
                transformation[row][column],
                row == column ? 1.0 : 0.0);
        }
    }
}

TEST(BeamTransformation, MapsGlobalComponentsIntoLocalDofOrder) {
    const fesa::Matrix12 transformation = fesa::beam_transformation({
        {0.0, 1.0, 0.0},
        {0.0, 0.0, 1.0},
        {1.0, 0.0, 0.0},
    });
    const std::array<double, 12> global{
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0,
        7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
    };

    EXPECT_EQ(
        multiply(transformation, global),
        (std::array<double, 12>{
            2.0, 3.0, 1.0, 5.0, 6.0, 4.0,
            8.0, 9.0, 7.0, 11.0, 12.0, 10.0,
        }));
}

TEST(BeamTransformation, PreservesVectorEnergyForArbitraryFrame) {
    const auto frame_result = fesa::make_beam_frame(
        {-2.0, 1.0, 3.0},
        {1.0, 5.0, 5.0},
        {2.0, -1.0, 4.0});
    ASSERT_TRUE(frame_result.frame.has_value());
    const fesa::Matrix12 transformation =
        fesa::beam_transformation(*frame_result.frame);
    const std::array<double, 12> global{
        3.0, -1.0, 2.0, 0.5, 4.0, -2.0,
        -3.0, 6.0, 1.5, 2.5, -0.25, 5.0,
    };

    const std::array<double, 12> local =
        multiply(transformation, global);
    const double global_energy = squared_norm(global);
    const double tolerance =
        256.0 * std::numeric_limits<double>::epsilon() *
        std::max(1.0, global_energy);
    EXPECT_NEAR(squared_norm(local), global_energy, tolerance);
}

}  // namespace
