#include "fesa/elements/finite_rotation.h"

#include <gtest/gtest.h>

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

namespace fesa {
namespace {

constexpr std::size_t kRotationDimension = 3U;

Matrix InitialFrame() {
  Matrix frame{kRotationDimension, kRotationDimension};
  frame(0U, 0U) = 1.0;
  frame(1U, 2U) = -1.0;
  frame(2U, 1U) = 1.0;
  return frame;
}

Matrix Transpose(const Matrix& matrix) {
  Matrix result{matrix.Columns(), matrix.Rows()};
  for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
    for (std::size_t column = 0U; column < matrix.Columns(); ++column) {
      result(column, row) = matrix(row, column);
    }
  }
  return result;
}

Vector3 Multiply(const Matrix& matrix, const Vector3& vector) {
  return Vector3{
      matrix(0U, 0U) * vector.X() + matrix(0U, 1U) * vector.Y() +
          matrix(0U, 2U) * vector.Z(),
      matrix(1U, 0U) * vector.X() + matrix(1U, 1U) * vector.Y() +
          matrix(1U, 2U) * vector.Z(),
      matrix(2U, 0U) * vector.X() + matrix(2U, 1U) * vector.Y() +
          matrix(2U, 2U) * vector.Z()};
}

double Determinant(const Matrix& matrix) {
  return matrix(0U, 0U) *
             (matrix(1U, 1U) * matrix(2U, 2U) -
              matrix(1U, 2U) * matrix(2U, 1U)) -
         matrix(0U, 1U) *
             (matrix(1U, 0U) * matrix(2U, 2U) -
              matrix(1U, 2U) * matrix(2U, 0U)) +
         matrix(0U, 2U) *
             (matrix(1U, 0U) * matrix(2U, 1U) -
              matrix(1U, 1U) * matrix(2U, 0U));
}

bool IsFinite(const Matrix& matrix) {
  for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
    for (std::size_t column = 0U; column < matrix.Columns(); ++column) {
      if (!std::isfinite(matrix(row, column))) {
        return false;
      }
    }
  }
  return true;
}

void ExpectMatrixNear(const Matrix& actual, const Matrix& expected,
                      const double tolerance = 1.0e-12) {
  ASSERT_EQ(actual.Rows(), expected.Rows());
  ASSERT_EQ(actual.Columns(), expected.Columns());
  for (std::size_t row = 0U; row < actual.Rows(); ++row) {
    for (std::size_t column = 0U; column < actual.Columns(); ++column) {
      EXPECT_NEAR(actual(row, column), expected(row, column), tolerance)
          << "matrix entry (" << row << ", " << column << ")";
    }
  }
}

void ExpectVectorNear(const Vector3& actual, const Vector3& expected,
                      const double tolerance = 1.0e-12) {
  EXPECT_NEAR(actual.X(), expected.X(), tolerance);
  EXPECT_NEAR(actual.Y(), expected.Y(), tolerance);
  EXPECT_NEAR(actual.Z(), expected.Z(), tolerance);
}

void ExpectProperOrthogonal(const Matrix& rotation) {
  const Matrix orthogonality = Transpose(rotation).Multiply(rotation);
  ASSERT_EQ(orthogonality.Rows(), kRotationDimension);
  ASSERT_EQ(orthogonality.Columns(), kRotationDimension);
  for (std::size_t row = 0U; row < kRotationDimension; ++row) {
    for (std::size_t column = 0U; column < kRotationDimension; ++column) {
      EXPECT_NEAR(orthogonality(row, column), row == column ? 1.0 : 0.0,
                  1.0e-12);
    }
  }
  EXPECT_NEAR(Determinant(rotation), 1.0, 1.0e-12);
}

TEST(FiniteRotation, ZeroCoordinatePreservesInitialFrameAndDirector) {
  const Matrix initial_frame = InitialFrame();

  const Result<FiniteRotationState> result =
      MakeFiniteRotationState(Vector3{}, initial_frame);

  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().phi, Vector3{});
  ExpectMatrixNear(result.Value().q, initial_frame);
  ExpectVectorNear(result.Value().director,
                   Multiply(initial_frame, Vector3{0.0, 0.0, 1.0}));
}

