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

namespace CartographerSharp.Mapping.D3D;

/// <summary>
/// Utility functions for HybridGrid indexing.
/// </summary>
internal static class HybridGridUtils
{
    /// <summary>
    /// Converts an 'index' with each dimension from 0 to 2^'bits' - 1 to a flat z-major index.
    /// </summary>
    public static int ToFlatIndex(Array3i index, int bits)
    {
        if (index.X < 0 || index.Y < 0 || index.Z < 0 ||
            index.X >= (1 << bits) || index.Y >= (1 << bits) || index.Z >= (1 << bits))
        {
            throw new ArgumentOutOfRangeException(nameof(index), 
                $"Index {index} is out of range for bits={bits}");
        }
        return (((index.Z << bits) + index.Y) << bits) + index.X;
    }

    /// <summary>
    /// Converts a flat z-major 'index' to a 3-dimensional index with each dimension
    /// from 0 to 2^'bits' - 1.
    /// </summary>
    public static Array3i To3DIndex(int index, int bits)
    {
        if (index < 0 || index >= (1 << (3 * bits)))
        {
            throw new ArgumentOutOfRangeException(nameof(index),
                $"Index {index} is out of range for bits={bits}");
        }
        int mask = (1 << bits) - 1;
        return new Array3i(
            index & mask,
            (index >> bits) & mask,
            (index >> bits) >> bits);
    }

    /// <summary>
    /// Checks if a value is the default value.
    /// </summary>
    public static bool IsDefaultValue<T>(T value) where T : struct
    {
        return EqualityComparer<T>.Default.Equals(value, default);
    }

    /// <summary>
    /// Checks if a list is empty (default value for collections).
    /// </summary>
    public static bool IsDefaultValue<T>(List<T>? value)
    {
        return value == null || value.Count == 0;
    }
}

/// <summary>
/// A flat grid of '2^kBits' x '2^kBits' x '2^kBits' voxels storing values of
/// type 'TValueType' in contiguous memory. Indices in each dimension are 0-based.
/// </summary>
internal class FlatGrid<TValueType> where TValueType : struct
{
    private const int kBits = 3; // Fixed at 3 bits = 8x8x8 = 512 cells
    private const int kGridSize = 1 << kBits; // 8
    private const int kTotalCells = 1 << (3 * kBits); // 512

    private readonly TValueType[] _cells;

    public FlatGrid()
    {
        _cells = new TValueType[kTotalCells];
        // Values are already default-initialized
    }

    /// <summary>
    /// Returns the number of voxels per dimension.
    /// </summary>
    public static int GridSize => kGridSize;

    /// <summary>
    /// Returns the value stored at 'index', each dimension of 'index' being
    /// between 0 and grid_size() - 1.
    /// </summary>
    public TValueType GetValue(Array3i index)
    {
        return _cells[HybridGridUtils.ToFlatIndex(index, kBits)];
    }

    /// <summary>
    /// Returns a reference to the value at 'index' to allow changing it.
    /// </summary>
    public ref TValueType GetMutableValue(Array3i index)
    {
        return ref _cells[HybridGridUtils.ToFlatIndex(index, kBits)];
    }

    /// <summary>
    /// Iterator for iterating over all values not comparing equal to the default constructed value.
    /// </summary>
    public class Iterator : IEnumerator<(Array3i Index, TValueType Value)>
    {
        private readonly FlatGrid<TValueType> _grid;
        private int _currentIndex;
        private (Array3i Index, TValueType Value)? _current;

        public Iterator(FlatGrid<TValueType> grid)
        {
            _grid = grid;
            _currentIndex = -1;
            MoveNext();
        }

        public bool MoveNext()
        {
            _currentIndex++;
            while (_currentIndex < _grid._cells.Length)
            {
                var value = _grid._cells[_currentIndex];
                if (!HybridGridUtils.IsDefaultValue(value))
                {
                    var index = HybridGridUtils.To3DIndex(_currentIndex, kBits);
                    _current = (index, value);
                    return true;
                }
                _currentIndex++;
            }
            _current = null;
            return false;
        }

