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

namespace CartographerSharp.Sensor;

/// <summary>
/// A compressed representation of a point cloud consisting of a collection of
/// points (Vector3) without time information.
/// Internally, points are grouped by blocks. Each block encodes a bit of meta
/// data (number of points in block, coordinates of the block) and encodes each
/// point with a fixed bit rate in relation to the block.
/// </summary>
public class CompressedPointCloud : IEnumerable<RangefinderPoint>
{
    private const double kPrecision = 0.001f; // in meters
    private const int kBitsPerCoordinate = 10;
    private const int kCoordinateMask = (1 << kBitsPerCoordinate) - 1;
    private const int kMaxBitsPerDirection = 23;

    private readonly List<int> _pointData;
    private readonly int _numPoints;

    /// <summary>
    /// Creates an empty compressed point cloud.
    /// </summary>
    public CompressedPointCloud()
    {
        _pointData = [];
        _numPoints = 0;
    }

    /// <summary>
    /// Creates a compressed point cloud from a point cloud.
    /// </summary>
    public CompressedPointCloud(PointCloud pointCloud)
    {
        _numPoints = pointCloud.Count;
        
        // Distribute points into blocks.
        // Using Dictionary to simulate HybridGrid behavior
        var blocks = new Dictionary<(int x, int y, int z), List<(Vector3 rasterPoint, int index)>>();
        
        for (int pointIndex = 0; pointIndex < pointCloud.Count; pointIndex++)
        {
            var point = pointCloud[pointIndex];
            var absMax = Math.Max(Math.Max(Math.Abs(point.Position.X), Math.Abs(point.Position.Y)), Math.Abs(point.Position.Z));
            
            if (absMax / kPrecision >= (1 << kMaxBitsPerDirection))
            {
                throw new ArgumentOutOfRangeException(nameof(pointCloud), 
                    $"Point out of bounds: {point.Position}");
            }

            var rasterPoint = new Vector3(
                Math.Round(point.Position.X / kPrecision),
                Math.Round(point.Position.Y / kPrecision),
                Math.Round(point.Position.Z / kPrecision)
            );

            var blockCoordinate = (
                (int)rasterPoint.X >> kBitsPerCoordinate,
                (int)rasterPoint.Y >> kBitsPerCoordinate,
                (int)rasterPoint.Z >> kBitsPerCoordinate
            );

            var relativePoint = new Vector3(
                (int)rasterPoint.X & kCoordinateMask,
                (int)rasterPoint.Y & kCoordinateMask,
                (int)rasterPoint.Z & kCoordinateMask
            );

            if (!blocks.TryGetValue(blockCoordinate, out var block))
            {
                block = [];
                blocks[blockCoordinate] = block;
            }
            block.Add((relativePoint, pointIndex));
        }

        // Encode blocks.
        _pointData = [];
        foreach (var (blockCoord, rasterPoints) in blocks)
        {
            if (rasterPoints.Count > int.MaxValue)
            {
                throw new ArgumentException("Block too large");
            }

            _pointData.Add(rasterPoints.Count);
            _pointData.Add(blockCoord.x);
            _pointData.Add(blockCoord.y);
            _pointData.Add(blockCoord.z);

            foreach (var (rasterPoint, _) in rasterPoints)
            {
                int encoded = (int)(((((int)rasterPoint.Z << kBitsPerCoordinate) + (int)rasterPoint.Y) << kBitsPerCoordinate) + (int)rasterPoint.X);
                _pointData.Add(encoded);
            }
        }
    }

    /// <summary>
    /// Creates a compressed point cloud from a proto.
    /// </summary>
    public CompressedPointCloud(CompressedPointCloud proto)
    {
        _numPoints = proto.NumPoints;
        _pointData = [.. proto.PointData];
    }

    /// <summary>
    /// Gets the number of points.
    /// </summary>
    public int NumPoints => _numPoints;

    /// <summary>
    /// Gets the point data.
    /// </summary>
    internal List<int> PointData => _pointData;

    /// <summary>
    /// Checks if the point cloud is empty.
    /// </summary>
    public bool IsEmpty => _numPoints == 0;

    /// <summary>
    /// Gets the number of points.
    /// </summary>
    public int Count => _numPoints;

