using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using RobotNet10.RobotApp.Devices;
using RobotNet10.Shared;
using RobotNet10.Shared.Geometry;
using RobotNet10.Shared.Sensor;
using System;
using System.Threading;
using System.Threading.Tasks;

namespace RobotNet10.RobotApp.Motion;

/// <summary>
/// EKFService - Extended Kalman Filter service để fusion odometry và IMU data
/// Cải thiện độ chính xác ước lượng trạng thái robot
/// </summary>
public class EKFService : IHostedService, IDisposable
{
    private readonly EKFConfiguration _config;
    private readonly IDeviceProvider _deviceProvider;
    private readonly OdometryService _odometryService;
    private readonly ILogger<EKFService> _logger;
    private readonly object _lock = new();

    // IMU device
    private IInertialMeasurementUnit? _imu;

    // Extended Kalman Filter
    private ExtendedKalmanFilter? _ekf;

    // Filtered odometry output
    private Odometry _filteredOdometry = new();
    private DateTime _lastUpdateTime = DateTime.UtcNow;
    private bool _isFirstUpdate = true;

    // Yaw offset calibration (applied at EKF level, not IMU level)
    private double _yawOffset = 0.0;
    private bool _yawOffsetCalibrated = false;

    // Sequence number for header
    private uint _sequenceNumber = 0;

    // EKF's own orientation (integrated from angular velocity)
    private Quaternion _ekfOrientation = new Quaternion(0, 0, 0, 1); // Identity quaternion
    private DateTime _lastOrientationUpdate = DateTime.UtcNow;
    
    // Dead zone for gyroscope to prevent drift when stationary
    private const double GYRO_DEADZONE_THRESHOLD = 0.001; // rad/s (~0.057 deg/s)

    // Timer for periodic updates
    private Timer? _updateTimer;
    private bool _disposed = false;

    /// <summary>
    /// Gets the filtered odometry data
    /// </summary>
    public Odometry FilteredOdometry
    {
        get
        {
            lock (_lock)
            {
                return _filteredOdometry;
            }
        }
    }

