/*
 * 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.
 */

namespace CartographerSharp.Mapping.Internal;

/// <summary>
/// A class that tracks the connectivity structure between trajectories.
///
/// Connectivity includes both the count ("How many times have I _directly_
/// connected trajectories i and j?") and the transitive connectivity.
///
/// Uses Union-Find (disjoint set forest) algorithm for efficient transitive
/// connectivity tracking.
///
/// Match C++ ConnectedComponents (connected_components.cc)
/// </summary>
public class ConnectedComponents
{
    private readonly object _lock = new();

    // Tracks transitive connectivity using a disjoint set forest, i.e. each
    // entry points towards the representative for the given trajectory.
    private readonly Dictionary<int, int> _forest = new();

    // Tracks the number of direct connections between a pair of trajectories.
    private readonly Dictionary<(int, int), int> _connectionMap = new();

    /// <summary>
    /// Add a trajectory which is initially connected to only itself.
    /// </summary>
    public void Add(int trajectoryId)
    {
        lock (_lock)
        {
            // Use TryAdd to avoid overwriting existing entries
            _forest.TryAdd(trajectoryId, trajectoryId);
        }
    }

    /// <summary>
    /// Connect two trajectories. If either trajectory is untracked, it will be
    /// tracked. This function is invariant to the order of its arguments. Repeated
    /// calls to Connect increment the connectivity count.
    /// </summary>
    public void Connect(int trajectoryIdA, int trajectoryIdB)
    {
        lock (_lock)
        {
            Union(trajectoryIdA, trajectoryIdB);
            var sortedPair = (Math.Min(trajectoryIdA, trajectoryIdB), Math.Max(trajectoryIdA, trajectoryIdB));
            if (!_connectionMap.TryGetValue(sortedPair, out var count))
            {
                count = 0;
            }
            _connectionMap[sortedPair] = count + 1;
        }
    }

    /// <summary>
    /// Determines if two trajectories have been (transitively) connected. If
    /// either trajectory is not being tracked, returns false, except when it is
    /// the same trajectory, where it returns true. This function is invariant to
    /// the order of its arguments.
    /// </summary>
    public bool TransitivelyConnected(int trajectoryIdA, int trajectoryIdB)
    {
        if (trajectoryIdA == trajectoryIdB)
        {
            return true;
        }

        lock (_lock)
        {
            if (!_forest.ContainsKey(trajectoryIdA) || !_forest.ContainsKey(trajectoryIdB))
            {
                return false;
            }
            return FindSet(trajectoryIdA) == FindSet(trajectoryIdB);
        }
    }

    /// <summary>
    /// Return the number of _direct_ connections between 'trajectoryIdA' and
    /// 'trajectoryIdB'. If either trajectory is not being tracked, returns 0.
    /// This function is invariant to the order of its arguments.
    /// </summary>
    public int ConnectionCount(int trajectoryIdA, int trajectoryIdB)
    {
        lock (_lock)
        {
            var sortedPair = (Math.Min(trajectoryIdA, trajectoryIdB), Math.Max(trajectoryIdA, trajectoryIdB));
            return _connectionMap.TryGetValue(sortedPair, out var count) ? count : 0;
        }
    }

    /// <summary>
    /// The trajectory IDs, grouped by connectivity.
    /// </summary>
    public List<List<int>> Components()
    {
        lock (_lock)
        {
            // Map from cluster exemplar -> growing cluster
            var map = new Dictionary<int, List<int>>();
            foreach (var entry in _forest)
            {
                var representative = FindSet(entry.Key);
                if (!map.TryGetValue(representative, out var component))
                {
                    component = [];
                    map[representative] = component;
                }
                component.Add(entry.Key);
            }

            return [.. map.Values];
        }
    }

    /// <summary>
    /// The list of trajectory IDs that belong to the same connected component as
    /// 'trajectoryId'.
    /// </summary>
    public List<int> GetComponent(int trajectoryId)
    {
        lock (_lock)
        {
            if (!_forest.ContainsKey(trajectoryId))
            {
                return [trajectoryId];
            }

            var setId = FindSet(trajectoryId);
            var trajectoryIds = new List<int>();
            foreach (var entry in _forest)
            {
                if (FindSet(entry.Key) == setId)
                {
                    trajectoryIds.Add(entry.Key);
                }
            }
            return trajectoryIds;
        }
    }

    /// <summary>
    /// Find the representative and compresses the path to it.
    /// Must be called with lock held.
    /// </summary>
    private int FindSet(int trajectoryId)
    {
        if (!_forest.TryGetValue(trajectoryId, out var parent))
        {
            return trajectoryId;
        }

        if (trajectoryId != parent)
        {
            // Path compression for efficiency
            _forest[trajectoryId] = FindSet(parent);
        }
        return _forest[trajectoryId];
    }

    /// <summary>
    /// Union two sets.
    /// Must be called with lock held.
    /// </summary>
    private void Union(int trajectoryIdA, int trajectoryIdB)
    {
        // Add trajectories if not already tracked
        _forest.TryAdd(trajectoryIdA, trajectoryIdA);
        _forest.TryAdd(trajectoryIdB, trajectoryIdB);

        var representativeA = FindSet(trajectoryIdA);
        var representativeB = FindSet(trajectoryIdB);
        _forest[representativeA] = representativeB;
    }
}