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

namespace CartographerSharp.Mapping.D2D;

/// <summary>
/// Represents a 2D grid of probabilities.
/// </summary>
public class ProbabilityGrid : Grid2D
{
    private readonly ValueConversionTables _conversionTables;

    public ProbabilityGrid(MapLimits limits, ValueConversionTables conversionTables)
        : base(limits, ProbabilityValues.kMinCorrespondenceCost, ProbabilityValues.kMaxCorrespondenceCost, conversionTables)
    {
        _conversionTables = conversionTables;
    }

    public ProbabilityGrid(Models.Mapping.Grid2D proto, ValueConversionTables conversionTables)
        : base(new MapLimits(proto.Limits.Resolution,
                             new Vector2(proto.Limits.Max.X, proto.Limits.Max.Y),
                             new CellLimits(proto.Limits.CellLimits.NumXCells, proto.Limits.CellLimits.NumYCells)),
               proto.MinCorrespondenceCost > 0 ? proto.MinCorrespondenceCost : ProbabilityValues.kMinCorrespondenceCost,
               proto.MaxCorrespondenceCost > 0 ? proto.MaxCorrespondenceCost : ProbabilityValues.kMaxCorrespondenceCost,
               conversionTables)
    {
        _conversionTables = conversionTables;
        
        // Copy cells from proto
        if (proto.Cells != null)
        {
            _correspondenceCostCells.Clear();
            foreach (var cell in proto.Cells)
            {
                _correspondenceCostCells.Add((ushort)cell);
            }
        }
        
        // Copy known cells box
        // Match C++: proto.has_known_cells_box() - use MinX <= MaxX to detect valid box,
        // which correctly handles boxes at origin (0,0)-(0,0).
        if (proto.KnownCellsBox.MinX <= proto.KnownCellsBox.MaxX &&
            proto.KnownCellsBox.MinY <= proto.KnownCellsBox.MaxY)
        {
            _knownCellsBox = (proto.KnownCellsBox.MinX, proto.KnownCellsBox.MinY,
                             proto.KnownCellsBox.MaxX, proto.KnownCellsBox.MaxY);
        }
    }

    /// <summary>
    /// Sets the probability of the cell at 'cell_index' to the given
    /// 'probability'. Only allowed if the cell was unknown before.
    /// </summary>
    public void SetProbability(Array2i cellIndex, double probability)
    {
        var flatIndex = ToFlatIndex(cellIndex);
        var cell = _correspondenceCostCells[flatIndex];
        
        const ushort kUnknownProbabilityValue = 0;
        if (cell != kUnknownProbabilityValue)
        {
            throw new InvalidOperationException("Cell must be unknown before setting probability");
        }

        _correspondenceCostCells[flatIndex] = ProbabilityValues.CorrespondenceCostToValue(
            ProbabilityValues.ProbabilityToCorrespondenceCost(probability));
        
        // Update known cells box
        UpdateKnownCellsBox(cellIndex);
    }

    /// <summary>
    /// Applies the 'odds' specified when calling ComputeLookupTableToApplyOdds()
    /// to the probability of the cell at 'cell_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.
    /// </summary>
    public bool ApplyLookupTable(Array2i cellIndex, List<ushort> table)
    {
        const ushort kUpdateMarker = (ushort)(1u << 15);
        const int kValueCount = 32768;
        
        if (table.Count != kValueCount)
        {
            throw new ArgumentException($"Table size must be {kValueCount}", nameof(table));
        }

        var flatIndex = ToFlatIndex(cellIndex);
        var cell = _correspondenceCostCells[flatIndex];
        
        if (cell >= kUpdateMarker)
        {
            return false; // Already updated
        }

        _updateIndices.Add(flatIndex);
        _correspondenceCostCells[flatIndex] = table[cell];
        // After applying lookup table, the cell value should have the update marker set
        // (value >= kUpdateMarker), which will be removed in FinishUpdate()
        UpdateKnownCellsBox(cellIndex);
        
        return true;
    }

    /// <summary>
    /// Gets the grid type.
    /// </summary>
    public override GridType GetGridType()
    {
        return GridType.ProbabilityGrid;
    }

    /// <summary>
    /// Returns the probability of the cell with 'cell_index'.
    /// </summary>
    public double GetProbability(Array2i cellIndex)
    {
        if (!_limits.Contains(cellIndex))
        {
            return ProbabilityValues.kMinProbability;
        }
        var flatIndex = ToFlatIndex(cellIndex);
        var value = _correspondenceCostCells[flatIndex];
        return ProbabilityValues.CorrespondenceCostToProbability(
            ProbabilityValues.ValueToCorrespondenceCost(value));
    }

    /// <summary>
    /// Converts to proto representation.
    /// </summary>
    public override Models.Mapping.Grid2D ToProto()
    {
        var proto = base.ToProto();
        proto.ProbabilityGrid2D = new Models.Mapping.ProbabilityGrid(); // Empty struct to indicate type
        return proto;
    }

