/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

using  RobotNet10.Shared.Numbers;

namespace CartographerSharp.Transform;

/// <summary>
/// Transform operations for Cartographer.
/// </summary>
public static class TransformOperations
{
    /// <summary>
    /// Returns the non-negative rotation angle in radians of the 3D transformation.
    /// </summary>
    public static double GetAngle(Rigid3d transform)
    {
        var vec = new Vector3(transform.Rotation.X, transform.Rotation.Y, transform.Rotation.Z);
        var vecNorm = vec.Length();
        return 2.0 * Math.Atan2(vecNorm, Math.Abs(transform.Rotation.W));
    }

    /// <summary>
    /// Returns the yaw component in radians of the given 3D rotation.
    /// Assuming rotation is composed of three rotations around X, then Y, then Z,
    /// returns the angle of the Z rotation.
    /// </summary>
    public static double GetYaw(Quaternion rotation)
    {
        var direction = Vector3.Transform(Vector3.UnitX, rotation);
        return Math.Atan2(direction.Y, direction.X);
    }

    /// <summary>
    /// Returns the yaw component in radians of the given 3D transformation.
    /// </summary>
    public static double GetYaw(Rigid3d transform)
    {
        return GetYaw(transform.Rotation);
    }

    /// <summary>
    /// Returns an angle-axis vector (a vector with the length of the rotation angle
    /// pointing to the direction of the rotation axis) representing the same
    /// rotation as the given quaternion.
    /// </summary>
    public static Vector3 RotationQuaternionToAngleAxisVector(Quaternion quaternion)
    {
        var normalized = Quaternion.Normalize(quaternion);
        
        // We choose the quaternion with positive 'w', i.e., the one with a smaller
        // angle that represents this orientation.
        if (normalized.W < 0.0)
        {
            normalized = new Quaternion(-normalized.X, -normalized.Y, -normalized.Z, -normalized.W);
        }
        
        // We convert the normalized_quaternion into a vector along the rotation axis
        // with length of the rotation angle.
        var vec = new Vector3(normalized.X, normalized.Y, normalized.Z);
        var vecNorm = vec.Length();
        const double kCutoffAngle = 1e-7; // We linearize below this angle.
        
        var angle = 2.0 * Math.Atan2(vecNorm, normalized.W);
        var scale = angle < kCutoffAngle ? 2.0 : angle / Math.Sin(angle / 2.0);
        
        return new Vector3(
            (scale * normalized.X),
            (scale * normalized.Y),
            (scale * normalized.Z)
        );
    }

    /// <summary>
    /// Returns a quaternion representing the same rotation as the given angle-axis vector.
    /// </summary>
    public static Quaternion AngleAxisVectorToRotationQuaternion(Vector3 angleAxis)
    {
        const double kCutoffAngle = 1e-8; // We linearize below this angle.
        
        var norm = angleAxis.Length();
        double scale = 0.5;
        double w = 1.0;
        
        if (norm * norm > kCutoffAngle)
        {
            scale = Math.Sin(norm / 2.0) / norm;
            w = Math.Cos(norm / 2.0);
        }
        
        return Quaternion.Normalize(new Quaternion(
            (scale * angleAxis.X),
            (scale * angleAxis.Y),
            (scale * angleAxis.Z),
            w
        ));
    }

    /// <summary>
    /// Projects 'transform' onto the XY plane.
    /// </summary>
    public static Rigid2d Project2D(Rigid3d transform)
    {
        var translation2D = new Vector2(transform.Translation.X, transform.Translation.Y);
        var yaw = GetYaw(transform);
        return new Rigid2d(translation2D, yaw);
    }

    /// <summary>
    /// Embeds 'transform' into 3D space in the XY plane.
    /// </summary>
    public static Rigid3d Embed3D(Rigid2d transform)
    {
        var translation3D = new Vector3(transform.Translation.X, transform.Translation.Y, 0.0);
        var rotation = Quaternion.CreateFromAxisAngle(Vector3.UnitZ, transform.Rotation);
        return new Rigid3d(translation3D, rotation);
    }

    /// <summary>
    /// Interpolates between two transforms at different times.
    /// </summary>
    public static Rigid3d Interpolate(
        Rigid3d startTransform,
        long startTime,
        Rigid3d endTransform,
        long endTime,
        long targetTime)
    {
        if (startTime == endTime)
        {
            return startTransform;
        }

        var factor = (double)(targetTime - startTime) / (endTime - startTime);
        factor = Math.Max(0.0, Math.Min(1.0, factor)); // Clamp to [0, 1]

        // Interpolate translation linearly
        var interpolatedTranslation = new Vector3(
            (startTransform.Translation.X + factor * (endTransform.Translation.X - startTransform.Translation.X)),
            (startTransform.Translation.Y + factor * (endTransform.Translation.Y - startTransform.Translation.Y)),
            (startTransform.Translation.Z + factor * (endTransform.Translation.Z - startTransform.Translation.Z))
        );

        // Interpolate rotation using SLERP
        var interpolatedRotation = SlerpQuaternions(
            startTransform.Rotation,
            endTransform.Rotation,
            factor
        );

        return new Rigid3d(interpolatedTranslation, interpolatedRotation);
    }

    /// <summary>
    /// Spherical linear interpolation of quaternions.
    /// </summary>
    private static Quaternion SlerpQuaternions(Quaternion start, Quaternion end, double factor)
    {
        start = Quaternion.Normalize(start);
        end = Quaternion.Normalize(end);

        var cosTheta = start.W * end.W + start.X * end.X + start.Y * end.Y + start.Z * end.Z;
        var absCosTheta = Math.Abs(cosTheta);

        double prevScale = 1.0 - factor;
        double nextScale = factor;

        if (absCosTheta < 1.0 - 1e-5)
        {
            var theta = Math.Acos(absCosTheta);
            var sinTheta = Math.Sin(theta);
            prevScale = Math.Sin((1.0 - factor) * theta) / sinTheta;
            nextScale = Math.Sin(factor * theta) / sinTheta;
        }

        if (cosTheta < 0.0)
        {
            nextScale = -nextScale;
        }

        return new Quaternion(
            (prevScale * start.X + nextScale * end.X),
            (prevScale * start.Y + nextScale * end.Y),
            (prevScale * start.Z + nextScale * end.Z),
            (prevScale * start.W + nextScale * end.W)
        );
    }
}