//! # Affine3
//!
//! A 3D affine transform, with a linear part (rotation, scale, shear) plus a translation.
//!
//! # Examples
//! ```rust
//! # use roofy::{Affine3, Vec3};
//! let t: Affine3 = Affine3::from_translation(Vec3::new(1.0, 2.0, 3.0));
//! let inv: Affine3 = t.inverse();
//!
//! let p: Vec3 = Vec3::new(4.0, 5.0, 6.0);
//! let round_tripped: Vec3 = inv.transform_point(t.transform_point(p));
//!
//! assert!(round_tripped.approx_eq(p));
//! ```

use std::{
    fmt::Display,
    ops::{Mul, MulAssign},
};

use crate::f32::{Mat3, Mat4, Quat, Vec3, Vec4};

/// A 3D affine transform, with a linear part (rotation, scale, shear) plus a translation.
#[derive(Clone, Copy, Debug, PartialEq)]
#[repr(C)]
pub struct Affine3 {
    /// The linear part of the transform (rotation, scale, shear).
    pub matrix: Mat3,
    /// The translation part of the transform.
    pub translation: Vec3,
}

impl Affine3 {
    /// Threshold for treating a value as zero.
    pub const EPSILON: f32 = 1e-6;
    /// [`EPSILON`](Self::EPSILON) squared.
    pub const EPSILON_SQ: f32 = Self::EPSILON * Self::EPSILON;

    /// An Affine3 with matrix and translation all `NaN`.
    pub const NAN: Self = Self::from_mat3_translation(Mat3::NAN, Vec3::NAN);

    /// An Affine3 with matrix and translation all `0.0`.
    pub const ZERO: Self = Self::from_mat3_translation(Mat3::ZERO, Vec3::ZERO);

    /// The identity transform.
    pub const IDENTITY: Self = Self::from_mat3_translation(Mat3::IDENTITY, Vec3::ZERO);

    /// Creates an `Affine3` from a `Mat3`.
    #[inline(always)]
    #[must_use]
    pub const fn from_mat3(m: Mat3) -> Self {
        Self::from_mat3_translation(m, Vec3::ZERO)
    }

    /// Creates an `Affine3` from a `matrix` and `translation`.
    #[inline(always)]
    #[must_use]
    pub const fn from_mat3_translation(matrix: Mat3, translation: Vec3) -> Self {
        Self {
            matrix,
            translation,
        }
    }

    /// Creates an `Affine3` from 3 columns and `translation`.
    #[inline]
    #[must_use]
    pub const fn from_cols(col0: Vec3, col1: Vec3, col2: Vec3, translation: Vec3) -> Self {
        Self::from_mat3_translation(Mat3::from_cols(col0, col1, col2), translation)
    }

    /// Creates an `Affine3` from a `[f32; 12]`, in column-major order:
    /// `[col0, col1, col2, translation]` with each column as `[x, y, z]`.
    #[inline]
    #[must_use]
    pub const fn from_array(a: [f32; 12]) -> Self {
        Self::from_mat3_translation(
            Mat3::from_cols(
                Vec3::new(a[0], a[1], a[2]),
                Vec3::new(a[3], a[4], a[5]),
                Vec3::new(a[6], a[7], a[8]),
            ),
            Vec3::new(a[9], a[10], a[11]),
        )
    }

    /// Returns a copy of the components as `[col0, col1, col2, translation]`,
    /// with each column as `[x, y, z]`, in column-major order.
    #[inline]
    #[must_use]
    pub const fn to_array(self) -> [f32; 12] {
        let m: Mat3 = self.matrix;
        let t: Vec3 = self.translation;

        [
            m.col0.x, m.col0.y, m.col0.z, m.col1.x, m.col1.y, m.col1.z, m.col2.x, m.col2.y,
            m.col2.z, t.x, t.y, t.z,
        ]
    }

    /// Creates an `Affine3` from first 12 values of `slice`, in column-major order.
    ///
    /// # Panics
    /// Panics if `slice` is less than 12 elements long.
    #[inline]
    #[must_use]
    pub const fn from_slice(slice: &[f32]) -> Self {
        Self::from_mat3_translation(
            Mat3::from_cols(
                Vec3::new(slice[0], slice[1], slice[2]),
                Vec3::new(slice[3], slice[4], slice[5]),
                Vec3::new(slice[6], slice[7], slice[8]),
            ),
            Vec3::new(slice[9], slice[10], slice[11]),
        )
    }

