// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

#pragma once

#include <utility>

#include "wpi/math/geometry/Rotation2d.hpp"
#include "wpi/math/geometry/Translation2d.hpp"
#include "wpi/util/SymbolExports.hpp"

namespace wpi::math {

class Pose2d;
struct Twist2d;

/**
 * Represents a transformation for a Pose2d in the pose's frame.
 */
class WPILIB_DLLEXPORT Transform2d final {
 public:
  /**
   * Constructs the transform that maps the initial pose to the final pose.
   *
   * @param initial The initial pose for the transformation.
   * @param final The final pose for the transformation.
   */
  constexpr Transform2d(const Pose2d& initial, const Pose2d& final);

  /**
   * Constructs a transform with the given translation and rotation components.
   *
   * @param translation Translational component of the transform.
   * @param rotation Rotational component of the transform.
   */
  constexpr Transform2d(Translation2d translation, Rotation2d rotation)
      : m_translation{std::move(translation)},
        m_rotation{std::move(rotation)} {}

  /**
   * Constructs a transform with x and y translations instead of a separate
   * Translation2d.
   *
   * @param x The x component of the translational component of the transform.
   * @param y The y component of the translational component of the transform.
   * @param rotation The rotational component of the transform.
   */
  constexpr Transform2d(wpi::units::meter_t x, wpi::units::meter_t y,
                        Rotation2d rotation)
      : m_translation{x, y}, m_rotation{std::move(rotation)} {}

  /**
   * Constructs a pose with the specified affine transformation matrix.
   *
   * @param matrix The affine transformation matrix.
   * @throws std::domain_error if the affine transformation matrix is invalid.
   */
  constexpr explicit Transform2d(const Eigen::Matrix3d& matrix)
      : m_translation{Eigen::Vector2d{{matrix(0, 2)}, {matrix(1, 2)}}},
        m_rotation{Eigen::Matrix2d{{matrix(0, 0), matrix(0, 1)},
                                   {matrix(1, 0), matrix(1, 1)}}} {
    if (matrix(2, 0) != 0.0 || matrix(2, 1) != 0.0 || matrix(2, 2) != 1.0) {
      throw std::domain_error("Affine transformation matrix is invalid");
    }
  }

  /**
   * Constructs the identity transform -- maps an initial pose to itself.
   */
  constexpr Transform2d() = default;

  /**
   * Returns the translation component of the transformation.
   *
   * @return Reference to the translational component of the transform.
   */
  constexpr const Translation2d& Translation() const { return m_translation; }

  /**
   * Returns the X component of the transformation's translation.
   *
   * @return The x component of the transformation's translation.
   */
  constexpr wpi::units::meter_t X() const { return m_translation.X(); }

  /**
   * Returns the Y component of the transformation's translation.
   *
   * @return The y component of the transformation's translation.
   */
  constexpr wpi::units::meter_t Y() const { return m_translation.Y(); }

  /**
   * Returns an affine transformation matrix representation of this
   * transformation.
   */
  constexpr Eigen::Matrix3d ToMatrix() const {
    auto vec = m_translation.ToVector();
    auto mat = m_rotation.ToMatrix();
    return Eigen::Matrix3d{{mat(0, 0), mat(0, 1), vec(0)},
                           {mat(1, 0), mat(1, 1), vec(1)},
                           {0.0, 0.0, 1.0}};
  }

  /**
   * Returns the rotational component of the transformation.
   *
   * @return Reference to the rotational component of the transform.
   */
  constexpr const Rotation2d& Rotation() const { return m_rotation; }

  /**
   * Returns a Twist2d of the current transform (pose delta). If b is the output
   * of {@code a.Log()}, then {@code b.Exp()} would yield a.
   *
   * @return The twist that maps the current transform.
   */
  constexpr Twist2d Log() const;

  /**
   * Invert the transformation. This is useful for undoing a transformation.
   *
   * @return The inverted transformation.
   */
  constexpr Transform2d Inverse() const {
    // We are rotating the difference between the translations
    // using a clockwise rotation matrix. This transforms the global
    // delta into a local delta (relative to the initial pose).
    return Transform2d{(-Translation()).RotateBy(-Rotation()), -Rotation()};
  }

  /**
   * Multiplies the transform by the scalar.
   *
   * @param scalar The scalar.
   * @return The scaled Transform2d.
   */
  constexpr Transform2d operator*(double scalar) const {
    return Transform2d(m_translation * scalar, m_rotation * scalar);
  }

  /**
   * Divides the transform by the scalar.
   *
   * @param scalar The scalar.
   * @return The scaled Transform2d.
   */
  constexpr Transform2d operator/(double scalar) const {
    return *this * (1.0 / scalar);
  }

  /**
   * Composes two transformations. The second transform is applied relative to
   * the orientation of the first.
   *
   * @param other The transform to compose with this one.
   * @return The composition of the two transformations.
   */
  constexpr Transform2d operator+(const Transform2d& other) const;

  /**
   * Checks equality between this Transform2d and another object.
   */
  constexpr bool operator==(const Transform2d&) const = default;

 private:
  Translation2d m_translation;
  Rotation2d m_rotation;
};

}  // namespace wpi::math

#include "wpi/math/geometry/Pose2d.hpp"
#include "wpi/math/geometry/Twist2d.hpp"

namespace wpi::math {

constexpr Transform2d::Transform2d(const Pose2d& initial, const Pose2d& final) {
  // To transform the global translation delta to be relative to the initial
  // pose, rotate by the inverse of the initial pose's orientation.
  m_translation = (final.Translation() - initial.Translation())
                      .RotateBy(-initial.Rotation());

  m_rotation = final.Rotation().RelativeTo(initial.Rotation());
}

constexpr Transform2d Transform2d::operator+(const Transform2d& other) const {
  return Transform2d{Pose2d{}, Pose2d{}.TransformBy(*this).TransformBy(other)};
}

constexpr Twist2d Transform2d::Log() const {
  const auto dtheta = m_rotation.Radians().value();
  const auto halfDtheta = dtheta / 2.0;

  const auto cosMinusOne = m_rotation.Cos() - 1;

  double halfThetaByTanOfHalfDtheta;

  if (gcem::abs(cosMinusOne) < 1E-9) {
    halfThetaByTanOfHalfDtheta = 1.0 - 1.0 / 12.0 * dtheta * dtheta;
  } else {
    halfThetaByTanOfHalfDtheta = -(halfDtheta * m_rotation.Sin()) / cosMinusOne;
  }

  const Translation2d translationPart =
      m_translation.RotateBy({halfThetaByTanOfHalfDtheta, -halfDtheta}) *
      gcem::hypot(halfThetaByTanOfHalfDtheta, halfDtheta);

  return {translationPart.X(), translationPart.Y(),
          wpi::units::radian_t{dtheta}};
}

}  // namespace wpi::math

#include "wpi/math/geometry/proto/Transform2dProto.hpp"
#include "wpi/math/geometry/struct/Transform2dStruct.hpp"