    /// <summary>
    /// Returns decompressed point cloud.
    /// </summary>
    public PointCloud Decompress()
    {
        var decompressed = new PointCloud();
        foreach (var point in this)
        {
            decompressed.Add(point);
        }
        return decompressed;
    }

    /// <summary>
    /// Gets an enumerator for the points.
    /// </summary>
    public IEnumerator<RangefinderPoint> GetEnumerator()
    {
        return new ConstIterator(this);
    }

    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }

    /// <summary>
    /// Converts to proto representation.
    /// </summary>
    public Models.Sensor.CompressedPointCloud ToProto()
    {
        return new Models.Sensor.CompressedPointCloud(_numPoints, [.. _pointData]);
    }

    /// <summary>
    /// Creates from proto representation.
    /// </summary>
    public static CompressedPointCloud FromProto(Models.Sensor.CompressedPointCloud proto)
    {
        return new CompressedPointCloud(proto.NumPoints, proto.PointData ?? []);
    }

    /// <summary>
    /// Creates a compressed point cloud from proto data.
    /// </summary>
    private CompressedPointCloud(int numPoints, List<int> pointData)
    {
        _numPoints = numPoints;
        _pointData = pointData ?? [];
    }

    /// <summary>
    /// Forward iterator for compressed point clouds.
    /// </summary>
    private class ConstIterator : IEnumerator<RangefinderPoint>
    {
        private readonly CompressedPointCloud _compressedPointCloud;
        private int _remainingPoints;
        private int _remainingPointsInCurrentBlock;
        private int _inputIndex;
        private Vector3 _currentPoint;
        private Vector3 _currentBlockCoordinates;

        public ConstIterator(CompressedPointCloud compressedPointCloud)
        {
            _compressedPointCloud = compressedPointCloud;
            _remainingPoints = compressedPointCloud._numPoints;
            _remainingPointsInCurrentBlock = 0;
            _inputIndex = 0;

            if (_remainingPoints > 0)
            {
                ReadNextPoint();
            }
        }

        public RangefinderPoint Current { get; private set; }

        object IEnumerator.Current => Current;

        public bool MoveNext()
        {
            if (_remainingPoints <= 0)
            {
                return false;
            }

            Current = new RangefinderPoint(_currentPoint);
            _remainingPoints--;

            if (_remainingPoints > 0)
            {
                ReadNextPoint();
            }

            return true;
        }

        public void Reset()
        {
            _remainingPoints = _compressedPointCloud._numPoints;
            _remainingPointsInCurrentBlock = 0;
            _inputIndex = 0;
            if (_remainingPoints > 0)
            {
                ReadNextPoint();
            }
        }

        public void Dispose()
        {
            // Nothing to dispose
        }

        private void ReadNextPoint()
        {
            if (_remainingPointsInCurrentBlock == 0)
            {
                if (_inputIndex >= _compressedPointCloud._pointData.Count)
                {
                    return;
                }

                _remainingPointsInCurrentBlock = _compressedPointCloud._pointData[_inputIndex++];
                
                if (_inputIndex + 3 > _compressedPointCloud._pointData.Count)
                {
                    return;
                }

                _currentBlockCoordinates = new Vector3(
                    _compressedPointCloud._pointData[_inputIndex++] << kBitsPerCoordinate,
                    _compressedPointCloud._pointData[_inputIndex++] << kBitsPerCoordinate,
                    _compressedPointCloud._pointData[_inputIndex++] << kBitsPerCoordinate
                );
            }

            _remainingPointsInCurrentBlock--;

            if (_inputIndex >= _compressedPointCloud._pointData.Count)
            {
                return;
            }

            int point = _compressedPointCloud._pointData[_inputIndex++];
            const int kMask = (1 << kBitsPerCoordinate) - 1;

            _currentPoint = new Vector3(
                ((_currentBlockCoordinates.X + (point & kMask)) * kPrecision),
                ((_currentBlockCoordinates.Y + ((point >> kBitsPerCoordinate) & kMask)) * kPrecision),
                ((_currentBlockCoordinates.Z + (point >> (2 * kBitsPerCoordinate))) * kPrecision)
            );
        }
    }
}