#include "fesa/elements/finite_rotation.h"

#include <cmath>
#include <cstddef>
#include <limits>
#include <string>
#include <utility>

namespace fesa {
namespace {

constexpr std::size_t kRotationDimension = 3U;
constexpr double kRotationStateTolerance = 1.0e-12;
constexpr double kSmallAngleThreshold = 1.0e-4;

/// @brief Returns the SO(3) skew-symmetric matrix for a global vector.
Matrix Skew(const Vector3& vector) {
  Matrix result{kRotationDimension, kRotationDimension};
  result(0U, 1U) = -vector.Z();
  result(0U, 2U) = vector.Y();
  result(1U, 0U) = vector.Z();
  result(1U, 2U) = -vector.X();
  result(2U, 0U) = -vector.Y();
  result(2U, 1U) = vector.X();
  return result;
}

/// @brief Returns the 3-by-3 identity matrix.
Matrix Identity() {
  Matrix result{kRotationDimension, kRotationDimension};
  for (std::size_t index = 0U; index < kRotationDimension; ++index) {
    result(index, index) = 1.0;
  }
  return result;
}

/// @brief Builds I plus linear and quadratic skew terms in stable row-major
/// order.
Matrix Polynomial(const Matrix& skew, const Matrix& skew_squared,
                  const double linear_coefficient,
                  const double quadratic_coefficient) {
  Matrix result = Identity();
  for (std::size_t row = 0U; row < kRotationDimension; ++row) {
    for (std::size_t column = 0U; column < kRotationDimension; ++column) {
      result(row, column) += linear_coefficient * skew(row, column) +
                             quadratic_coefficient *
                                 skew_squared(row, column);
    }
  }
  return result;
}

/// @brief Computes the determinant of one 3-by-3 matrix.
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));
}

/// @brief Checks that all matrix entries are finite.
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;
}

/// @brief Checks the exact 3-by-3 proper-orthogonal frame contract.
bool IsProperOrthogonal(const Matrix& matrix) {
  if (matrix.Rows() != kRotationDimension ||
      matrix.Columns() != kRotationDimension || !IsFinite(matrix)) {
    return false;
  }

  for (std::size_t row = 0U; row < kRotationDimension; ++row) {
    for (std::size_t column = 0U; column < kRotationDimension; ++column) {
      double value = 0.0;
      for (std::size_t index = 0U; index < kRotationDimension; ++index) {
        value += matrix(index, row) * matrix(index, column);
      }
      if (!std::isfinite(value) ||
          std::abs(value - (row == column ? 1.0 : 0.0)) >
              kRotationStateTolerance) {
        return false;
      }
    }
  }

  const double determinant = Determinant(matrix);
  return std::isfinite(determinant) &&
         std::abs(determinant - 1.0) <= kRotationStateTolerance;
}

/// @brief Checks that the coordinate is finite and away from the pi branch.
bool IsPrincipalCoordinate(const Vector3& phi) {
  if (!phi.IsFinite()) {
    return false;
  }

  const double angle = phi.Norm();
  const double pi = std::acos(-1.0);
  const double margin = std::sqrt(std::numeric_limits<double>::epsilon());
  return std::isfinite(angle) && pi - angle > margin;
}

/// @brief Multiplies a 3-by-3 frame by a global vector.
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()};
}

/// @brief Checks componentwise agreement without scale-dependent repair.
bool IsNear(const Matrix& actual, const Matrix& expected) {
  if (actual.Rows() != expected.Rows() ||
      actual.Columns() != expected.Columns()) {
    return false;
  }
  for (std::size_t row = 0U; row < actual.Rows(); ++row) {
    for (std::size_t column = 0U; column < actual.Columns(); ++column) {
      if (!std::isfinite(actual(row, column)) ||
          !std::isfinite(expected(row, column)) ||
          std::abs(actual(row, column) - expected(row, column)) >
              kRotationStateTolerance) {
        return false;
      }
    }
  }
  return true;
}

/// @brief Checks componentwise agreement for a derived director.
bool IsNear(const Vector3& actual, const Vector3& expected) {
  return actual.IsFinite() && expected.IsFinite() &&
         std::abs(actual.X() - expected.X()) <= kRotationStateTolerance &&
         std::abs(actual.Y() - expected.Y()) <= kRotationStateTolerance &&
         std::abs(actual.Z() - expected.Z()) <= kRotationStateTolerance;
}

/// @brief Returns the stable failure used when a nonlinear state is invalid.
Result<FiniteRotationState> InvalidAcceptedState(const std::string& message) {
  return Result<FiniteRotationState>::Failure(Status::Failure(
      FailureCategory::kSolver,
      {{Severity::kError, "invalid-accepted-state", {{}, 0U},
        "FINITE_ROTATION", "", message}}));
}