TEST(FiniteRotation, FiniteCoordinateBuildsProperTriadAndDirector) {
  const Matrix initial_frame = InitialFrame();
  const Vector3 phi{0.3, -0.4, 0.5};

  const Result<FiniteRotationState> result =
      MakeFiniteRotationState(phi, initial_frame);

  ASSERT_TRUE(result.HasValue());
  const Matrix expected_q = ExpSkew(phi).Multiply(initial_frame);
  ExpectMatrixNear(result.Value().q, expected_q);
  ExpectProperOrthogonal(result.Value().q);
  ExpectVectorNear(result.Value().director,
                   Multiply(result.Value().q, Vector3{0.0, 0.0, 1.0}));
}

TEST(FiniteRotation, GlobalIncrementRebuildsFixedReferenceOrientation) {
  const Matrix initial_frame = InitialFrame();
  const Vector3 phi{0.2, -0.1, 0.3};
  const Vector3 delta_phi{-0.05, 0.15, 0.1};
  const Result<FiniteRotationState> committed =
      MakeFiniteRotationState(phi, initial_frame);
  ASSERT_TRUE(committed.HasValue());

  const Result<FiniteRotationState> updated =
      ApplyGlobalRotationIncrement(committed.Value(), delta_phi, initial_frame);

  ASSERT_TRUE(updated.HasValue());
  const Vector3 expected_phi = phi + delta_phi;
  EXPECT_EQ(updated.Value().phi, expected_phi);
  ExpectMatrixNear(updated.Value().q,
                   ExpSkew(expected_phi).Multiply(initial_frame));
  ExpectVectorNear(updated.Value().director,
                   Multiply(updated.Value().q, Vector3{0.0, 0.0, 1.0}));
  ExpectMatrixNear(committed.Value().q, ExpSkew(phi).Multiply(initial_frame));
}