    /// Writes the columns of `self` to the first 12 elements of `slice`.
    ///
    /// # Panics
    /// Panics if `slice` is less than 12 elements long.
    #[inline]
    pub const fn write_to_slice(self, slice: &mut [f32]) {
        slice[0] = self.matrix.col0.x;
        slice[1] = self.matrix.col0.y;
        slice[2] = self.matrix.col0.z;
        slice[3] = self.matrix.col1.x;
        slice[4] = self.matrix.col1.y;
        slice[5] = self.matrix.col1.z;
        slice[6] = self.matrix.col2.x;
        slice[7] = self.matrix.col2.y;
        slice[8] = self.matrix.col2.z;
        slice[9] = self.translation.x;
        slice[10] = self.translation.y;
        slice[11] = self.translation.z;
    }

    /// Creates an `Affine3` that translates by `translation`.
    #[inline]
    #[must_use]
    pub const fn from_translation(translation: Vec3) -> Self {
        Self::from_mat3_translation(Mat3::IDENTITY, translation)
    }

    /// Creates an `Affine3` that scales by `x`, `y` and `z` components of `scale` along each axis.
    #[inline]
    #[must_use]
    pub const fn from_scale(scale: Vec3) -> Self {
        Self::from_mat3_translation(Mat3::from_diagonal(scale), Vec3::ZERO)
    }