        public void Reset()
        {
            _currentIndex = -1;
            _current = null;
        }

        public (Array3i Index, TValueType Value) Current => _current!.Value;

        object IEnumerator.Current => Current;

        public void Dispose()
        {
            GC.SuppressFinalize(this);
        }
    }

    /// <summary>
    /// Gets an enumerator for all non-default values.
    /// </summary>
    public Iterator GetEnumerator()
    {
        return new Iterator(this);
    }
}

/// <summary>
/// A grid consisting of '2^kBits' x '2^kBits' x '2^kBits' grids of type 'FlatGrid'.
/// Wrapped grids are constructed on first access via 'GetMutableValue()'.
/// This is a concrete implementation for the specific case: NestedGrid<FlatGrid<TValueType>, 3>
/// </summary>
internal class NestedGrid<TValueType> where TValueType : struct
{
    private const int kBits = 3; // Fixed at 3 bits = 8x8x8 = 512 meta cells
    private const int kWrappedGridSize = 8; // FlatGrid<TValueType>.GridSize = 8
    private readonly FlatGrid<TValueType>?[] _metaCells;

    public NestedGrid()
    {
        _metaCells = new FlatGrid<TValueType>?[1 << (3 * kBits)]; // 512
    }

    public static int GridSize => kWrappedGridSize << kBits; // 8 * 8 = 64

    public TValueType GetValue(Array3i index)
    {
        var metaIndex = NestedGrid<TValueType>.GetMetaIndex(index);
        var metaCell = _metaCells[HybridGridUtils.ToFlatIndex(metaIndex, kBits)];
        if (metaCell == null)
        {
            return default;
        }
        var innerIndex = index - metaIndex * kWrappedGridSize;
        return metaCell.GetValue(innerIndex);
    }

    public ref TValueType GetMutableValue(Array3i index)
    {
        var metaIndex = NestedGrid<TValueType>.GetMetaIndex(index);
        var flatIndex = HybridGridUtils.ToFlatIndex(metaIndex, kBits);
        if (_metaCells[flatIndex] == null)
        {
            _metaCells[flatIndex] = new FlatGrid<TValueType>();
        }
        var innerIndex = index - metaIndex * kWrappedGridSize;
        return ref _metaCells[flatIndex]!.GetMutableValue(innerIndex);
    }

    public IEnumerator<(Array3i Index, TValueType Value)> GetEnumerator()
    {
        return new Iterator(this);
    }

    private static Array3i GetMetaIndex(Array3i index)
    {
        if (index.X < 0 || index.Y < 0 || index.Z < 0)
        {
            throw new ArgumentOutOfRangeException(nameof(index), $"Index {index} has negative components");
        }
        var metaIndex = index / kWrappedGridSize;
        if (metaIndex.X >= (1 << kBits) || metaIndex.Y >= (1 << kBits) || metaIndex.Z >= (1 << kBits))
        {
            throw new ArgumentOutOfRangeException(nameof(index), $"Meta index {metaIndex} is out of range");
        }
        return metaIndex;
    }

    /// <summary>
    /// Iterator for iterating over all non-default values.
    /// </summary>
    public class Iterator : IEnumerator<(Array3i Index, TValueType Value)>
    {
        private readonly NestedGrid<TValueType> _grid;
        private int _currentMetaIndex;
        private FlatGrid<TValueType>.Iterator? _nestedIterator;
        private (Array3i Index, TValueType Value)? _current;

        public Iterator(NestedGrid<TValueType> grid)
        {
            _grid = grid;
            _currentMetaIndex = -1;
            AdvanceToValidNestedIterator();
        }

