/*
 * 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 CartographerSharp.Transform;
using System;
using  RobotNet10.Shared.Numbers;

namespace CartographerSharp.Mapping;

/// <summary>
/// Keeps track of the orientation using angular velocities and linear
/// accelerations from an IMU. Because averaged linear acceleration (assuming
/// slow movement) is a direct measurement of gravity, roll/pitch does not drift,
/// though yaw does.
/// </summary>
public class ImuTracker
{
    private readonly double _imuGravityTimeConstant;
    private long _time;
    private long _lastLinearAccelerationTime;
    private Quaternion _orientation;
    private Vector3 _gravityVector;
    private Vector3 _imuAngularVelocity;

    // Recovery tracking: count consecutive invalid gravity states
    private int _invalidGravityCount;
    private const int MaxInvalidGravityBeforeReset = 10; // Reset after 10 consecutive invalid states

    public ImuTracker(double imuGravityTimeConstant, long time)
    {
        _imuGravityTimeConstant = imuGravityTimeConstant;
        _time = time;
        _lastLinearAccelerationTime = long.MinValue;
        _orientation = Quaternion.Identity;
        _gravityVector = Vector3.UnitZ;
        _imuAngularVelocity = Vector3.Zero;
        _invalidGravityCount = 0;
    }

    /// <summary>
    /// Copy constructor.
    /// </summary>
    public ImuTracker(ImuTracker other)
    {
        _imuGravityTimeConstant = other._imuGravityTimeConstant;
        _time = other._time;
        _lastLinearAccelerationTime = other._lastLinearAccelerationTime;
        _orientation = other._orientation;
        _gravityVector = other._gravityVector;
        _imuAngularVelocity = other._imuAngularVelocity;
        _invalidGravityCount = other._invalidGravityCount;
    }

    /// <summary>
    /// Advances to the given 'time' and updates the orientation to reflect this.
    /// </summary>
    public void Advance(long time)
    {
        if (time < _time)
        {
            // DEBUG: Log detailed timestamp information
            var timeDiffMs = (_time - time) / TimeSpan.TicksPerMillisecond;
            
            // If the difference is small (< 100ms), it's likely due to synchronization issues
            // between multiple sensors or old scans being reprocessed. In this case, we advance
            // to current time instead of throwing. This is a workaround for RangeDataCollator
            // synchronization issues and old scan reprocessing.
            if (timeDiffMs >= 100)
            {
                // For larger differences, throw exception as it indicates a real problem
                throw new ArgumentException($"time ({time}) must be >= current time ({_time}), diff={timeDiffMs}ms", nameof(time));
            }
        }

        var deltaT = (time - _time) / 10_000_000.0; // Convert ticks to seconds (10 million ticks per second)
        var rotation = TransformOperations.AngleAxisVectorToRotationQuaternion(
            _imuAngularVelocity * deltaT);
        _orientation = Quaternion.Normalize(_orientation * rotation);
        
        // Rotate gravity vector by inverse rotation (conjugate)
        // In C++: gravity_vector_ = rotation.conjugate() * gravity_vector_
        // In C#: equivalent to Transform with conjugate quaternion
        var rotationConjugate = Quaternion.Conjugate(rotation);
        _gravityVector = Vector3.Transform(_gravityVector, rotationConjugate);
        
        _time = time;
    }

    /// <summary>
    /// Updates from an IMU reading (in the IMU frame).
    /// </summary>
    public void AddImuLinearAccelerationObservation(Vector3 imuLinearAcceleration)
    {
        // Validate input: reject zero or near-zero acceleration (invalid sensor data)
        var inputMagnitude = imuLinearAcceleration.Length();
        if (inputMagnitude < 0.1) // Less than 0.1 m/s² is invalid (should be ~9.81 when stationary)
        {
            // Skip invalid reading, don't update state
            return;
        }

        // Update the 'gravity_vector_' with an exponential moving average using the
        // 'imu_gravity_time_constant'.
        var deltaT = _lastLinearAccelerationTime > long.MinValue
            ? (_time - _lastLinearAccelerationTime) / 10_000_000.0 // Convert ticks to seconds (10 million ticks per second)
            : double.PositiveInfinity;
        _lastLinearAccelerationTime = _time;

        var alpha = 1.0 - Math.Exp(-deltaT / _imuGravityTimeConstant);
        _gravityVector = (1.0 - alpha) * _gravityVector + alpha * imuLinearAcceleration;

        // Change the 'orientation_' so that it agrees with the current 'gravity_vector_'.
        // Match C++: FromTwoVectors(gravity_vector_, orientation_.conjugate() * Eigen::Vector3d::UnitZ())
        // This computes rotation from gravity_vector_ (in IMU frame) to UnitZ transformed to IMU frame
        var unitZInImuFrame = Vector3.Transform(Vector3.UnitZ, Quaternion.Inverse(_orientation));
        var rotation = FromTwoVectors(_gravityVector, unitZInImuFrame);
        _orientation = Quaternion.Normalize(_orientation * rotation);

        // Validate: gravity vector transformed by orientation should point up (positive Z in world frame)
        // When IMU is level, gravity in world frame should be (0, 0, +g) after transform
        var transformedGravity = Vector3.Transform(_gravityVector, _orientation);
        var normalizedTransformedGravity = Vector3.Normalize(transformedGravity);

        if (transformedGravity.Z <= 0 || normalizedTransformedGravity.Z < 0.9)
        {
            _invalidGravityCount++;

            // Recovery: if too many consecutive invalid states, reset to known good state
            if (_invalidGravityCount >= MaxInvalidGravityBeforeReset)
            {
                Console.WriteLine($"[ImuTracker] RECOVERY: Resetting after {_invalidGravityCount} consecutive invalid gravity states. " +
                    $"TransformedGravity=({transformedGravity.X:F3},{transformedGravity.Y:F3},{transformedGravity.Z:F3})");

                // Reset gravity vector to point in the direction of current acceleration
                // (assuming robot is mostly stationary, acceleration ≈ gravity)
                _gravityVector = Vector3.Normalize(imuLinearAcceleration) * 9.81;

                // Reset orientation to align gravity with world Z-axis
                var gravityDirection = Vector3.Normalize(_gravityVector);
                _orientation = FromTwoVectors(gravityDirection, Vector3.UnitZ);

                _invalidGravityCount = 0;
            }
        }
        else
        {
            // Valid state - reset counter
            _invalidGravityCount = 0;
        }
    }

    /// <summary>
    /// Updates from an IMU reading (in the IMU frame).
    /// </summary>
    public void AddImuAngularVelocityObservation(Vector3 imuAngularVelocity)
    {
        _imuAngularVelocity = imuAngularVelocity;
    }

    /// <summary>
    /// Query the current time.
    /// </summary>
    public long Time => _time;

    /// <summary>
    /// Query the current orientation estimate.
    /// </summary>
    public Quaternion Orientation => _orientation;

    /// <summary>
    /// Computes a quaternion that rotates vector 'a' to vector 'b'.
    /// Equivalent to Eigen::Quaterniond::FromTwoVectors().
    /// </summary>
    private static Quaternion FromTwoVectors(Vector3 a, Vector3 b)
    {
        // Normalize input vectors
        a = Vector3.Normalize(a);
        b = Vector3.Normalize(b);
        
        // If vectors are parallel, return identity
        var dot = Vector3.Dot(a, b);
        if (Math.Abs(dot - 1.0) < 1e-6)
        {
            return Quaternion.Identity;
        }
        
        // If vectors are opposite, need special handling
        if (Math.Abs(dot + 1.0) < 1e-6)
        {
            // Find an orthogonal vector to 'a'
            Vector3 orthogonal;
            if (Math.Abs(a.X) < Math.Abs(a.Y))
            {
                orthogonal = Vector3.UnitX;
            }
            else
            {
                orthogonal = Vector3.UnitY;
            }
            orthogonal = Vector3.Normalize(Vector3.Cross(a, orthogonal));
            
            // Create 180-degree rotation around orthogonal axis
            return Quaternion.CreateFromAxisAngle(orthogonal, Math.PI);
        }
        
        // General case: compute rotation axis and angle
        var axis = Vector3.Cross(a, b);
        var axisLength = axis.Length();
        
        if (axisLength < 1e-6)
        {
            return Quaternion.Identity;
        }
        
        axis = Vector3.Normalize(axis);
        var angle = Math.Acos(Math.Clamp(dot, -1.0, 1.0));
        
        return Quaternion.CreateFromAxisAngle(axis, angle);
    }
}