using System;
using System.Collections.Generic;
using DataStructures.DisjointSet;

namespace Algorithms.Graph.MinimumSpanningTree;

/// <summary>
///     Algorithm to determine the minimum spanning forest of an undirected graph.
/// </summary>
/// <remarks>
///     Kruskal's algorithm is a greedy algorithm that can determine the
///     minimum spanning tree or minimum spanning forest of any undirected
///     graph. Unlike Prim's algorithm, Kruskal's algorithm will work on
///     graphs that are unconnected. This algorithm will always have a
///     running time of O(E log V) where E is the number of edges and V is
///     the number of vertices/nodes.
///     More information: https://en.wikipedia.org/wiki/Kruskal%27s_algorithm .
///     Pseudocode and analysis: https://www.personal.kent.edu/~rmuhamma/Algorithms/MyAlgorithms/GraphAlgor/primAlgor.htm .
/// </remarks>
public static class Kruskal
{
    /// <summary>
    ///     Determine the minimum spanning tree/forest of the given graph.
    /// </summary>
    /// <param name="adjacencyMatrix">Adjacency matrix representing the graph.</param>
    /// <returns>Adjacency matrix of the minimum spanning tree/forest.</returns>
    public static float[,] Solve(float[,] adjacencyMatrix)
    {
        ValidateGraph(adjacencyMatrix);

        var numNodes = adjacencyMatrix.GetLength(0);
        var set = new DisjointSet<int>();
        var nodes = new Node<int>[numNodes];
        var edgeWeightList = new List<float>();
        var nodeConnectList = new List<(int, int)>();

        // Add nodes to disjoint set
        for (var i = 0; i < numNodes; i++)
        {
            nodes[i] = set.MakeSet(i);
        }

        // Create lists with edge weights and associated connectivity
        for (var i = 0; i < numNodes - 1; i++)
        {
            for (var j = i + 1; j < numNodes; j++)
            {
                if (float.IsFinite(adjacencyMatrix[i, j]))
                {
                    edgeWeightList.Add(adjacencyMatrix[i, j]);
                    nodeConnectList.Add((i, j));
                }
            }
        }

        var edges = Solve(set, nodes, edgeWeightList.ToArray(), nodeConnectList.ToArray());

        // Initialize minimum spanning tree
        var mst = new float[numNodes, numNodes];
        for (var i = 0; i < numNodes; i++)
        {
            mst[i, i] = float.PositiveInfinity;

            for (var j = i + 1; j < numNodes; j++)
            {
                mst[i, j] = float.PositiveInfinity;
                mst[j, i] = float.PositiveInfinity;
            }
        }

        foreach (var (node1, node2) in edges)
        {
            mst[node1, node2] = adjacencyMatrix[node1, node2];
            mst[node2, node1] = adjacencyMatrix[node1, node2];
        }

        return mst;
    }

    /// <summary>
    ///     Determine the minimum spanning tree/forest of the given graph.
    /// </summary>
    /// <param name="adjacencyList">Adjacency list representing the graph.</param>
    /// <returns>Adjacency list of the minimum spanning tree/forest.</returns>
    public static Dictionary<int, float>[] Solve(Dictionary<int, float>[] adjacencyList)
    {
        ValidateGraph(adjacencyList);

        var numNodes = adjacencyList.Length;
        var set = new DisjointSet<int>();
        var nodes = new Node<int>[numNodes];
        var edgeWeightList = new List<float>();
        var nodeConnectList = new List<(int, int)>();

        // Add nodes to disjoint set and create list of edge weights and associated connectivity
        for (var i = 0; i < numNodes; i++)
        {
            nodes[i] = set.MakeSet(i);

            foreach(var (node, weight) in adjacencyList[i])
            {
                edgeWeightList.Add(weight);
                nodeConnectList.Add((i, node));
            }
        }

        var edges = Solve(set, nodes, edgeWeightList.ToArray(), nodeConnectList.ToArray());

        // Create minimum spanning tree
        var mst = new Dictionary<int, float>[numNodes];
        for (var i = 0; i < numNodes; i++)
        {
            mst[i] = new Dictionary<int, float>();
        }

        foreach (var (node1, node2) in edges)
        {
            mst[node1].Add(node2, adjacencyList[node1][node2]);
            mst[node2].Add(node1, adjacencyList[node1][node2]);
        }

        return mst;
    }

    /// <summary>
    ///     Ensure that the given graph is undirected.
    /// </summary>
    /// <param name="adj">Adjacency matrix of graph to check.</param>
    private static void ValidateGraph(float[,] adj)
    {
        if (adj.GetLength(0) != adj.GetLength(1))
        {
            throw new ArgumentException("Matrix must be square!");
        }

        for (var i = 0; i < adj.GetLength(0) - 1; i++)
        {
            for (var j = i + 1; j < adj.GetLength(1); j++)
            {
                if (Math.Abs(adj[i, j] - adj[j, i]) > 1e-6)
                {
                    throw new ArgumentException("Matrix must be symmetric!");
                }
            }
        }
    }

    /// <summary>
    ///     Ensure that the given graph is undirected.
    /// </summary>
    /// <param name="adj">Adjacency list of graph to check.</param>
    private static void ValidateGraph(Dictionary<int, float>[] adj)
    {
        for (var i = 0; i < adj.Length; i++)
        {
            foreach (var edge in adj[i])
            {
                if (!adj[edge.Key].ContainsKey(i) || Math.Abs(edge.Value - adj[edge.Key][i]) > 1e-6)
                {
                    throw new ArgumentException("Graph must be undirected!");
                }
            }
        }
    }

    /// <summary>
    ///     Determine the minimum spanning tree/forest.
    /// </summary>
    /// <param name="set">Disjoint set needed for set operations.</param>
    /// <param name="nodes">List of nodes in disjoint set associated with each node.</param>
    /// <param name="edgeWeights">Weights of each edge.</param>
    /// <param name="connections">Nodes associated with each item in the <paramref name="edgeWeights"/> parameter.</param>
    /// <returns>Array of edges in the minimum spanning tree/forest.</returns>
    private static (int, int)[] Solve(DisjointSet<int> set, Node<int>[] nodes, float[] edgeWeights, (int, int)[] connections)
    {
        var edges = new List<(int, int)>();

        Array.Sort(edgeWeights, connections);

        foreach (var (node1, node2) in connections)
        {
            if (set.FindSet(nodes[node1]) != set.FindSet(nodes[node2]))
            {
                set.UnionSet(nodes[node1], nodes[node2]);
                edges.Add((node1, node2));
            }
        }

        return edges.ToArray();
    }
}