        private void AdvanceToValidNestedIterator()
        {
            while (_currentMetaIndex < _grid._metaCells.Length - 1)
            {
                _currentMetaIndex++;
                if (_currentMetaIndex >= _grid._metaCells.Length)
                {
                    _nestedIterator = null;
                    _current = null;
                    return;
                }

                var metaCell = _grid._metaCells[_currentMetaIndex];
                if (metaCell != null)
                {
                    _nestedIterator = metaCell.GetEnumerator();
                    if (_nestedIterator.MoveNext())
                    {
                        var (innerIndex, value) = _nestedIterator.Current;
                        var metaIndex = HybridGridUtils.To3DIndex(_currentMetaIndex, kBits);
                        var fullIndex = metaIndex * kWrappedGridSize + innerIndex;
                        _current = (fullIndex, value);
                        return;
                    }
                }
            }
            _nestedIterator = null;
            _current = null;
        }

        public bool MoveNext()
        {
            if (_nestedIterator != null && _nestedIterator.MoveNext())
            {
                var (innerIndex, value) = _nestedIterator.Current;
                var metaIndex = HybridGridUtils.To3DIndex(_currentMetaIndex, kBits);
                var fullIndex = metaIndex * kWrappedGridSize + innerIndex;
                _current = (fullIndex, value);
                return true;
            }

            AdvanceToValidNestedIterator();
            return _current.HasValue;
        }

        public void Reset()
        {
            _currentMetaIndex = -1;
            _nestedIterator = null;
            _current = null;
            AdvanceToValidNestedIterator();
        }

        public (Array3i Index, TValueType Value) Current => _current!.Value;

        object IEnumerator.Current => Current;

        public void Dispose()
        {
            _nestedIterator?.Dispose();
            _current = null;
            GC.SuppressFinalize(this);
        }
    }
}

/// <summary>
/// A grid consisting of 2x2x2 grids of type 'NestedGrid' initially. Wrapped grids
/// are constructed on first access via 'GetMutableValue()'. If necessary, the grid
/// grows to twice the size in each dimension. The range of indices is (almost)
/// symmetric around the origin, i.e. negative indices are allowed.
/// </summary>
internal class DynamicGrid<TValueType> where TValueType : struct
{
    private const int kWrappedGridSize = 64; // NestedGrid<TValueType>.GridSize = 64
    private int _bits; // Starts at 1 (2x2x2 = 8 meta cells)
    private NestedGrid<TValueType>?[] _metaCells;

    public DynamicGrid()
    {
        _bits = 1;
        _metaCells = new NestedGrid<TValueType>?[8]; // 2^3 = 8
    }

    /// <summary>
    /// Returns the current number of voxels per dimension.
    /// </summary>
    public int GridSize => kWrappedGridSize << _bits;

    /// <summary>
    /// Returns the value stored at 'index'.
    /// </summary>
    public TValueType GetValue(Array3i index)
    {
        var shiftedIndex = index + new Array3i(GridSize >> 1, GridSize >> 1, GridSize >> 1);
        
        // Check bounds using unsigned comparison for performance
        if (shiftedIndex.X < 0 || shiftedIndex.Y < 0 || shiftedIndex.Z < 0 ||
            shiftedIndex.X >= GridSize || shiftedIndex.Y >= GridSize || shiftedIndex.Z >= GridSize)
        {
            return default;
        }

        var metaIndex = GetMetaIndex(shiftedIndex);
        var metaCell = _metaCells[HybridGridUtils.ToFlatIndex(metaIndex, _bits)];
        if (metaCell == null)
        {
            return default;
        }
        var innerIndex = shiftedIndex - metaIndex * kWrappedGridSize;
        return metaCell.GetValue(innerIndex);
    }