    /// Creates an `Affine3` that rotates around the X axis by `angle` in radians.
    #[inline]
    #[must_use]
    pub fn from_rotation_x(angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_rotation_x(angle), Vec3::ZERO)
    }

    /// Creates an `Affine3` that rotates around the X axis by `angle` in degrees.
    #[inline]
    #[must_use]
    pub fn from_rotation_x_degrees(angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_rotation_x_degrees(angle), Vec3::ZERO)
    }

    /// Creates an `Affine3` that rotates around the Y axis by `angle` in radians.
    #[inline]
    #[must_use]
    pub fn from_rotation_y(angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_rotation_y(angle), Vec3::ZERO)
    }

    /// Creates an `Affine3` that rotates around the Y axis by `angle` in degrees.
    #[inline]
    #[must_use]
    pub fn from_rotation_y_degrees(angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_rotation_y_degrees(angle), Vec3::ZERO)
    }

    /// Creates an `Affine3 that rotates around the Z axis by `angle` in radians.
    #[inline]
    #[must_use]
    pub fn from_rotation_z(angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_rotation_z(angle), Vec3::ZERO)
    }

    /// Creates an `Affine3 that rotates around the Z axis by `angle` in degrees.
    #[inline]
    #[must_use]
    pub fn from_rotation_z_degrees(angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_rotation_z_degrees(angle), Vec3::ZERO)
    }

    /// Creates an `Affine3` from a `translation`, a `rotation` as a `Quat` and a `scale`.
    ///
    /// Applied in order: scale -> rotate -> translate.
    #[inline]
    #[must_use]
    pub fn from_trs(translation: Vec3, rotation: Quat, scale: Vec3) -> Self {
        Self::from_mat3_translation(Mat3::from_rs(rotation, scale), translation)
    }

    /// Creates an `Affine3` from a `translation`, a `rotation` as euler angles in a `Vec3`
    /// in radians and a `scale`.
    ///
    /// Applied in order: scale -> rotate -> translate.
    #[inline]
    #[must_use]
    pub fn from_trs_euler(translation: Vec3, rotation: Vec3, scale: Vec3) -> Self {
        Self::from_mat3_translation(Mat3::from_rs_euler(rotation, scale), translation)
    }

    /// Creates an `Affine3` from a `translation`, a `rotation` as euler angles in a `Vec3`
    /// in degrees and a `scale`.
    ///
    /// Applied in order: scale -> rotate -> translate.
    #[inline]
    #[must_use]
    pub fn from_trs_euler_degrees(translation: Vec3, rotation: Vec3, scale: Vec3) -> Self {
        Self::from_mat3_translation(Mat3::from_rs_euler_degrees(rotation, scale), translation)
    }

    /// Creates an `Affine3` that rotates by `Quat`.
    #[inline]
    #[must_use]
    pub const fn from_quat(quat: Quat) -> Self {
        Self::from_mat3_translation(Mat3::from_quat(quat), Vec3::ZERO)
    }

    /// Creates an `Affine3` that rotates around `axis` by `angle` in radians.
    ///
    /// # Preconditions
    /// Expects `axis` to be normalized.
    #[inline]
    #[must_use]
    pub fn from_axis_angle(axis: Vec3, angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_axis_angle(axis, angle), Vec3::ZERO)
    }

    /// Creates an `Affine3` that rotates around `axis` by `angle` in degrees.
    ///
    /// # Preconditions
    /// Expects `axis` to be normalized.
    #[inline]
    #[must_use]
    pub fn from_axis_angle_degrees(axis: Vec3, angle: f32) -> Self {
        Self::from_mat3_translation(Mat3::from_axis_angle_degrees(axis, angle), Vec3::ZERO)
    }

    /// Creates an `Affine3` from a `Mat4`, discarding its projective row.
    #[inline]
    #[must_use]
    pub const fn from_mat4(m: Mat4) -> Self {
        Self::from_mat3_translation(
            Mat3::from_cols(
                Vec3::new(m.col0.x, m.col0.y, m.col0.z),
                Vec3::new(m.col1.x, m.col1.y, m.col1.z),
                Vec3::new(m.col2.x, m.col2.y, m.col2.z),
            ),
            m.translation(),
        )
    }

    /// Returns the equivalent `Mat4` of `self`.
    #[inline]
    #[must_use]
    pub const fn to_mat4(self) -> Mat4 {
        let m: Mat3 = self.matrix;
        let t: Vec3 = self.translation;

        Mat4::from_cols(
            Vec4::new(m.col0.x, m.col0.y, m.col0.z, 0.0),
            Vec4::new(m.col1.x, m.col1.y, m.col1.z, 0.0),
            Vec4::new(m.col2.x, m.col2.y, m.col2.z, 0.0),
            Vec4::new(t.x, t.y, t.z, 1.0),
        )
    }

    /// Creates an `Affine3` that is a view transform
    /// positioned at `eye`, looking to `dir`, with `up` as the up direction.
    ///
    /// # Convention
    /// Uses right-handed coordinate system.
    ///
    /// # Preconditions
    /// Expects `dir` to be normalized.
    /// Expects the length of `dir.cross(up)` to not be approximately `0.0` and finite.
    #[inline]
    #[must_use]
    pub fn look_to_rh(eye: Vec3, dir: Vec3, up: Vec3) -> Self {
        let f: Vec3 = dir;
        let r: Vec3 = f.cross(up).normalized();
        let u: Vec3 = r.cross(f);

        Self::from_mat3_translation(
            Mat3::from_cols(
                Vec3::new(r.x, u.x, -f.x),
                Vec3::new(r.y, u.y, -f.y),
                Vec3::new(r.z, u.z, -f.z),
            ),
            Vec3::new(-r.dot(eye), -u.dot(eye), f.dot(eye)),
        )
    }

    /// Creates an `Affine3` that is a view transform
    /// positioned at `eye`, looking to `dir`, with `up` as the up direction.
    ///
    /// # Convention
    /// Uses left-handed coordinate system.
    ///
    /// # Preconditions
    /// Expects `dir` to be normalized.
    /// Expects the length of `dir.cross(up)` to not be approximately `0.0` and finite.
    #[inline]
    #[must_use]
    pub fn look_to_lh(eye: Vec3, dir: Vec3, up: Vec3) -> Self {
        Self::look_to_rh(eye, -dir, up)
    }

    /// Creates an `Affine3` that is a view transform
    /// positioned at `eye`, looking at `center`, with `up` as the up direction.
    ///
    /// # Convention
    /// Uses right-handed coordinate system.
    ///
    /// # Preconditions
    /// Expects the length of `(center - eye)` to not be approximately `0.0` and finite.
    /// Expects the length of `(center - eye).normalized().cross(up)` to not be approximately `0.0` and finite.
    #[inline]
    #[must_use]
    pub fn look_at_rh(eye: Vec3, center: Vec3, up: Vec3) -> Self {
        Self::look_to_rh(eye, (center - eye).normalized(), up)
    }

    /// Creates an `Affine3` that is a view transform
    /// positioned at `eye`, looking at `center`, with `up` as the up direction.
    ///
    /// # Convention
    /// Uses left-handed coordinate system.
    ///
    /// # Preconditions
    /// Expects the length of `(center - eye)` to not be approximately `0.0` and finite.
    /// Expects the length of `(center - eye).normalized().cross(up)` to not be approximately `0.0` and finite.
    #[inline]
    #[must_use]
    pub fn look_at_lh(eye: Vec3, center: Vec3, up: Vec3) -> Self {
        Self::look_to_lh(eye, (center - eye).normalized(), up)
    }

    /// Returns the composition of `self` and `rhs`, applying `rhs` first.
    #[inline]
    #[must_use]
    pub const fn mul(self, rhs: Self) -> Self {
        Self::from_mat3_translation(
            self.matrix.mul(rhs.matrix),
            self.matrix.transform(rhs.translation).add(self.translation),
        )
    }

    /// Sets `self` to the composition of `self` and `rhs`, applying `rhs` first.
    #[inline]
    pub const fn mul_assign(&mut self, rhs: Self) {
        *self = self.mul(rhs);
    }

    /// Returns `rhs` transformed by `self`, applying rotation, scale, shear and translation.
    ///
    /// Equivalent to [`transform_point`](Self::transform_point).
    #[inline]
    #[must_use]
    pub const fn mul_vec3(self, rhs: Vec3) -> Vec3 {
        self.transform_point(rhs)
    }

    /// Returns the inverse of `self` or `None` if `self` is not invertible.
    ///
    /// # Edge cases
    /// If the linear part of `self` is not invertible, returns `None`.
    #[inline]
    #[must_use]
    pub const fn try_inverse(self) -> Option<Self> {
        let Some(inv_matrix) = self.matrix.try_inverse() else {
            return None;
        };
        let inv_translation: Vec3 = inv_matrix.mul_vec3(self.translation).neg();

        Some(Self::from_mat3_translation(inv_matrix, inv_translation))
    }

    /// Returns the inverse of `self`.
    ///
    /// # Panics
    /// Panics if `self` is not invertible.
    #[inline]
    #[must_use]
    pub const fn inverse(self) -> Self {
        self.try_inverse().expect("affine transform not invertible")
    }

    /// Returns `true` if `self` is invertible.
    #[inline]
    #[must_use]
    pub const fn is_invertible(self) -> bool {
        self.matrix.is_invertible()
    }

    /// Returns `v` transformed as a point by `self`.
    ///
    /// Applies the full transform including translation.
    #[inline]
    #[must_use]
    pub const fn transform_point(self, v: Vec3) -> Vec3 {
        Vec3::new(
            self.matrix.col0.x * v.x
                + self.matrix.col1.x * v.y
                + self.matrix.col2.x * v.z
                + self.translation.x,
            self.matrix.col0.y * v.x
                + self.matrix.col1.y * v.y
                + self.matrix.col2.y * v.z
                + self.translation.y,
            self.matrix.col0.z * v.x
                + self.matrix.col1.z * v.y
                + self.matrix.col2.z * v.z
                + self.translation.z,
        )
    }

    /// Returns `v` transformed by `self`, applying rotation, scale and shear (NOT translation).
    #[inline]
    #[must_use]
    pub const fn transform_vector(self, v: Vec3) -> Vec3 {
        self.matrix.transform(v)
    }

    /// Returns `true` if `self` and `other` approximately represent the same transform.
    #[inline]
    #[must_use]
    pub fn approx_eq(self, other: Self) -> bool {
        self.matrix.approx_eq(other.matrix) && self.translation.approx_eq(other.translation)
    }

    /// Returns `true` if all components of `self` are finite.
    #[inline]
    #[must_use]
    pub const fn is_finite(self) -> bool {
        self.matrix.is_finite() && self.translation.is_finite()
    }

    /// Returns `true` if any component of `self` is `NaN`.
    #[inline]
    #[must_use]
    pub const fn is_nan(self) -> bool {
        self.matrix.is_nan() || self.translation.is_nan()
    }
}

