/*
 * Copyright 2018 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.Mapping.Internal.D2D;
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 truncated signed distances and weights.
/// </summary>
public class TSDF2D : Grid2D
{
    private readonly ValueConversionTables _conversionTables;
    private readonly TSDValueConverter _valueConverter;
    private List<ushort> _weightCells;

    public TSDF2D(MapLimits limits, double truncationDistance, double maxWeight,
                   ValueConversionTables conversionTables)
        : base(limits, -truncationDistance, truncationDistance, conversionTables)
    {
        _conversionTables = conversionTables;
        _valueConverter = new TSDValueConverter(truncationDistance, maxWeight, conversionTables);
        
        var size = limits.CellLimits.NumXCells * limits.CellLimits.NumYCells;
        _weightCells = new List<ushort>(size);
        for (int i = 0; i < size; i++)
        {
            _weightCells.Add(TSDValueConverter.GetUnknownWeightValue());
        }
    }

    public TSDF2D(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 : -0.3,
               proto.MaxCorrespondenceCost > 0 ? proto.MaxCorrespondenceCost : 0.3,
               conversionTables)
    {
        if (proto.Tsdf2D == null)
        {
            throw new ArgumentException("Proto must have TSDF2D data", nameof(proto));
        }

        _conversionTables = conversionTables;
        var tsdfProto = proto.Tsdf2D.Value;
        _valueConverter = new TSDValueConverter(
            tsdfProto.TruncationDistance, tsdfProto.MaxWeight, conversionTables);

        // Match C++: Grid2D(proto, ...) copies cells from proto.cells() into correspondence_cost_cells_
        // Since we call base constructor with new MapLimits (not proto), we must copy manually.
        // Without this, all TSD values remain 0 (unknown) after deserialization.
        if (proto.Cells != null)
        {
            _correspondenceCostCells.Clear();
            foreach (var cell in proto.Cells)
            {
                _correspondenceCostCells.Add((ushort)cell);
            }
        }

        _weightCells = new List<ushort>(tsdfProto.WeightCells.Count);
        foreach (var cell in tsdfProto.WeightCells)
        {
            if (cell > ushort.MaxValue)
            {
                throw new ArgumentException("Weight cell value exceeds ushort.MaxValue");
            }
            _weightCells.Add((ushort)cell);
        }

        // Match C++: Grid2D constructor copies known_cells_box from proto if present
        // Since we call base constructor with new MapLimits, we need to copy known_cells_box manually
        // C++ Grid2D constructor: if (proto.has_known_cells_box()) { known_cells_box_ = ... }
        // Check if known_cells_box is not empty (minX > maxX indicates empty in C++ AlignedBox2i)
        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 TSD and weight of the cell at 'cell_index'.
    /// Only allowed if the cell was not updated in the current update cycle.
    /// </summary>
    public void SetCell(Array2i cellIndex, double tsd, double weight)
    {
        var flatIndex = ToFlatIndex(cellIndex);
        var tsdfCell = _correspondenceCostCells[flatIndex];
        
        ushort kUpdateMarker = TSDValueConverter.GetUpdateMarker();
        if (tsdfCell >= kUpdateMarker)
        {
            return; // Already updated
        }

        _updateIndices.Add(flatIndex);
        UpdateKnownCellsBox(cellIndex);

        // Set TSD with update marker
        _correspondenceCostCells[flatIndex] = (ushort)(_valueConverter.TSDToValue(tsd) + kUpdateMarker);
        
        // Set weight
        _weightCells[flatIndex] = _valueConverter.WeightToValue(weight);
    }

    /// <summary>
    /// Returns the TSD of the cell with 'cell_index'.
    /// </summary>
    public double GetTSD(Array2i cellIndex)
    {
        if (_limits.Contains(cellIndex))
        {
            return _valueConverter.ValueToTSD(_correspondenceCostCells[ToFlatIndex(cellIndex)]);
        }
        return _valueConverter.GetMinTSD();
    }

    /// <summary>
    /// Returns the weight of the cell with 'cell_index'.
    /// </summary>
    public double GetWeight(Array2i cellIndex)
    {
        if (_limits.Contains(cellIndex))
        {
            var flatIndex = ToFlatIndex(cellIndex);
            return _valueConverter.ValueToWeight(_weightCells[flatIndex]);
        }
        return _valueConverter.GetMinWeight();
    }

    /// <summary>
    /// Returns the TSD and weight of the cell with 'cell_index'.
    /// </summary>
    public (double tsd, double weight) GetTSDAndWeight(Array2i cellIndex)
    {
        if (_limits.Contains(cellIndex))
        {
            var flatIndex = ToFlatIndex(cellIndex);
            return (_valueConverter.ValueToTSD(_correspondenceCostCells[flatIndex]),
                    _valueConverter.ValueToWeight(_weightCells[flatIndex]));
        }
        return (_valueConverter.GetMinTSD(), _valueConverter.GetMinWeight());
    }

    /// <summary>
    /// Returns true if the cell at 'cell_index' was updated in the current update cycle.
    /// </summary>
    public bool CellIsUpdated(Array2i cellIndex)
    {
        var flatIndex = ToFlatIndex(cellIndex);
        var tsdfCell = _correspondenceCostCells[flatIndex];
        return tsdfCell >= TSDValueConverter.GetUpdateMarker();
    }

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

    /// <summary>
    /// Grows the map as necessary to include 'point'.
    /// </summary>
    public override void GrowLimits(Vector2 point)
    {
        GrowLimits(point, [_correspondenceCostCells, _weightCells],
                   [TSDValueConverter.GetUnknownTSDValue(), TSDValueConverter.GetUnknownWeightValue()]);
    }

    protected override void UpdateGridReferences(List<ushort>[] grids)
    {
        base.UpdateGridReferences(grids);
        _weightCells = grids[1];
    }

    /// <summary>
    /// Converts to proto representation.
    /// </summary>
    public override Models.Mapping.Grid2D ToProto()
    {
        var proto = base.ToProto();
        
        var weightCells = new List<int>(_weightCells.Count);
        foreach (var cell in _weightCells)
        {
            weightCells.Add(cell);
        }

        proto.Tsdf2D = new Models.Mapping.TSDF2D(
            _valueConverter.GetMaxTSD(),
            _valueConverter.GetMaxWeight(),
            weightCells);

        return proto;
    }

    /// <summary>
    /// Computes a cropped grid containing only known cells.
    /// </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 TSDF2D(
            new MapLimits(resolution, max, cellLimits),
            _valueConverter.GetMaxTSD(),
            _valueConverter.GetMaxWeight(),
            _conversionTables);

        // Copy known cells
        for (int y = 0; y < cellLimits.NumYCells; y++)
        {
            for (int x = 0; x < cellLimits.NumXCells; x++)
            {
                var oldIndex = new Array2i(offset.X + x, offset.Y + y);
                if (_limits.Contains(oldIndex) && IsKnown(oldIndex))
                {
                    var newIndex = new Array2i(x, y);
                    var (tsd, weight) = GetTSDAndWeight(oldIndex);
                    croppedGrid.SetCell(newIndex, tsd, weight);
                }
            }
        }

        croppedGrid.FinishUpdate();
        return croppedGrid;
    }

    /// <summary>
    /// Draws the TSDF grid to a submap texture for visualization.
    /// Match C++: TSDF2D::DrawToSubmapTexture
    ///
    /// TSDF convention:
    /// - tsd > 0: Free space (away from obstacles)
    /// - tsd < 0: Occupied space (inside/near obstacles)
    /// - tsd = 0: On obstacle surface
    ///
    /// Output texture convention (matching ProbabilityGrid):
    /// - delta > 0 → value high, alpha 0 → bright pixel → FREE
    /// - delta < 0 → value 0, alpha high → dark pixel → OCCUPIED
    /// </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);

        var maxTsd = _valueConverter.GetMaxTSD();
        var maxWeight = _valueConverter.GetMaxWeight();

        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
                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 tsd = GetTSD(sourceIndex);
            var weight = GetWeight(sourceIndex);
            var normalizedWeight = weight / maxWeight;

            // FIXED: Keep the sign of TSD to distinguish free vs occupied space
            // tsd > 0 (free) → normalizedTsd > 0 → delta > 0 → value high (bright)
            // tsd < 0 (occupied) → normalizedTsd < 0 → delta < 0 → alpha high (dark)
            //
            // Normalize TSD to [-1, 1] range while preserving sign
            var normalizedTsd = Math.Clamp(tsd / maxTsd, -1.0, 1.0);

            // Apply sqrt scaling to magnitude only, preserve sign for better visualization
            // This makes the gradient more visible near the surface
            var magnitude = Math.Pow(Math.Abs(normalizedTsd), 0.5);
            var signedMagnitude = normalizedTsd >= 0 ? magnitude : -magnitude;

            // Scale by weight and convert to delta
            // delta range: [-127, 127] scaled by weight
            var delta = (int)Math.Round(normalizedWeight * signedMagnitude * 127.0,
                                        MidpointRounding.AwayFromZero);

            byte alpha = (byte)(delta < 0 ? Math.Min(255, -delta) : 0);
            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++: AlignedBox2i::isEmpty() logic: min > max indicates empty
    /// </summary>
    private void UpdateKnownCellsBox(Array2i cellIndex)
    {
        // Match C++ AlignedBox2i::isEmpty() logic: min > max indicates empty
        if (_knownCellsBox.minX > _knownCellsBox.maxX || _knownCellsBox.minY > _knownCellsBox.maxY)
        {
            _knownCellsBox = (cellIndex.X, cellIndex.Y, cellIndex.X, cellIndex.Y);
        }
        else
        {
            _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)
            );
        }
    }
}