    /// <summary>
    /// Returns a reference to the value at 'index' to allow changing it, dynamically
    /// growing the DynamicGrid and constructing new NestedGrids as needed.
    /// </summary>
    public ref TValueType GetMutableValue(Array3i index)
    {
        var shiftedIndex = index + new Array3i(GridSize >> 1, GridSize >> 1, GridSize >> 1);
        
        // Check bounds using unsigned comparison for performance
        if (shiftedIndex.X < 0 || shiftedIndex.Y < 0 || shiftedIndex.Z < 0 ||
            shiftedIndex.X >= GridSize || shiftedIndex.Y >= GridSize || shiftedIndex.Z >= GridSize)
        {
            // SAFEGUARD: Store old bits to detect if Grow() actually increased size
            var oldBits = _bits;
            
            // Grow the grid
            Grow();
            
            // SAFEGUARD: Check if Grow() actually increased size (prevent infinite recursion)
            if (_bits == oldBits)
            {
                throw new InvalidOperationException(
                    $"Cannot grow grid further. Index {index} is out of bounds even after grow attempt. " +
                    $"Current bits={_bits}, GridSize={GridSize}, shiftedIndex={shiftedIndex}");
            }
            
            // SAFEGUARD: Recalculate shiftedIndex after grow and check bounds again
            shiftedIndex = index + new Array3i(GridSize >> 1, GridSize >> 1, GridSize >> 1);
            
            // SAFEGUARD: If still out of bounds after grow, throw exception instead of infinite recursion
            if (shiftedIndex.X < 0 || shiftedIndex.Y < 0 || shiftedIndex.Z < 0 ||
                shiftedIndex.X >= GridSize || shiftedIndex.Y >= GridSize || shiftedIndex.Z >= GridSize)
            {
                throw new ArgumentOutOfRangeException(nameof(index),
                    $"Index {index} is out of bounds even after growing grid to maximum size. " +
                    $"GridSize={GridSize}, shiftedIndex={shiftedIndex}, bits={_bits}");
            }
            
            return ref GetMutableValue(index); // Recursive call after grow (now safe)
        }

        var metaIndex = GetMetaIndex(shiftedIndex);
        var flatIndex = HybridGridUtils.ToFlatIndex(metaIndex, _bits);
        if (_metaCells[flatIndex] == null)
        {
            _metaCells[flatIndex] = new NestedGrid<TValueType>();
        }
        var innerIndex = shiftedIndex - metaIndex * kWrappedGridSize;
        return ref _metaCells[flatIndex]!.GetMutableValue(innerIndex);
    }

    /// <summary>
    /// Iterator for iterating over all values not comparing equal to the default constructed value.
    /// </summary>
    public class Iterator : IEnumerator<(Array3i Index, TValueType Value)>
    {
        private readonly DynamicGrid<TValueType> _grid;
        private readonly int _bits;
        private int _currentMetaIndex;
        private IEnumerator<(Array3i Index, TValueType Value)>? _nestedIterator;
        private (Array3i Index, TValueType Value)? _current;

        public Iterator(DynamicGrid<TValueType> grid)
        {
            _grid = grid;
            _bits = grid._bits;
            _currentMetaIndex = -1;
            AdvanceToValidNestedIterator();
        }

        private void AdvanceToValidNestedIterator()
        {
            while (_currentMetaIndex < _grid._metaCells.Length - 1)
            {
                _currentMetaIndex++;
                if (_currentMetaIndex >= _grid._metaCells.Length)
                {
                    _nestedIterator = null;
                    _current = null;
                    return;
                }

                var metaCell = _grid._metaCells[_currentMetaIndex];
                if (metaCell != null)
                {
                    _nestedIterator = metaCell.GetEnumerator();
                    if (_nestedIterator.MoveNext())
                    {
                        var (innerIndex, value) = _nestedIterator.Current;
                        var metaIndex = HybridGridUtils.To3DIndex(_currentMetaIndex, _bits);
                        var shiftedIndex = metaIndex * kWrappedGridSize + innerIndex;
                        var originalIndex = shiftedIndex - new Array3i(
                            (1 << (_bits - 1)) * kWrappedGridSize,
                            (1 << (_bits - 1)) * kWrappedGridSize,
                            (1 << (_bits - 1)) * kWrappedGridSize);
                        _current = (originalIndex, value);
                        return;
                    }
                }
            }
            _nestedIterator = null;
            _current = null;
        }

