/// This module implements a generic [union-find](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) data structure.
/// This structure is used to manage disjoint sets, often used in partitioning or clustering problems.
/// 
/// Here is a summary of the features and members of this data structure:
/// 
/// 1. **Constructor**: The `UnionFind` constructor initializes a new instance of the Union-Find data structure with
///    default elements of the original set specified as a sequence. This set can be extended.
///
/// 2. **Public Members** :
///  - `UnionFind.AddElement` Adds a new element to the set.
///  - `UnionFind.AddRange` Adds a sequence of elements to the set.
///  - `UnionFind.Connect`: Combines two equivalence classes to which two given representatives belong into a single class.
///  - `UnionFind.ConnectRange`: Combines all equivalence classes to which the given representatives belong into a single class.
///  - `UnionFind.GetConnections`: Retrieves all elements of the equivalence class to which a given element belongs.
///  - `UnionFind.GetAllConnections`: Retrieves all equivalence classes.
///  - `UnionFind.AreConnected`: Checks whether two elements belong to the same equivalence class.
/// 
/// 3. **Inline functions**: These functions allow you to access members of the `UnionFind` instance more concisely.
/// 
/// Example:
///     <example>
///     let uf = UnionFind [| 1..6 |]
///     connectRange uf [| 1; 2; 3 |]
///     connectRange uf [| 4; 5 |]
///     </example>
///
module public UnionFind =
    open System.Collections.Generic
    open System.Linq

    /// Generic [union-find](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) data structure
    /// using the path compression find amelioration, making `Find` and `Connect` O(log n) in the worst case.
    type public UnionFind<'t when 't: equality>(defaultElements: 't seq) =
        let mutable parents = Dictionary<'t, 't>()
        let mutable ranks = Dictionary<'t, int>()

        do
            Seq.iter
                (fun element ->
                    parents.Add(element, element)
                    ranks.Add(element, 1))
                defaultElements

        let find (element: 't) =
            let rec inner (element: 't) =
                match parents.TryGetValue element with
                | true, parent when parent <> element ->
                    let grandParent = parents.[parent]
                    parents.[element] <- grandParent
                    inner grandParent
                | _ -> element

            inner element

        let treeSize (element: 't) = ranks.[element]

        member internal __.GetParents
            with get () = parents
            and set value = parents <- value

        member internal __.GetRanks
            with get () = ranks
            and set value = ranks <- value

        /// Adds an element to the original set.
        member public __.AddElement(element: 't) =
            match parents.TryGetValue element with
            | false, _ ->
                parents.Add(element, element)
                ranks.Add(element, 1)
            | _ -> ()

        /// Adds the list of elements given in the original set.
        member public this.AddRange(elements: 't seq) = Seq.iter this.AddElement elements

        /// Combines the two equivalence classes to which the two given representatives belong into a single class.
        member public __.Connect (left: 't) (right: 't) =
            assert (parents.ContainsKey left)
            assert (parents.ContainsKey right)

            let leftTree = find left
            let rightTree = find right

            if leftTree = rightTree then
                ()

            let leftTreeSize = treeSize leftTree
            let rightTreeSize = treeSize rightTree

            if leftTreeSize < rightTreeSize then
                parents.[leftTree] <- rightTree
                ranks.[rightTree] <- leftTreeSize + rightTreeSize
            else
                parents.[rightTree] <- leftTree
                ranks.[leftTree] <- leftTreeSize + rightTreeSize

        /// Combines all the equivalence classes to which the list of given representatives belong into a single class.
        member public this.ConnectRange(xs: 't seq) =
            if Seq.length xs < 2 then
                ()

            let head = Seq.head xs
            let tail = Seq.tail xs

            Seq.iter (this.Connect head) tail

        /// Retrieves the equivalence class to which the given element belongs.
        member public __.GetConnections(item: 't) =
            if not (parents.ContainsKey item) then
                Array.empty
            else
                let representative = find item

                parents
                    .GroupBy(fun kvp -> kvp.Value)
                    .Where(fun x -> x.Key = representative)
                    .Select(fun gp -> gp.Select(fun kvp -> kvp.Key) |> Seq.toArray)
                |> Seq.toArray

        /// Gets the set of equivalence classes.
        member public __.GetAllConnections() =
            parents
                .GroupBy(fun kvp -> kvp.Value)
                .Select(fun gp -> gp.Select(fun kvp -> kvp.Key) |> Seq.toArray)
            |> Seq.toArray

        /// Check whether two elements belong to the same equivalence class.
        member public __.AreConnected (left: 't) (right: 't) = find left = find right

    /// Adds an element to the original set.
    let inline public addElement (uf: 't UnionFind) = uf.AddElement
    /// Adds the list of elements given in the original set.
    let inline public addRange (uf: 't UnionFind) = uf.AddRange
    /// Combines the two equivalence classes to which the two given representatives belong into a single class.
    let inline public connect (uf: 't UnionFind) = uf.Connect
    /// Combines all the equivalence classes to which the list of given representatives belong into a single class.
    let inline public connectRange (uf: 't UnionFind) = uf.ConnectRange
    /// Retrieves the equivalence class to which the given element belongs.
    let inline public getConnections (uf: 't UnionFind) = uf.GetConnections
    /// Gets the set of equivalence classes.
    let inline public getAllConnections (uf: 't UnionFind) = uf.GetAllConnections()
    /// Check whether two elements belong to the same equivalence class.
    let inline public areConnected (uf: 't UnionFind) = uf.AreConnected