impl Mul for Affine3 {
    type Output = Self;

    /// Returns the composition of `self` and `rhs`, applying `rhs` first.
    ///
    /// Equivalent to [`mul`](Self::mul()).
    #[inline]
    fn mul(self, rhs: Self) -> Self::Output {
        self.mul(rhs)
    }
}

impl MulAssign for Affine3 {
    /// Sets `self` to the composition of `self` and `rhs`, applying `rhs` first.
    ///
    /// Equivalent to [`mul_assign`](Self::mul_assign()).
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        self.mul_assign(rhs);
    }
}

impl Mul<Vec3> for Affine3 {
    type Output = Vec3;

    /// Returns `rhs` transformed by `self`, applying rotation, scale, shear and translation.
    ///
    /// Equivalent to [`mul_vec3`](Self::mul_vec3()).
    #[inline]
    fn mul(self, rhs: Vec3) -> Self::Output {
        self.mul_vec3(rhs)
    }
}

impl Display for Affine3 {
    /// Formats `self` as `[col0, col1, col2, translation]`,
    /// with each column as `(x, y, z)`.
    ///
    /// Respects the precision of `f`.
    #[inline]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(p) = f.precision() {
            write!(
                f,
                "[{:.*}, {:.*}, {:.*}, {:.*}]",
                p, self.matrix.col0, p, self.matrix.col1, p, self.matrix.col2, p, self.translation
            )
        } else {
            write!(
                f,
                "[{}, {}, {}, {}]",
                self.matrix.col0, self.matrix.col1, self.matrix.col2, self.translation
            )
        }
    }
}