        public bool MoveNext()
        {
            if (_nestedIterator != null && _nestedIterator.MoveNext())
            {
                var (innerIndex, value) = _nestedIterator.Current;
                var metaIndex = HybridGridUtils.To3DIndex(_currentMetaIndex, _bits);
                var shiftedIndex = metaIndex * kWrappedGridSize + innerIndex;
                var originalIndex = shiftedIndex - new Array3i(
                    (1 << (_bits - 1)) * kWrappedGridSize,
                    (1 << (_bits - 1)) * kWrappedGridSize,
                    (1 << (_bits - 1)) * kWrappedGridSize);
                _current = (originalIndex, value);
                return true;
            }

            AdvanceToValidNestedIterator();
            return _current.HasValue;
        }

        public void Reset()
        {
            _currentMetaIndex = -1;
            _nestedIterator = null;
            _current = null;
            AdvanceToValidNestedIterator();
        }

        public (Array3i Index, TValueType Value) Current => _current!.Value;

        object IEnumerator.Current => Current;

        public void Dispose()
        {
            _nestedIterator?.Dispose();
            _current = null;
            GC.SuppressFinalize(this);
        }

        /// <summary>
        /// Advances iterator to end (for end() implementation).
        /// </summary>
        public void AdvanceToEnd()
        {
            _currentMetaIndex = _grid._metaCells.Length;
            _nestedIterator = null;
            _current = null;
        }
    }

    /// <summary>
    /// Gets an enumerator for all non-default values.
    /// </summary>
    public Iterator GetEnumerator()
    {
        return new Iterator(this);
    }

    private Array3i GetMetaIndex(Array3i index)
    {
        if (index.X < 0 || index.Y < 0 || index.Z < 0)
        {
            throw new ArgumentOutOfRangeException(nameof(index), $"Index {index} has negative components");
        }
        var metaIndex = index / kWrappedGridSize;
        if (metaIndex.X >= (1 << _bits) || metaIndex.Y >= (1 << _bits) || metaIndex.Z >= (1 << _bits))
        {
            throw new ArgumentOutOfRangeException(nameof(index), $"Meta index {metaIndex} is out of range");
        }
        return metaIndex;
    }

    /// <summary>
    /// Grows this grid by a factor of 2 in each of the 3 dimensions.
    /// </summary>
    private void Grow()
    {
        var newBits = _bits + 1;
        if (newBits > 8)
        {
            throw new InvalidOperationException($"Cannot grow grid beyond bits=8 (current bits={_bits})");
        }

        var newMetaCells = new NestedGrid<TValueType>?[8 * _metaCells.Length];
        
        for (int z = 0; z < (1 << _bits); z++)
        {
            for (int y = 0; y < (1 << _bits); y++)
            {
                for (int x = 0; x < (1 << _bits); x++)
                {
                    var originalMetaIndex = new Array3i(x, y, z);
                    var newMetaIndex = originalMetaIndex + new Array3i(1 << (_bits - 1), 1 << (_bits - 1), 1 << (_bits - 1));
                    var originalFlatIndex = HybridGridUtils.ToFlatIndex(originalMetaIndex, _bits);
                    var newFlatIndex = HybridGridUtils.ToFlatIndex(newMetaIndex, newBits);
                    newMetaCells[newFlatIndex] = _metaCells[originalFlatIndex];
                }
            }
        }

        _metaCells = newMetaCells;
        _bits = newBits;
    }
}

/// <summary>
/// Represents a 3D grid as a wide, shallow tree.
/// This is the base class for HybridGrid and IntensityHybridGrid.
/// </summary>
/// <remarks>
/// Creates a new tree-based grid with voxels having edge length 'resolution'
/// around the origin which becomes the center of the cell at index (0, 0, 0).
/// </remarks>
public class HybridGridBase<TValueType>(double resolution) where TValueType : struct
{
    private readonly DynamicGrid<TValueType> _grid = new();