    /// <summary>
    /// Computes a cropped grid containing only known cells.
    /// Match C++ implementation: only copy known cells using SetProbability.
    /// </summary>
    public override Grid2D ComputeCroppedGrid()
    {
        ComputeCroppedLimits(out var offset, out var cellLimits);
        var resolution = _limits.Resolution;
        var max = new Vector2(
            (_limits.Max.X - resolution * offset.Y),
            (_limits.Max.Y - resolution * offset.X));
        
        var croppedGrid = new ProbabilityGrid(
            new MapLimits(resolution, max, cellLimits),
            _conversionTables);

        // Match C++: for (const Eigen::Array2i& xy_index : XYIndexRangeIterator(cell_limits)) {
        //              if (!IsKnown(xy_index + offset)) continue;
        //              cropped_grid->SetProbability(xy_index, GetProbability(xy_index + offset));
        //            }
        // Only copy known cells using SetProbability (which updates known_cells_box)
        for (int y = 0; y < cellLimits.NumYCells; y++)
        {
            for (int x = 0; x < cellLimits.NumXCells; x++)
            {
                var xyIndex = new Array2i(x, y);
                var oldIndex = new Array2i(offset.X + x, offset.Y + y);
                if (!IsKnown(oldIndex))
                {
                    continue; // Skip unknown cells
                }
                croppedGrid.SetProbability(xyIndex, GetProbability(oldIndex));
            }
        }

        return croppedGrid;
    }

    /// <summary>
    /// Draws the probability grid to a submap texture for visualization.
    /// Match C++: ProbabilityGrid::DrawToSubmapTexture
    /// </summary>
    /// <param name="localPose">The local pose of the submap.</param>
    /// <returns>A texture containing the visualization data.</returns>
    public SubmapQuery.Texture DrawToSubmapTexture(Rigid3d localPose)
    {
        ComputeCroppedLimits(out var offset, out var cellLimits);

        // Build the cells data (value + alpha pairs)
        var cellsData = new List<byte>(cellLimits.NumXCells * cellLimits.NumYCells * 2);

        foreach (var xyIndex in new XYIndexRange(cellLimits))
        {
            var sourceIndex = new Array2i(xyIndex.X + offset.X, xyIndex.Y + offset.Y);

            if (!IsKnown(sourceIndex))
            {
                cellsData.Add(0); // value (unknown log odds value)
                cellsData.Add(0); // alpha
                continue;
            }

            // We would like to add 'delta' but this is not possible using a value and
            // alpha. We use premultiplied alpha, so when 'delta' is positive we can
            // add it by setting 'alpha' to zero. If it is negative, we set 'value' to
            // zero, and use 'alpha' to subtract. This is only correct when the pixel
            // is currently white, so walls will look too gray. This should be hard to
            // detect visually for the user, though.
            var probability = GetProbability(sourceIndex);
            var delta = 128 - SubmapProbabilityUtils.ProbabilityToLogOddsInteger(probability);

            byte alpha = (byte)(delta > 0 ? 0 : Math.Min(255, -delta));
            byte value = (byte)(delta > 0 ? Math.Min(255, delta) : 0);

            cellsData.Add(value);
            cellsData.Add((value != 0 || alpha != 0) ? alpha : (byte)1);
        }

        // Compress using GZip
        var compressedCells = CompressGzip(cellsData.ToArray());

        // Calculate slice pose
        var resolution = _limits.Resolution;
        var maxX = _limits.Max.X - resolution * offset.Y;
        var maxY = _limits.Max.Y - resolution * offset.X;

        var slicePose = localPose.Inverse() *
                        Rigid3d.FromTranslation(new Vector3(maxX, maxY, 0));

        return new SubmapQuery.Texture(
            compressedCells,
            cellLimits.NumXCells,
            cellLimits.NumYCells,
            resolution,
            slicePose);
    }

    /// <summary>
    /// Compresses data using GZip.
    /// Match C++: common::FastGzipString
    /// </summary>
    private static List<byte> CompressGzip(byte[] data)
    {
        using var memoryStream = new MemoryStream();
        using (var gzipStream = new GZipStream(memoryStream, CompressionMode.Compress, leaveOpen: true))
        {
            gzipStream.Write(data, 0, data.Length);
        }
        return new List<byte>(memoryStream.ToArray());
    }

    /// <summary>
    /// Updates the known cells box to include the given cell index.
    /// Match C++: mutable_known_cells_box()->extend(cell_index.matrix())
    /// </summary>
    private void UpdateKnownCellsBox(Array2i cellIndex)
    {
        // Match C++: AlignedBox2i::extend() - if empty, sets min=max=point, otherwise extends bounds
        // Empty box is marked by minX > maxX (or minY > maxY)
        if (_knownCellsBox.minX > _knownCellsBox.maxX || _knownCellsBox.minY > _knownCellsBox.maxY)
        {
            // Box is empty, set min and max to cellIndex
            _knownCellsBox = (cellIndex.X, cellIndex.Y, cellIndex.X, cellIndex.Y);
        }
        else
        {
            // Box is not empty, extend bounds
            _knownCellsBox = (
                Math.Min(_knownCellsBox.minX, cellIndex.X),
                Math.Min(_knownCellsBox.minY, cellIndex.Y),
                Math.Max(_knownCellsBox.maxX, cellIndex.X),
                Math.Max(_knownCellsBox.maxY, cellIndex.Y)
            );
        }
    }
}