using System;

namespace Algorithms.Graph.MinimumSpanningTree;

/// <summary>
///     Class that uses Prim's (Jarnik's algorithm) to determine the minimum
///     spanning tree (MST) of a given graph. Prim's algorithm is a greedy
///     algorithm that can determine the MST of a weighted undirected graph
///     in O(V^2) time where V is the number of nodes/vertices when using an
///     adjacency matrix representation.
///     More information: https://en.wikipedia.org/wiki/Prim%27s_algorithm
///     Pseudocode and runtime analysis: https://www.personal.kent.edu/~rmuhamma/Algorithms/MyAlgorithms/GraphAlgor/primAlgor.htm .
/// </summary>
public static class PrimMatrix
{
    /// <summary>
    ///     Determine the minimum spanning tree for a given weighted undirected graph.
    /// </summary>
    /// <param name="adjacencyMatrix">Adjacency matrix for graph to find MST of.</param>
    /// <param name="start">Node to start search from.</param>
    /// <returns>Adjacency matrix of the found MST.</returns>
    public static float[,] Solve(float[,] adjacencyMatrix, int start)
    {
        ValidateMatrix(adjacencyMatrix);

        var numNodes = adjacencyMatrix.GetLength(0);

        // Create array to represent minimum spanning tree
        var mst = new float[numNodes, numNodes];

        // Create array to keep track of which nodes are in the MST already
        var added = new bool[numNodes];

        // Create array to keep track of smallest edge weight for node
        var key = new float[numNodes];

        // Create array to store parent of node
        var parent = new int[numNodes];

        for (var i = 0; i < numNodes; i++)
        {
            mst[i, i] = float.PositiveInfinity;
            key[i] = float.PositiveInfinity;

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

        // Ensures that the starting node is added first
        key[start] = 0;

        // Keep looping until all nodes are in tree
        for (var i = 0; i < numNodes - 1; i++)
        {
            GetNextNode(adjacencyMatrix, key, added, parent);
        }

        // Build adjacency matrix for tree
        for (var i = 0; i < numNodes; i++)
        {
            if (i == start)
            {
                continue;
            }

            mst[i, parent[i]] = adjacencyMatrix[i, parent[i]];
            mst[parent[i], i] = adjacencyMatrix[i, parent[i]];
        }

        return mst;
    }

    /// <summary>
    ///     Ensure that the given adjacency matrix represents a weighted undirected graph.
    /// </summary>
    /// <param name="adjacencyMatrix">Adjacency matric to check.</param>
    private static void ValidateMatrix(float[,] adjacencyMatrix)
    {
        // Matrix should be square
        if (adjacencyMatrix.GetLength(0) != adjacencyMatrix.GetLength(1))
        {
            throw new ArgumentException("Adjacency matrix must be square!");
        }

        // Graph needs to be undirected and connected
        for (var i = 0; i < adjacencyMatrix.GetLength(0); i++)
        {
            var connection = false;
            for (var j = 0; j < adjacencyMatrix.GetLength(0); j++)
            {
                if (Math.Abs(adjacencyMatrix[i, j] - adjacencyMatrix[j, i]) > 1e-6)
                {
                    throw new ArgumentException("Adjacency matrix must be symmetric!");
                }

                if (!connection && float.IsFinite(adjacencyMatrix[i, j]))
                {
                    connection = true;
                }
            }

            if (!connection)
            {
                throw new ArgumentException("Graph must be connected!");
            }
        }
    }

    /// <summary>
    ///     Determine which node should be added next to the MST.
    /// </summary>
    /// <param name="adjacencyMatrix">Adjacency matrix of graph.</param>
    /// <param name="key">Currently known minimum edge weight connected to each node.</param>
    /// <param name="added">Whether or not a node has been added to the MST.</param>
    /// <param name="parent">The node that added the node to the MST. Used for building MST adjacency matrix.</param>
    private static void GetNextNode(float[,] adjacencyMatrix, float[] key, bool[] added, int[] parent)
    {
        var numNodes = adjacencyMatrix.GetLength(0);
        var minWeight = float.PositiveInfinity;

        var node = -1;

        // Find node with smallest node with known edge weight not in tree. Will always start with starting node
        for (var i = 0; i < numNodes; i++)
        {
            if (!added[i] && key[i] < minWeight)
            {
                minWeight = key[i];
                node = i;
            }
        }

        // Add node to mst
        added[node] = true;

        // Update smallest found edge weights and parent for adjacent nodes
        for (var i = 0; i < numNodes; i++)
        {
            if (!added[i] && adjacencyMatrix[node, i] < key[i])
            {
                key[i] = adjacencyMatrix[node, i];
                parent[i] = node;
            }
        }
    }
}