    public EKFService(
        IConfiguration configuration,
        IDeviceProvider deviceProvider,
        OdometryService odometryService,
        ILogger<EKFService> logger)
    {
        _deviceProvider = deviceProvider ?? throw new ArgumentNullException(nameof(deviceProvider));
        _odometryService = odometryService ?? throw new ArgumentNullException(nameof(odometryService));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));

        // Load configuration
        var configSection = configuration.GetSection("Motion:EKF");
        if (!configSection.Exists())
        {
            throw new InvalidOperationException("Configuration section 'Motion:EKF' not found in appsettings.json");
        }

        _config = new EKFConfiguration();
        configSection.Bind(_config);

        // Validate configuration
        ValidateConfiguration();
    }

    private void ValidateConfiguration()
    {
        if (string.IsNullOrWhiteSpace(_config.ImuDeviceId))
            throw new InvalidOperationException("ImuDeviceId is required");

        if (_config.UpdateRateHz <= 0)
            throw new InvalidOperationException("UpdateRateHz must be greater than 0");

        if (string.IsNullOrWhiteSpace(_config.FrameId))
            _config.FrameId = "odom";

        if (string.IsNullOrWhiteSpace(_config.ChildFrameId))
            _config.ChildFrameId = "base_link";
    }

    public async Task StartAsync(CancellationToken cancellationToken)
    {
        _logger.LogInformation("Starting EKFService...");

        try
        {
            // Wait for devices to be connected
            _logger.LogInformation("Waiting for all devices to be connected...");
            var connected = await _deviceProvider.WaitForDevicesConnectedAsync(TimeSpan.FromMinutes(5), cancellationToken);

            if (!connected)
            {
                _logger.LogWarning("Timeout waiting for devices to connect. EKFService will not be initialized.");
                return;
            }

            // Get IMU device
            var imuDevice = _deviceProvider.GetDevice(_config.ImuDeviceId);
            if (imuDevice == null)
            {
                _logger.LogError("IMU device '{DeviceId}' not found", _config.ImuDeviceId);
                return;
            }

            if (imuDevice is not IInertialMeasurementUnit imu)
            {
                _logger.LogError("Device '{DeviceId}' is not an IInertialMeasurementUnit", _config.ImuDeviceId);
                return;
            }

            if (!imuDevice.IsConnected)
            {
                _logger.LogWarning("IMU device '{DeviceId}' is not connected. Status: {Status}. EKFService will wait for connection.",
                    _config.ImuDeviceId, imuDevice.Status);
                // Start retry task
                _ = Task.Run(async () => await RetryInitializationWhenDevicesReadyAsync(cancellationToken));
                return;
            }

            _imu = imu;

            // Initialize EKF
            InitializeEKF();

            // Start periodic update timer
            var updateInterval = TimeSpan.FromMilliseconds(1000.0 / _config.UpdateRateHz);
            var LastTime = DateTime.UtcNow;
            _updateTimer = new Timer(OnTimerCallback, null, updateInterval, updateInterval);

            _logger.LogInformation("EKFService started successfully with update rate {UpdateRate} Hz", _config.UpdateRateHz);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error starting EKFService");
        }

        await Task.CompletedTask;
    }

    private async Task RetryInitializationWhenDevicesReadyAsync(CancellationToken cancellationToken)
    {
        const int maxRetries = 60; // 5 minutes with 5 second intervals
        int retryCount = 0;

        while (retryCount < maxRetries && !cancellationToken.IsCancellationRequested)
        {
            await Task.Delay(TimeSpan.FromSeconds(5), cancellationToken);

            var imuDevice = _deviceProvider.GetDevice(_config.ImuDeviceId);

            if (imuDevice != null && imuDevice.IsConnected && imuDevice is IInertialMeasurementUnit imu)
            {
                _logger.LogInformation("IMU device is now connected. Initializing EKFService...");

                lock (_lock)
                {
                    _imu = imu;
                }

                InitializeEKF();

                var updateInterval = TimeSpan.FromMilliseconds(1000.0 / _config.UpdateRateHz);
                _updateTimer = new Timer(OnTimerCallback, null, updateInterval, updateInterval);

                _logger.LogInformation("EKFService initialized successfully after device connection");
                return;
            }

            retryCount++;
            if (retryCount % 12 == 0) // Log every minute
            {
                _logger.LogInformation("Still waiting for IMU device to connect... (attempt {Attempt}/{MaxAttempts})",
                    retryCount, maxRetries);
            }
        }

        if (retryCount >= maxRetries)
        {
            _logger.LogWarning("Timeout waiting for IMU device to connect. EKFService will not be initialized.");
        }
    }

    public async Task StopAsync(CancellationToken cancellationToken)
    {
        _logger.LogInformation("Stopping EKFService...");

        try
        {
            _updateTimer?.Dispose();
            _updateTimer = null;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error stopping EKFService");
        }

        await Task.CompletedTask;
    }

    private void OnTimerCallback(object? state)
    {
        try
        {
            UpdateEKF();
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error updating EKF");
        }
    }

    private void InitializeEKF()
    {
        lock (_lock)
        {
            // Get initial state from odometry
            var odom = _odometryService.CurrentOdometry;
            var initialX = odom.Pose.Pose.Position.X;
            var initialY = odom.Pose.Pose.Position.Y;
            var initialTheta = QuaternionToYaw(odom.Pose.Pose.Orientation);

            // Initial state: [x, y, theta, vx, vy, omega]
            double[] initialState = new double[6]
            {
                initialX,
                initialY,
                initialTheta,
                0.0, // vx
                0.0, // vy
                0.0  // omega
            };

            // Initial covariance (diagonal)
            double[] initialCovDiag = new double[6]
            {
                _config.InitialPositionUncertainty,      // x
                _config.InitialPositionUncertainty,      // y
                _config.InitialOrientationUncertainty,   // theta
                _config.InitialVelocityUncertainty,      // vx
                _config.InitialVelocityUncertainty,      // vy
                _config.InitialVelocityUncertainty       // omega
            };
            var initialCovariance = MatrixHelper.CreateDiagonal(initialCovDiag);

            // Process noise covariance (diagonal)
            double[] processNoiseDiag = new double[6]
            {
                _config.ProcessNoisePosition,      // x
                _config.ProcessNoisePosition,      // y
                _config.ProcessNoiseOrientation,   // theta
                _config.ProcessNoiseVelocity,      // vx
                _config.ProcessNoiseVelocity,      // vy
                _config.ProcessNoiseVelocity       // omega
            };
            var processNoise = MatrixHelper.CreateDiagonal(processNoiseDiag);

            // Measurement noise covariance (diagonal) - [x, y, theta]
            double[] measurementNoiseDiag = new double[3]
            {
                _config.MeasurementNoisePosition,      // x
                _config.MeasurementNoisePosition,      // y
                _config.MeasurementNoiseOrientation    // theta
            };
            var measurementNoise = MatrixHelper.CreateDiagonal(measurementNoiseDiag);

            // Create EKF
            _ekf = new ExtendedKalmanFilter(
                initialState,
                initialCovariance,
                processNoise,
                measurementNoise);

            _lastUpdateTime = DateTime.UtcNow;
            _isFirstUpdate = true;

            // Reset yaw offset - will be calibrated on first update
            _yawOffset = 0.0;
            _yawOffsetCalibrated = false;

            // Initialize EKF orientation (start from identity quaternion)
            _ekfOrientation = new Quaternion(0, 0, 0, 1);
            _lastOrientationUpdate = DateTime.UtcNow;

            _logger.LogInformation("EKF initialized with state: X={X:F3}, Y={Y:F3}, Theta={Theta:F3}",
                initialX, initialY, initialTheta);
        }
    }

    private void UpdateEKF()
    {
        if (_ekf == null || _imu == null)
            return;

        // Check if IMU is still connected
        var imuDevice = _deviceProvider.GetDevice(_config.ImuDeviceId);
        if (imuDevice == null || !imuDevice.IsConnected)
        {
            return;
        }

        lock (_lock)
        {
            var currentTime = DateTime.UtcNow;
            var dt = (currentTime - _lastUpdateTime).TotalSeconds;

            if (dt <= 0 || _isFirstUpdate)
            {
                _lastUpdateTime = currentTime;
                _isFirstUpdate = false;
                return;
            }

            // Get odometry data (for velocities)
            var odom = _odometryService.CurrentOdometry;
            var vx = odom.Twist.Twist.Linear.X;
            var vy = odom.Twist.Twist.Linear.Y;
            var omega = odom.Twist.Twist.Angular.Z;

            // PREDICTION STEP: Propagate state using odometry velocities
            _ekf.Predict(vx, vy, omega, dt);

            // Get IMU orientation (raw from sensor)
            // double rawYaw = 0.0;
            // var quaternion = _imu.CachedQuaternion;
            // if (quaternion.HasValue)
            // {
            //     rawYaw = QuaternionToYaw(quaternion.Value.Quaternion);
            // }
            // else
            // {
            //     // Fallback to Euler angles
            //     var orientation = _imu.CachedOrientation;
            //     rawYaw = orientation.Vector.Z; // Yaw is Z component
            // }

            // Calibrate yaw offset on first reading (set current yaw as zero reference)
            // if (!_yawOffsetCalibrated)
            // {
            //     _yawOffset = rawYaw;
            //     _yawOffsetCalibrated = true;
            //     _logger.LogInformation("IMU Yaw offset calibrated at EKF level: {Offset:F3} rad ({OffsetDeg:F1}°)",
            //         _yawOffset, _yawOffset * 180.0 / Math.PI);
            // }

            // ===== EKF TÍCH PHÂN ANGULAR VELOCITY ĐỂ TÍNH ORIENTATION =====
            // Get angular velocity from IMU (raw data, không dùng CachedOrientation)
            var angularVelocity = _imu.CachedAngularVelocity.Vector;
            double wx = angularVelocity.X;
            double wy = angularVelocity.Y;
            double wz = angularVelocity.Z;
            
            // Transform from IMU frame to Odom frame (180° rotation around X-axis)
            // Rotation matrix: [1, 0, 0; 0, -1, 0; 0, 0, -1]
            // This accounts for IMU mounting orientation (upside down relative to base_link)
            wx = wx;      // X unchanged
            wy = -wy;     // Y negated
            wz = -wz;     // Z negated
            
            // Apply dead zone to filter gyro noise/drift when stationary
            // if (Math.Abs(wx) < GYRO_DEADZONE_THRESHOLD) wx = 0.0;
            // if (Math.Abs(wy) < GYRO_DEADZONE_THRESHOLD) wy = 0.0;
            // if (Math.Abs(wz) < GYRO_DEADZONE_THRESHOLD) wz = 0.0;

            // Integrate quaternion using angular velocity
            // Formula: q_new = q_old + 0.5 * dt * [0, wx, wy, wz] * q_old
            var q = _ekfOrientation;
            
            // Half of angular velocity components
            double halfWx = 0.5 * wx;
            double halfWy = 0.5 * wy;
            double halfWz = 0.5 * wz;
            
            // Quaternion derivative: dq/dt = 0.5 * [0, wx, wy, wz] * q
            double dqx = halfWx * q.W + halfWy * q.Z - halfWz * q.Y;
            double dqy = halfWy * q.W - halfWx * q.Z + halfWz * q.X;
            double dqz = halfWz * q.W + halfWx * q.Y - halfWy * q.X;
            double dqw = -halfWx * q.X - halfWy * q.Y - halfWz * q.Z;
            
            // Integrate: q_new = q_old + dq * dt
            double newQx = q.X + dqx * dt;
            double newQy = q.Y + dqy * dt;
            double newQz = q.Z + dqz * dt;
            double newQw = q.W + dqw * dt;
            
            // Normalize quaternion to maintain unit length
            double norm = Math.Sqrt(newQx * newQx + newQy * newQy + newQz * newQz + newQw * newQw);
            if (norm > 1e-6)  // Avoid division by zero
            {
                newQx /= norm;
                newQy /= norm;
                newQz /= norm;
                newQw /= norm;
            }
            
            // Update EKF orientation
            _ekfOrientation = new Quaternion(newQx, newQy, newQz, newQw);
            _lastOrientationUpdate = currentTime;

            // Convert quaternion to yaw for EKF update
            double measuredTheta = QuaternionToYaw(_ekfOrientation);
            measuredTheta = NormalizeAngle(measuredTheta);

            // UPDATE STEP: Correct state using integrated orientation from angular velocity
            _ekf.UpdateOrientation(measuredTheta);
            // Console.WriteLine($"EKF Update: X={_ekf.X}, Y={_ekf.Y}, Theta={_ekf.Theta}, LinearVel={_ekf.Vx}, AngularVel={_ekf.Omega}");    
        
            // Build filtered odometry message
            var filteredPose = new Pose
            {
                Position = new Point
                {
                    X = _ekf.X,
                    Y = _ekf.Y,
                    Z = 0.0
                },
                Orientation = new Quaternion
                {
                    X = 0.0,
                    Y = 0.0,
                    Z = Math.Sin(_ekf.Theta / 2.0),
                    W = Math.Cos(_ekf.Theta / 2.0)
                }
            };

            var filteredTwist = new Twist
            {
                Linear = new Vector3
                {
                    X = _ekf.Vx,
                    Y = _ekf.Vy,
                    Z = 0.0
                },
                Angular = new Vector3
                {
                    X = 0.0,
                    Y = 0.0,
                    Z = _ekf.Omega
                }
            };

            // Extract covariance from EKF state covariance
            var stateCov = _ekf.Covariance;
            var poseCovariance = new double[36]; // 6x6
            var twistCovariance = new double[36]; // 6x6

            // Map state covariance to pose covariance
            // Pose: [x, y, z, roll, pitch, yaw]

            poseCovariance[0] = stateCov[0, 0];   // x
            poseCovariance[7] = stateCov[1, 1];   // y
            poseCovariance[14] = 0.0;             // z (not estimated)
            poseCovariance[21] = 0.0;             // roll (not estimated)
            poseCovariance[28] = 0.0;             // pitch (not estimated)
            poseCovariance[35] = stateCov[2, 2];  // yaw (theta)

            // Map state covariance to twist covariance
            // Twist: [vx, vy, vz, wx, wy, ywz]

            twistCovariance[0] = stateCov[3, 3];  // vx
            twistCovariance[7] = stateCov[4, 4];  // vy
            twistCovariance[14] = 0.0;            // vz (not estimated)
            twistCovariance[21] = 0.0;            // wx (not estimated)
            twistCovariance[28] = 0.0;            // wy (not estimated)
            twistCovariance[35] = stateCov[5, 5]; // wz (omega)

            _filteredOdometry = new Odometry
            {
                Header = new Header
                {
                    Seq = _sequenceNumber++,
                    Stamp = currentTime,
                    FrameId = _config.FrameId
                },
                ChildFrameId = _config.ChildFrameId,
                Pose = new PoseWithCovariance
                {
                    Pose = filteredPose,
                    Covariance = poseCovariance
                },
                Twist = new TwistWithCovariance
                {
                    Twist = filteredTwist,
                    Covariance = twistCovariance
                }
            };

            _lastUpdateTime = currentTime;

            // Log periodically (every 2 seconds)
            if (_sequenceNumber % 40 == 0) // At 20Hz, 40 updates = 2 seconds
            {
                _logger.LogDebug("EKF State: X={X:F3}, Y={Y:F3}, Theta={Theta:F3}, Vx={Vx:F3}, Omega={Omega:F3}",
                    _ekf.X, _ekf.Y, _ekf.Theta, _ekf.Vx, _ekf.Omega);
            }
        }
    }

    private static double QuaternionToYaw(Quaternion q)
    {
        // Yaw = atan2(2*(w*z + x*y), 1 - 2*(y^2 + z^2))
        var sinYaw = 2.0 * (q.W * q.Z + q.X * q.Y);
        var cosYaw = 1.0 - 2.0 * (q.Y * q.Y + q.Z * q.Z);
        return Math.Atan2(sinYaw, cosYaw);
    }

    /// <summary>
    /// Normalize angle to [-π, π] range
    /// </summary>
    private static double NormalizeAngle(double angle)
    {
        while (angle > Math.PI) angle -= 2.0 * Math.PI;
        while (angle < -Math.PI) angle += 2.0 * Math.PI;
        return angle;
    }

    /// <summary>
    // /// Manually recalibrate yaw offset - reset current yaw as new zero reference
    // /// </summary>
    // public void RecalibrateYawOffset()
    // {
    //     lock (_lock)
    //     {
    //         _yawOffsetCalibrated = false;
    //         _logger.LogInformation("Yaw offset recalibration requested. Will recalibrate on next update.");
    //     }
    // }

    public void Dispose()
    {
        if (_disposed)
            return;

        _disposed = true;

        try
        {
            _updateTimer?.Dispose();
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error disposing EKFService");
        }
    }
}