/*
 * 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.Mapping.D2D;
using CartographerSharp.Transform;
using  RobotNet10.Shared.Numbers;

namespace CartographerSharp.Mapping.Internal.D2D;

/// <summary>
/// Trims submaps from the pose graph based on overlap area.
/// Removes older submaps that overlap significantly with newer ones,
/// keeping only the freshest submaps while ensuring minimum coverage.
/// </summary>
public class OverlappingSubmapsTrimmer2D : PoseGraphTrimmer
{
    private readonly int _freshSubmapsCount;
    private readonly double _minCoveredArea;
    private readonly int _minAddedSubmapsCount;

    // Current finished submap count (matches C++ current_submap_count_)
    private int _currentSubmapCount = 0;

    public OverlappingSubmapsTrimmer2D(
        int freshSubmapsCount,
        double minCoveredArea,
        int minAddedSubmapsCount)
    {
        if (freshSubmapsCount < 0)
            throw new ArgumentException("freshSubmapsCount must be non-negative", nameof(freshSubmapsCount));
        if (minCoveredArea < 0)
            throw new ArgumentException("minCoveredArea must be non-negative", nameof(minCoveredArea));
        if (minAddedSubmapsCount < 0)
            throw new ArgumentException("minAddedSubmapsCount must be non-negative", nameof(minAddedSubmapsCount));

        _freshSubmapsCount = freshSubmapsCount;
        _minCoveredArea = minCoveredArea;
        _minAddedSubmapsCount = minAddedSubmapsCount;
    }

    public override void Trim(ITrimmable trimmable)
    {
        var submapData = trimmable.GetOptimizedSubmapData();
        
        // Match C++: if (submap_data.size() - current_submap_count_ <= min_added_submaps_count_)
        if (submapData.Count - _currentSubmapCount <= _minAddedSubmapsCount)
        {
            return;
        }

        // Get first submap's map limits to initialize coverage grid
        if (submapData.Count == 0)
        {
            return;
        }

        var firstSubmapData = submapData.First();
        if (firstSubmapData.Data.Submap is not Mapping.D2D.Submap2D firstSubmap || firstSubmap.Grid == null)
        {
            return;
        }

        var firstSubmapMapLimits = firstSubmap.Grid.Limits;
        var coverageGrid = new SubmapCoverageGrid2D(firstSubmapMapLimits);

        // Compute submap freshness from intra-submap constraints
        var submapFreshness = ComputeSubmapFreshness(
            submapData,
            trimmable.GetTrajectoryNodes(),
            trimmable.GetConstraints());

        // Add all submaps to coverage grid
        var allSubmapIds = AddSubmapsToSubmapCoverageGrid2D(
            submapFreshness,
            submapData,
            coverageGrid);

        // Find submaps to trim
        // Match C++: min_covered_area_ / common::Pow2(coverage_grid.resolution())
        var minCoveredCellsCount = (int)Math.Round(_minCoveredArea / MathUtils.Pow2(coverageGrid.Resolution));
        var submapIdsToRemove = FindSubmapIdsToTrim(
            coverageGrid,
            allSubmapIds,
            _freshSubmapsCount,
            minCoveredCellsCount);

        // Update current submap count (matches C++: current_submap_count_ = submap_data.size() - submap_ids_to_remove.size())
        _currentSubmapCount = submapData.Count - submapIdsToRemove.Count;

        // Trim the submaps
        foreach (var id in submapIdsToRemove)
        {
            trimmable.TrimSubmap(id);
        }
    }

    /// <summary>
    /// Tracks which submaps cover which cells in a global coordinate system.
    /// </summary>
    private class SubmapCoverageGrid2D(Mapping.D2D.MapLimits mapLimits)
    {
        // Aliases for documentation only (no type-safety).
        public record CellId(long X, long Y);

        public record StoredType(SubmapId SubmapId, long Time);

        private readonly Vector2 _offset = mapLimits.Max;
        private readonly double _resolution = mapLimits.Resolution;
        private readonly Dictionary<CellId, List<StoredType>> _cells = [];

        public void AddPoint(Vector2 point, SubmapId submapId, long time)
        {
            var cellId = new CellId(
                (long)Math.Round((_offset.X - point.X) / _resolution, MidpointRounding.AwayFromZero),
                (long)Math.Round((_offset.Y - point.Y) / _resolution, MidpointRounding.AwayFromZero));
            
            if (!_cells.TryGetValue(cellId, out var storedTypes))
            {
                storedTypes = [];
                _cells[cellId] = storedTypes;
            }
            storedTypes.Add(new StoredType(submapId, time));
        }

        public Dictionary<CellId, List<StoredType>> Cells => _cells;
        public double Resolution => _resolution;
    }

    /// <summary>
    /// Uses intra-submap constraints and trajectory node timestamps to identify time
    /// of the last range data insertion to the submap.
    /// </summary>
    private static Dictionary<SubmapId, long> ComputeSubmapFreshness(
        MapById<SubmapId, IPoseGraph.SubmapData> submapData,
        MapById<NodeId, TrajectoryNode> trajectoryNodes,
        List<IPoseGraph.Constraint> constraints)
    {
        var submapFreshness = new Dictionary<SubmapId, long>();

        // Find the node with the largest NodeId per SubmapId.
        var submapToLatestNode = new Dictionary<SubmapId, NodeId>();
        foreach (var constraint in constraints)
        {
            if (constraint.ConstraintTag != IPoseGraph.Constraint.Tag.IntraSubmap)
            {
                continue;
            }

            if (!submapToLatestNode.TryGetValue(constraint.SubmapId, out var existingNodeId))
            {
                submapToLatestNode[constraint.SubmapId] = constraint.NodeId;
                continue;
            }

            // Keep the maximum NodeId (matches C++: std::max)
            if (CompareNodeIds(constraint.NodeId, existingNodeId) > 0)
            {
                submapToLatestNode[constraint.SubmapId] = constraint.NodeId;
            }
        }

        // Find timestamp of every latest node.
        foreach (var (submapId, nodeId) in submapToLatestNode)
        {
            if (!submapData.Contains(submapId))
            {
                // Log warning equivalent (C++: LOG(WARNING))
                continue;
            }

            if (!trajectoryNodes.Contains(nodeId))
            {
                continue;
            }
            var trajectoryNode = trajectoryNodes[nodeId];

            if (trajectoryNode.ConstantData == null)
            {
                continue;
            }

            submapFreshness[submapId] = trajectoryNode.ConstantData.Time;
        }

        return submapFreshness;
    }

    /// <summary>
    /// Compares two NodeIds. Returns positive if lhs > rhs, negative if lhs < rhs, 0 if equal.
    /// </summary>
    private static int CompareNodeIds(NodeId lhs, NodeId rhs)
    {
        var trajectoryCompare = lhs.TrajectoryId.CompareTo(rhs.TrajectoryId);
        if (trajectoryCompare != 0)
        {
            return trajectoryCompare;
        }
        return lhs.NodeIndex.CompareTo(rhs.NodeIndex);
    }

    /// <summary>
    /// Iterates over every cell in a submap, transforms the center of the cell to
    /// the global frame and then adds the submap id and the timestamp of the most
    /// recent range data insertion into the global grid.
    /// </summary>
    private static HashSet<SubmapId> AddSubmapsToSubmapCoverageGrid2D(
        Dictionary<SubmapId, long> submapFreshness,
        MapById<SubmapId, IPoseGraph.SubmapData> submapData,
        SubmapCoverageGrid2D coverageGrid)
    {
        var allSubmapIds = new HashSet<SubmapId>();

        foreach (var submap in submapData)
        {
            if (!submapFreshness.TryGetValue(submap.Id, out var freshness))
            {
                continue;
            }

            if (submap.Data.Submap is not Mapping.D2D.Submap2D submap2D || !submap2D.InsertionFinished)
            {
                continue;
            }

            if (submap2D.Grid == null)
            {
                continue;
            }

            allSubmapIds.Add(submap.Id);

            var grid = submap2D.Grid;
            
            // Iterate over every cell in a submap.
            grid.ComputeCroppedLimits(out var offset, out var cellLimits);
            if (cellLimits.NumXCells == 0 || cellLimits.NumYCells == 0)
            {
                // Log warning equivalent (C++: LOG(WARNING))
                continue;
            }

            var globalFrameFromSubmapFrame = submap.Data.Pose;
            var submapFrameFromLocalFrame = submap2D.LocalPose.Inverse();

            foreach (var xyIndex in new XYIndexRange(cellLimits))
            {
                var index = xyIndex + offset;
                if (!grid.IsKnown(index))
                {
                    continue;
                }

                // Match C++: center_of_cell_in_local_frame calculation
                // C++: grid.limits().max().x() - grid.limits().resolution() * (index.y() + 0.5)
                // C++: grid.limits().max().y() - grid.limits().resolution() * (index.x() + 0.5)
                var centerOfCellInLocalFrame = new Rigid3d(
                    new Vector3(
                        (grid.Limits.Max.X - grid.Limits.Resolution * (index.Y + 0.5)),
                        (grid.Limits.Max.Y - grid.Limits.Resolution * (index.X + 0.5)),
                        0.0),
                    Quaternion.Identity);

                // Match C++: transform::Project2D(global_frame_from_submap_frame * submap_frame_from_local_frame * center_of_cell_in_local_frame)
                var centerOfCellInGlobalFrame = TransformOperations.Project2D(
                    globalFrameFromSubmapFrame *
                    submapFrameFromLocalFrame *
                    centerOfCellInLocalFrame);

                coverageGrid.AddPoint(
                    centerOfCellInGlobalFrame.Translation,
                    submap.Id,
                    freshness);
            }
        }

        return allSubmapIds;
    }

    /// <summary>
    /// Returns IDs of submaps that have less than 'min_covered_cells_count' cells
    /// not overlapped by at least 'fresh_submaps_count' submaps.
    /// </summary>
    private static List<SubmapId> FindSubmapIdsToTrim(
        SubmapCoverageGrid2D coverageGrid,
        HashSet<SubmapId> allSubmapIds,
        int freshSubmapsCount,
        int minCoveredCellsCount)
    {
        var submapToCoveredCellsCount = new Dictionary<SubmapId, int>();

        foreach (var (cellId, storedTypes) in coverageGrid.Cells)
        {
            var submapsPerCell = new List<(SubmapId SubmapId, long Time)>();
            foreach (var storedType in storedTypes)
            {
                submapsPerCell.Add((storedType.SubmapId, storedType.Time));
            }

            // In case there are several submaps covering the cell, only the freshest
            // submaps are kept.
            if (submapsPerCell.Count > freshSubmapsCount)
            {
                // Sort by time in descending order (matches C++: std::sort with > comparison)
                submapsPerCell.Sort((left, right) => right.Time.CompareTo(left.Time));
                submapsPerCell = [.. submapsPerCell.Take(freshSubmapsCount)];
            }

            foreach (var (submapId, _) in submapsPerCell)
            {
                if (!submapToCoveredCellsCount.TryGetValue(submapId, out var count))
                {
                    count = 0;
                }
                submapToCoveredCellsCount[submapId] = count + 1;
            }
        }

        var submapIdsToKeep = new List<SubmapId>();
        foreach (var (submapId, cellsCount) in submapToCoveredCellsCount)
        {
            if (cellsCount < minCoveredCellsCount)
            {
                continue;
            }
            submapIdsToKeep.Add(submapId);
        }

        // Match C++: std::set_difference(all_submap_ids, submap_ids_to_keep)
        submapIdsToKeep.Sort((a, b) =>
        {
            var trajectoryCompare = a.TrajectoryId.CompareTo(b.TrajectoryId);
            if (trajectoryCompare != 0)
            {
                return trajectoryCompare;
            }
            return a.SubmapIndex.CompareTo(b.SubmapIndex);
        });

        var result = new List<SubmapId>();
        foreach (var submapId in allSubmapIds)
        {
            if (!submapIdsToKeep.Contains(submapId))
            {
                result.Add(submapId);
            }
        }

        return result;
    }
}