TEST(FiniteRotation, RejectsCoordinatesInsidePrincipalBranchMargin) {
  const Matrix initial_frame = InitialFrame();
  const double pi = std::acos(-1.0);
  const double margin = std::sqrt(std::numeric_limits<double>::epsilon());

  EXPECT_TRUE(MakeFiniteRotationState(Vector3{pi - 2.0 * margin, 0.0, 0.0},
                                      initial_frame)
                  .HasValue());
  const Result<FiniteRotationState> ambiguous =
      MakeFiniteRotationState(Vector3{pi - margin, 0.0, 0.0}, initial_frame);
  ASSERT_FALSE(ambiguous.HasValue());
  ASSERT_FALSE(ambiguous.GetStatus().IsOk());
  ASSERT_EQ(ambiguous.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(ambiguous.GetStatus().Diagnostics()[0U].code,
            "invalid-accepted-state");
}

TEST(FiniteRotation, SpatialMomentSolvePreservesKnownLoadAndVirtualWork) {
  const Vector3 moment{2.0, -3.0, 4.0};
  const Vector3 variation{0.3, -0.5, 0.7};
  const double pi = std::acos(-1.0);
  for (const double angle : {0.0, 1e-9, 0.7, pi - 1e-6}) {
    const double a = angle == 0.0 ? 1.0 : std::sin(angle) / angle;
    const double b = angle < 1e-8 ? angle / 2.0 :
        (1.0 - std::cos(angle)) / angle;
    // Independent z-axis J^T*m; the axial component is unchanged.
    const Vector3 generalized{a * moment.X() + b * moment.Y(),
                              -b * moment.X() + a * moment.Y(), moment.Z()};
    const auto solved = SpatialMomentSO3(Vector3{0, 0, angle}, generalized);
    ASSERT_TRUE(solved.HasValue());
    ExpectVectorNear(solved.Value(), moment, 1e-11);
    const Vector3 spatial_variation{a * variation.X() - b * variation.Y(),
                                    b * variation.X() + a * variation.Y(),
                                    variation.Z()};
    EXPECT_NEAR(generalized.Dot(variation),
                solved.Value().Dot(spatial_variation), 1e-11);
  }
}

TEST(FiniteRotation,
     JacobiansAreFiniteInverseTransposeAndDerivativeConsistent) {
  const double pi = std::acos(-1.0);
  const std::array<Vector3, 3> rotations{
      Vector3{1.0e-9, -2.0e-9, 3.0e-9}, Vector3{0.4, -0.2, 0.3},
      Vector3{0.0, pi - 1.0e-3, 0.0}};
  const Vector3 direction{0.2, -0.3, 0.5};
  const double perturbation = 1.0e-7;

  for (const Vector3& phi : rotations) {
    const Matrix left_jacobian = LeftJacobianSO3(phi);
    EXPECT_TRUE(IsFinite(left_jacobian));

    // Preserve the original complete inverse-action check through the solve API.
    for (std::size_t column = 0; column < kRotationDimension; ++column) {
      const Vector3 basis{column == 0 ? 1.0 : 0.0,
                          column == 1 ? 1.0 : 0.0,
                          column == 2 ? 1.0 : 0.0};
      const auto solved = SpatialMomentSO3(phi, basis);
      ASSERT_TRUE(solved.HasValue());
      EXPECT_TRUE(std::isfinite(solved.Value().X()));
      EXPECT_TRUE(std::isfinite(solved.Value().Y()));
      EXPECT_TRUE(std::isfinite(solved.Value().Z()));
      const auto identity_column = Multiply(Transpose(left_jacobian), solved.Value());
      for (std::size_t row = 0; row < kRotationDimension; ++row)
        EXPECT_NEAR(identity_column[row], row == column ? 1.0 : 0.0, 1.0e-11);
    }

    const Vector3 load{0.7, -0.2, 0.4};
    const auto spatial = SpatialMomentSO3(phi, load);
    ASSERT_TRUE(spatial.HasValue());
    ExpectVectorNear(Multiply(Transpose(left_jacobian), spatial.Value()),
                     load, 1e-11);

    const Vector3 director = Multiply(ExpSkew(phi), Vector3{0.0, 0.0, 1.0});
    const Vector3 positive =
        Multiply(ExpSkew(phi + perturbation * direction),
                 Vector3{0.0, 0.0, 1.0});
    const Vector3 negative =
        Multiply(ExpSkew(phi - perturbation * direction),
                 Vector3{0.0, 0.0, 1.0});
    const Vector3 observed = (positive - negative) / (2.0 * perturbation);
    const Vector3 expected =
        Multiply(left_jacobian, direction).Cross(director);
    ExpectVectorNear(observed, expected, 1.0e-8);
  }
}

TEST(FiniteRotation, RejectsNonfiniteCoordinatesAndInvalidInitialFrames) {
  const Matrix initial_frame = InitialFrame();
  const double infinity = std::numeric_limits<double>::infinity();
  const Result<FiniteRotationState> nonfinite =
      MakeFiniteRotationState(Vector3{infinity, 0.0, 0.0}, initial_frame);
  ASSERT_FALSE(nonfinite.HasValue());
  ASSERT_EQ(nonfinite.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(nonfinite.GetStatus().Diagnostics()[0U].code,
            "invalid-accepted-state");

  Matrix invalid_frame = initial_frame;
  invalid_frame(0U, 0U) = 2.0;
  EXPECT_FALSE(MakeFiniteRotationState(Vector3{}, invalid_frame).HasValue());

  const Result<FiniteRotationState> committed =
      MakeFiniteRotationState(Vector3{}, initial_frame);
  ASSERT_TRUE(committed.HasValue());
  const Result<FiniteRotationState> invalid_increment =
      ApplyGlobalRotationIncrement(committed.Value(),
                                   Vector3{0.0, infinity, 0.0}, initial_frame);
  ASSERT_FALSE(invalid_increment.HasValue());
  ASSERT_EQ(invalid_increment.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(invalid_increment.GetStatus().Diagnostics()[0U].code,
            "invalid-accepted-state");
}

}  // namespace
}  // namespace fesa