impl Default for Affine3 {
    /// Default is [`IDENTITY`](Self::IDENTITY).
    #[inline]
    fn default() -> Self {
        Self::IDENTITY
    }
}

impl From<[f32; 12]> for Affine3 {
    /// Creates an `Affine3` from a `[f32; 12]`, in order:
    /// `[col0, col1, col2, translation]` with each column as `[x, y, z]`.
    ///
    /// Equivalent to [`from_array`](Self::from_array()).
    #[inline]
    fn from(v: [f32; 12]) -> Self {
        Self::from_array(v)
    }
}

impl From<Affine3> for [f32; 12] {
    /// Returns a copy of the components as `[col0, col1, col2, translation]`,
    /// with each column as `[x, y, z]`.
    ///
    /// Equivalent to [`to_array`](Affine3::to_array()).
    #[inline]
    fn from(v: Affine3) -> Self {
        v.to_array()
    }
}

impl From<Mat4> for Affine3 {
    /// Creates an `Affine3` from a `Mat4`, discarding its projective row.
    ///
    /// Equivalent to [`from_mat4`](Self::from_mat4()).
    #[inline]
    fn from(v: Mat4) -> Self {
        Self::from_mat4(v)
    }
}

#[cfg(feature = "serde")]
impl serde::Serialize for Affine3 {
    /// Serializes as a `[f32; 12]`, the first 9 elements are the linear part (`Mat3`)
    /// in column-major order, the last 3 are the translation.
    ///
    /// # Format
    /// Encoded as a flat 12-element array. Component order matches [`to_array`](Self::to_array()).
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        self.to_array().serialize(serializer)
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for Affine3 {
    /// Deserializes from a `[f32; 12]`, the first 9 elements are the linear part (`Mat3`)
    /// in column-major order, the last 3 are the translation.
    ///
    /// # Format
    /// Decoded as a flat 12-element array. Component order matches [`from_array`](Self::from_array()).
    ///
    /// # Errors
    /// Returns an error if array is not exactly 12 elements.
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        <[f32; 12]>::deserialize(deserializer).map(Self::from_array)
    }
}

#[cfg(feature = "bytemuck")]
unsafe impl bytemuck::Zeroable for Affine3 {}

#[cfg(feature = "bytemuck")]
unsafe impl bytemuck::Pod for Affine3 {}