    /// <summary>
    /// Returns the resolution (edge length of each voxel).
    /// </summary>
    public double Resolution => resolution;

    /// <summary>
    /// Returns the value stored at 'index'.
    /// </summary>
    protected TValueType GetValue(Array3i index)
    {
        return _grid.GetValue(index);
    }

    /// <summary>
    /// Returns a reference to the value at 'index' to allow changing it.
    /// </summary>
    protected ref TValueType GetMutableValue(Array3i index)
    {
        return ref _grid.GetMutableValue(index);
    }

    /// <summary>
    /// Returns the index of the cell containing the 'point'. Indices are integer
    /// vectors identifying cells, for this the coordinates are rounded to the next
    /// multiple of the resolution.
    /// </summary>
    public Array3i GetCellIndex(Vector3 point)
    {
        var index = new Vector3(point.X / resolution, point.Y / resolution, point.Z / resolution);
        return new Array3i(
            (int)System.Math.Round(index.X),
            (int)System.Math.Round(index.Y),
            (int)System.Math.Round(index.Z));
    }

    /// <summary>
    /// Returns one of the octants, (0, 0, 0), (1, 0, 0), ..., (1, 1, 1).
    /// </summary>
    public static Array3i GetOctant(int i)
    {
        if (i < 0 || i >= 8)
        {
            throw new ArgumentOutOfRangeException(nameof(i), $"Octant index {i} must be in range [0, 7]");
        }
        return new Array3i(
            (i & 1) != 0 ? 1 : 0,
            (i & 2) != 0 ? 1 : 0,
            (i & 4) != 0 ? 1 : 0);
    }

    /// <summary>
    /// Returns the center of the cell at 'index'.
    /// </summary>
    public Vector3 GetCenterOfCell(Array3i index)
    {
        return new Vector3(
            index.X * resolution,
            index.Y * resolution,
            index.Z * resolution);
    }

    /// <summary>
    /// Gets an enumerator for all non-default values.
    /// </summary>
    public IEnumerator<(Array3i Index, TValueType Value)> GetEnumerator()
    {
        return _grid.GetEnumerator();
    }
}

/// <summary>
/// A grid containing probability values stored using 15 bits, and an update
/// marker per voxel.
/// Points are expected to be close to the origin. Points far from the origin
/// require the grid to grow dynamically. For centimeter resolution, points
/// can only be tens of meters from the origin.
/// The hard limit of cell indexes is +/- 8192 around the origin.
/// </summary>
public class HybridGrid : HybridGridBase<ushort>
{
    private const ushort kUpdateMarker = (ushort)(1u << 15);
    private readonly List<Array3i> _updateIndices;

    /// <summary>
    /// Creates a new HybridGrid with the specified resolution.
    /// </summary>
    public HybridGrid(double resolution) : base(resolution)
    {
        _updateIndices = [];
    }

    /// <summary>
    /// Creates a HybridGrid from a proto.
    /// </summary>
    public HybridGrid(Models.Mapping.HybridGrid proto) : base(proto.Resolution)
    {
        _updateIndices = [];
        
        if (proto.XIndices == null || proto.YIndices == null || proto.ZIndices == null || proto.Values == null)
        {
            throw new ArgumentException("Proto must have valid indices and values", nameof(proto));
        }

        if (proto.XIndices.Count != proto.Values.Count ||
            proto.YIndices.Count != proto.Values.Count ||
            proto.ZIndices.Count != proto.Values.Count)
        {
            throw new ArgumentException(
                $"Proto indices and values count mismatch: X={proto.XIndices.Count}, Y={proto.YIndices.Count}, Z={proto.ZIndices.Count}, Values={proto.Values.Count}",
                nameof(proto));
        }

        for (int i = 0; i < proto.Values.Count; i++)
        {
            var index = new Array3i(proto.XIndices[i], proto.YIndices[i], proto.ZIndices[i]);
            var probability = ProbabilityValues.ValueToProbability((ushort)proto.Values[i]);
            SetProbability(index, probability);
        }
    }

