namespace StellaOps.ElkSharp;

internal static class ElkNodeOrdering
{
    internal static ElkNode[][] OptimizeLayerOrdering(
        ElkNode[][] initialLayers,
        IReadOnlyDictionary<string, List<string>> incomingNodeIds,
        IReadOnlyDictionary<string, List<string>> outgoingNodeIds,
        IReadOnlyDictionary<string, int> inputOrder,
        int iterationCount = 8)
    {
        if (initialLayers.Length <= 2)
        {
            return initialLayers;
        }

        var layers = initialLayers
            .Select(layer => layer.ToList())
            .ToArray();

        for (var iteration = 0; iteration < iterationCount; iteration++)
        {
            for (var layerIndex = 1; layerIndex < layers.Length; layerIndex++)
            {
                OrderLayer(layers, layerIndex, incomingNodeIds, inputOrder);
            }

            for (var layerIndex = layers.Length - 2; layerIndex >= 0; layerIndex--)
            {
                OrderLayer(layers, layerIndex, outgoingNodeIds, inputOrder);
            }
        }

        return layers
            .Select(layer => layer.ToArray())
            .ToArray();
    }

    internal static void OrderLayer(
        IReadOnlyList<List<ElkNode>> layers,
        int layerIndex,
        IReadOnlyDictionary<string, List<string>> adjacentNodeIds,
        IReadOnlyDictionary<string, int> inputOrder)
    {
        var positions = BuildNodeOrderPositions(layers);
        var currentLayer = layers[layerIndex];
        currentLayer.Sort((left, right) =>
        {
            var leftRank = ResolveOrderingRank(left.Id, adjacentNodeIds, positions);
            var rightRank = ResolveOrderingRank(right.Id, adjacentNodeIds, positions);
            var comparison = leftRank.CompareTo(rightRank);
            if (comparison != 0)
            {
                return comparison;
            }

            comparison = positions[left.Id].CompareTo(positions[right.Id]);
            if (comparison != 0)
            {
                return comparison;
            }

            return inputOrder[left.Id].CompareTo(inputOrder[right.Id]);
        });
    }

    internal static Dictionary<string, int> BuildNodeOrderPositions(IReadOnlyList<List<ElkNode>> layers)
    {
        var positions = new Dictionary<string, int>(StringComparer.Ordinal);
        foreach (var layer in layers)
        {
            for (var index = 0; index < layer.Count; index++)
            {
                positions[layer[index].Id] = index;
            }
        }

        return positions;
    }

    internal static double ResolveOrderingRank(
        string nodeId,
        IReadOnlyDictionary<string, List<string>> adjacentNodeIds,
        IReadOnlyDictionary<string, int> positions)
    {
        if (!adjacentNodeIds.TryGetValue(nodeId, out var neighbors) || neighbors.Count == 0)
        {
            return double.PositiveInfinity;
        }

        var ordered = neighbors
            .Where(positions.ContainsKey)
            .Select(neighborId => (double)positions[neighborId])
            .OrderBy(value => value)
            .ToArray();

        if (ordered.Length == 0)
        {
            return double.PositiveInfinity;
        }

        var middle = ordered.Length / 2;
        return ordered.Length % 2 == 1
            ? ordered[middle]
            : (ordered[middle - 1] + ordered[middle]) / 2d;
    }
}