/*
 * Copyright 2017 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>
/// Tracks the connectivity state between trajectories and the last time a global
/// constraint connected two trajectories.
///
/// Compared to ConnectedComponents it tracks additionally the last time that a global
/// constraint connected two trajectories.
///
/// Match C++ TrajectoryConnectivityState (trajectory_connectivity_state.cc)
/// </summary>
public class TrajectoryConnectivityState
{
    // ConnectedComponents is thread safe
    private readonly ConnectedComponents _connectedComponents = new();

    // Tracks the last time a direct connection between two trajectories has
    // been added. The exception is when a connection between two trajectories
    // connects two formerly unconnected connected components. In this case all
    // bipartite trajectories entries for these components are updated with the
    // new connection time.
    private readonly Dictionary<(int, int), long> _lastConnectionTimeMap = new();

    /// <summary>
    /// Add a trajectory which is initially connected to only itself.
    /// </summary>
    public void Add(int trajectoryId)
    {
        _connectedComponents.Add(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 and update the last
    /// connected time.
    /// </summary>
    public void Connect(int trajectoryIdA, int trajectoryIdB, long time)
    {
        if (TransitivelyConnected(trajectoryIdA, trajectoryIdB))
        {
            // The trajectories are transitively connected, i.e. they belong to the same
            // connected component. In this case we only update the last connection time
            // of those two trajectories.
            var sortedPair = (Math.Min(trajectoryIdA, trajectoryIdB), Math.Max(trajectoryIdA, trajectoryIdB));
            if (!_lastConnectionTimeMap.TryGetValue(sortedPair, out var existing) || existing < time)
            {
                _lastConnectionTimeMap[sortedPair] = time;
            }
        }
        else
        {
            // The connection between these two trajectories is about to join two
            // connected components. Here we update all bipartite trajectory pairs for
            // the two connected components with the connection time. This is to quickly
            // change to a more efficient loop closure search (by constraining the
            // search window) when connected components are joined.
            var componentA = _connectedComponents.GetComponent(trajectoryIdA);
            var componentB = _connectedComponents.GetComponent(trajectoryIdB);
            foreach (var idA in componentA)
            {
                foreach (var idB in componentB)
                {
                    var idPair = (Math.Min(idA, idB), Math.Max(idA, idB));
                    _lastConnectionTimeMap[idPair] = time;
                }
            }
        }
        _connectedComponents.Connect(trajectoryIdA, trajectoryIdB);
    }

    /// <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)
    {
        return _connectedComponents.TransitivelyConnected(trajectoryIdA, trajectoryIdB);
    }

    /// <summary>
    /// The trajectory IDs, grouped by connectivity.
    /// </summary>
    public List<List<int>> Components()
    {
        return _connectedComponents.Components();
    }

    /// <summary>
    /// Returns the last connection time between the two trajectories.
    /// If either of the trajectories is untracked or they have never been
    /// connected returns 0 (beginning of time).
    /// </summary>
    public long LastConnectionTime(int trajectoryIdA, int trajectoryIdB)
    {
        var sortedPair = (Math.Min(trajectoryIdA, trajectoryIdB), Math.Max(trajectoryIdA, trajectoryIdB));
        return _lastConnectionTimeMap.TryGetValue(sortedPair, out var t) ? t : 0L;
    }
}