    /// <summary>
    /// Sets the probability of the cell at 'index' to the given 'probability'.
    /// </summary>
    public void SetProbability(Array3i index, double probability)
    {
        var clampedProbability = ProbabilityValues.ClampProbability(probability);
        var value = ProbabilityValues.ProbabilityToValue(clampedProbability);
        GetMutableValue(index) = value;
    }

    /// <summary>
    /// Finishes the update sequence by removing update markers from all updated cells.
    /// </summary>
    public void FinishUpdate()
    {
        foreach (var index in _updateIndices)
        {
            ref var cell = ref GetMutableValue(index);
            if (cell >= kUpdateMarker)
            {
                cell = (ushort)(cell - kUpdateMarker);
            }
        }
        _updateIndices.Clear();
    }

    /// <summary>
    /// Applies the 'table' (lookup table from ComputeLookupTableToApplyOdds) to the
    /// probability of the cell at 'index' if the cell has not already been updated.
    /// Multiple updates of the same cell will be ignored until FinishUpdate() is called.
    /// Returns true if the cell was updated.
    /// 
    /// If this is the first call to ApplyLookupTable() for the specified cell, its value
    /// will be set to probability corresponding to the table entry.
    /// </summary>
    public bool ApplyLookupTable(Array3i index, List<ushort> table)
    {
        if (table == null || table.Count != kUpdateMarker)
        {
            throw new ArgumentException($"Table must have size {kUpdateMarker}", nameof(table));
        }

        ref var cell = ref GetMutableValue(index);
        if (cell >= kUpdateMarker)
        {
            return false; // Already updated
        }

        _updateIndices.Add(index);
        cell = table[cell];
        return true;
    }

    /// <summary>
    /// Returns the probability of the cell with 'index'.
    /// </summary>
    public double GetProbability(Array3i index)
    {
        return ProbabilityValues.ValueToProbability(GetValue(index));
    }

    /// <summary>
    /// Returns true if the probability at the specified 'index' is known.
    /// </summary>
    public bool IsKnown(Array3i index)
    {
        return GetValue(index) != 0;
    }

    /// <summary>
    /// Converts this HybridGrid to a proto.
    /// </summary>
    public Models.Mapping.HybridGrid ToProto()
    {
        if (_updateIndices.Count > 0)
        {
            throw new InvalidOperationException(
                "Serializing a grid during an update is not supported. Finish the update first.");
        }

        var result = new Models.Mapping.HybridGrid
        {
            Resolution = Resolution,
            XIndices = [],
            YIndices = [],
            ZIndices = [],
            Values = []
        };

        foreach (var (index, value) in this)
        {
            result.XIndices.Add(index.X);
            result.YIndices.Add(index.Y);
            result.ZIndices.Add(index.Z);
            result.Values.Add(value);
        }

        return result;
    }
}

/// <summary>
/// Average intensity data structure for IntensityHybridGrid.
/// </summary>
public struct AverageIntensityData
{
    public double Sum { get; set; }
    public int Count { get; set; }
}

/// <summary>
/// Hybrid grid for storing intensity data (average intensity per voxel).
/// </summary>
/// <remarks>
/// Creates a new IntensityHybridGrid with the specified resolution.
/// </remarks>
public class IntensityHybridGrid(double resolution) : HybridGridBase<AverageIntensityData>(resolution)
{

    /// <summary>
    /// Adds intensity value to the cell at 'index'.
    /// </summary>
    public void AddIntensity(Array3i index, double intensity)
    {
        ref var cell = ref GetMutableValue(index);
        cell.Count += 1;
        cell.Sum += intensity;
    }

    /// <summary>
    /// Returns the average intensity of the cell at 'index'.
    /// </summary>
    public double GetIntensity(Array3i index)
    {
        var cell = GetValue(index);
        if (cell.Count == 0)
        {
            return 0.0;
        }
        return cell.Sum / cell.Count;
    }
}