/// @brief Validates that an existing state remains derived from its coordinate.
bool IsConsistentState(const FiniteRotationState& state,
                       const Matrix& initial_frame) {
  if (!IsPrincipalCoordinate(state.phi) || !IsProperOrthogonal(state.q) ||
      !state.director.IsFinite()) {
    return false;
  }

  const Matrix expected_q = ExpSkew(state.phi).Multiply(initial_frame);
  return IsNear(state.q, expected_q) &&
         IsNear(state.director, Multiply(state.q, Vector3{0.0, 0.0, 1.0}));
}

}  // namespace

Matrix ExpSkew(const Vector3& phi) {
  const Matrix skew = Skew(phi);
  const Matrix skew_squared = skew.Multiply(skew);
  const double theta = phi.Norm();
  const double theta_squared = theta * theta;

  double linear_coefficient = 1.0;
  double quadratic_coefficient = 0.5;
  if (theta >= kSmallAngleThreshold) {
    linear_coefficient = std::sin(theta) / theta;
    quadratic_coefficient = (1.0 - std::cos(theta)) / theta_squared;
  } else {
    const double theta_fourth = theta_squared * theta_squared;
    const double theta_sixth = theta_fourth * theta_squared;
    linear_coefficient =
        1.0 - theta_squared / 6.0 + theta_fourth / 120.0 -
        theta_sixth / 5040.0;
    quadratic_coefficient =
        0.5 - theta_squared / 24.0 + theta_fourth / 720.0 -
        theta_sixth / 40320.0;
  }
  return Polynomial(skew, skew_squared, linear_coefficient,
                    quadratic_coefficient);
}

Matrix LeftJacobianSO3(const Vector3& phi) {
  const Matrix skew = Skew(phi);
  const Matrix skew_squared = skew.Multiply(skew);
  const double theta = phi.Norm();
  const double theta_squared = theta * theta;

  double linear_coefficient = 0.5;
  double quadratic_coefficient = 1.0 / 6.0;
  if (theta >= kSmallAngleThreshold) {
    linear_coefficient = (1.0 - std::cos(theta)) / theta_squared;
    quadratic_coefficient = (theta - std::sin(theta)) /
                            (theta_squared * theta);
  } else {
    const double theta_fourth = theta_squared * theta_squared;
    const double theta_sixth = theta_fourth * theta_squared;
    linear_coefficient =
        0.5 - theta_squared / 24.0 + theta_fourth / 720.0 -
        theta_sixth / 40320.0;
    quadratic_coefficient =
        1.0 / 6.0 - theta_squared / 120.0 + theta_fourth / 5040.0 -
        theta_sixth / 362880.0;
  }
  return Polynomial(skew, skew_squared, linear_coefficient,
                    quadratic_coefficient);
}

Matrix InverseTransposeLeftJacobianSO3(const Vector3& phi) {
  const Matrix skew = Skew(phi);
  const Matrix skew_squared = skew.Multiply(skew);
  const double theta = phi.Norm();
  const double theta_squared = theta * theta;

  double quadratic_coefficient = 1.0 / 12.0;
  if (theta >= kSmallAngleThreshold) {
    const double half_theta = 0.5 * theta;
    quadratic_coefficient =
        (1.0 - half_theta / std::tan(half_theta)) / theta_squared;
  } else {
    const double theta_fourth = theta_squared * theta_squared;
    const double theta_sixth = theta_fourth * theta_squared;
    quadratic_coefficient =
        1.0 / 12.0 + theta_squared / 720.0 + theta_fourth / 30240.0 +
        theta_sixth / 1209600.0;
  }
  return Polynomial(skew, skew_squared, 0.5, quadratic_coefficient);
}

Result<FiniteRotationState> MakeFiniteRotationState(
    const Vector3& phi, const Matrix& initial_frame) {
  if (!IsPrincipalCoordinate(phi)) {
    return InvalidAcceptedState(
        "Finite rotation coordinates must be finite and remain inside the "
        "principal branch margin.");
  }
  if (!IsProperOrthogonal(initial_frame)) {
    return InvalidAcceptedState(
        "The initial finite-rotation frame must be a finite proper "
        "orthogonal triad.");
  }

  Matrix q = ExpSkew(phi).Multiply(initial_frame);
  const Vector3 director = Multiply(q, Vector3{0.0, 0.0, 1.0});
  if (!IsProperOrthogonal(q) || !director.IsFinite()) {
    return InvalidAcceptedState(
        "Finite rotation state construction produced an invalid triad.");
  }
  return Result<FiniteRotationState>::Success(
      FiniteRotationState{phi, std::move(q), director});
}

Result<FiniteRotationState> ApplyGlobalRotationIncrement(
    const FiniteRotationState& committed, const Vector3& delta_phi,
    const Matrix& initial_frame) {
  if (!IsProperOrthogonal(initial_frame) ||
      !IsConsistentState(committed, initial_frame) ||
      !delta_phi.IsFinite()) {
    return InvalidAcceptedState(
        "The accepted finite rotation state and its correction must remain "
        "finite and consistent with the initial frame.");
  }

  return MakeFiniteRotationState(committed.phi + delta_phi, initial_frame);
}

}  // namespace fesa