RobotNet/RobotNet.RobotManager/Services/Planner/Space/RTree.cs

using RobotNet.MapShares;
using RobotNet.MapShares.Dtos;
using RobotNet.MapShares.Enums;

namespace RobotNet.RobotManager.Services.Planner.Space;

public class Rectangle(double minX, double minY, double maxX, double maxY)
{
    public double MinX { get; set; } = minX;
    public double MinY { get; set; } = minY;
    public double MaxX { get; set; } = maxX;
    public double MaxY { get; set; } = maxY;

    public double DistanceToPoint(double x, double y)
    {
        double dx = Math.Max(Math.Max(MinX - x, 0), x - MaxX);
        double dy = Math.Max(Math.Max(MinY - y, 0), y - MaxY);
        return Math.Sqrt(dx * dx + dy * dy);
    }

    public bool Contains(double x, double y)
    {
        return x >= MinX && x <= MaxX && y >= MinY && y <= MaxY;
    }
}

public class RTreeNode
{
    public Rectangle Bounds { get; set; } = new Rectangle(0, 0, 0, 0);
    public List<RTreeNode> Children { get; set; } = [];
    public List<(EdgeDto Edge, Rectangle Rect)> Entries { get; set; } = [];
    public bool IsLeaf => Children.Count == 0;
}

public class RTree
{
    private readonly RTreeNode Root;
    private const int MaxEntries = 4;
    private readonly List<NodeDto> Nodes;

    public RTree(List<NodeDto> nodes, List<EdgeDto> edges)
    {
        Nodes = nodes;
        Root = new RTreeNode();
        foreach (var edge in edges)
        {
            var rect = CalculateMBR(edge);
            Insert(Root, (edge, rect));
        }
    }

    private Rectangle CalculateMBR(EdgeDto edge)
    {
        var startNode = Nodes.FirstOrDefault(n => n.Id == edge.StartNodeId);
        var endNode = Nodes.FirstOrDefault(n => n.Id == edge.EndNodeId);
        if (startNode == null || endNode == null) return new Rectangle(0, 0, 0, 0);

        double minX = Math.Min(startNode.X, endNode.X);
        double maxX = Math.Max(startNode.X, endNode.X);
        double minY = Math.Min(startNode.Y, endNode.Y);
        double maxY = Math.Max(startNode.Y, endNode.Y);

        // Mở rộng MBR nếu edge là đường cong
        if (edge.TrajectoryDegree != TrajectoryDegree.One)
        {
            minX = Math.Min(minX, Math.Min(edge.ControlPoint1X, edge.ControlPoint2X));
            maxX = Math.Max(maxX, Math.Max(edge.ControlPoint1X, edge.ControlPoint2X));
            minY = Math.Min(minY, Math.Min(edge.ControlPoint1Y, edge.ControlPoint2Y));
            maxY = Math.Max(maxY, Math.Max(edge.ControlPoint1Y, edge.ControlPoint2Y));
        }

        return new Rectangle(minX, minY, maxX, maxY);
    }

    private static void Insert(RTreeNode node, (EdgeDto Edge, Rectangle Rect) entry)
    {
        if (node.IsLeaf)
        {
            node.Entries.Add(entry);
            if (node.Entries.Count > MaxEntries)
            {
                SplitNode(node);
            }
        }
        else
        {
            var bestChild = ChooseBestChild(node, entry.Rect);
            Insert(bestChild, entry);
            AdjustBounds(bestChild);
        }

        node.Bounds.MinX = Math.Min(node.Bounds.MinX, entry.Rect.MinX);
        node.Bounds.MinY = Math.Min(node.Bounds.MinY, entry.Rect.MinY);
        node.Bounds.MaxX = Math.Max(node.Bounds.MaxX, entry.Rect.MaxX);
        node.Bounds.MaxY = Math.Max(node.Bounds.MaxY, entry.Rect.MaxY);
    }

    private static RTreeNode ChooseBestChild(RTreeNode node, Rectangle rect)
    {
        RTreeNode best = node.Children[0];
        double minEnlargement = CalculateEnlargement(best.Bounds, rect);
        foreach (var child in node.Children.Skip(1))
        {
            double enlargement = CalculateEnlargement(child.Bounds, rect);
            if (enlargement < minEnlargement)
            {
                minEnlargement = enlargement;
                best = child;
            }
        }
        return best;
    }

    private static double CalculateEnlargement(Rectangle bounds, Rectangle rect)
    {
        double newArea = (Math.Max(bounds.MaxX, rect.MaxX) - Math.Min(bounds.MinX, rect.MinX)) *
                         (Math.Max(bounds.MaxY, rect.MaxY) - Math.Min(bounds.MinY, rect.MinY));
        double oldArea = (bounds.MaxX - bounds.MinX) * (bounds.MaxY - bounds.MinY);
        return newArea - oldArea;
    }

    private static void SplitNode(RTreeNode node)
    {
        var (group1, group2) = SplitEntries(node.Entries);
        node.Children.Add(new RTreeNode { Entries = group1 });
        node.Children.Add(new RTreeNode { Entries = group2 });
        node.Entries.Clear();
        foreach (var child in node.Children)
        {
            child.Bounds = CalculateBounds(child.Entries);
        }
    }

    private static (List<(EdgeDto, Rectangle)>, List<(EdgeDto, Rectangle)>) SplitEntries(List<(EdgeDto Edge, Rectangle Rect)> entries)
    {
        entries.Sort((a, b) => a.Rect.MinX.CompareTo(b.Rect.MinX));
        int mid = entries.Count / 2;
        return (entries.GetRange(0, mid), entries.GetRange(mid, entries.Count - mid));
    }

    private static Rectangle CalculateBounds(List<(EdgeDto Edge, Rectangle Rect)> entries)
    {
        if (entries.Count == 0) return new(0, 0, 0, 0);
        var first = entries[0].Rect;
        double minX = first.MinX, minY = first.MinY, maxX = first.MaxX, maxY = first.MaxY;
        foreach (var (Edge, Rect) in entries.Skip(1))
        {
            minX = Math.Min(minX, Rect.MinX);
            minY = Math.Min(minY, Rect.MinY);
            maxX = Math.Max(maxX, Rect.MaxX);
            maxY = Math.Max(maxY, Rect.MaxY);
        }
        return new Rectangle(minX, minY, maxX, maxY);
    }

    private static void AdjustBounds(RTreeNode node)
    {
        if (node.IsLeaf)
        {
            node.Bounds = CalculateBounds(node.Entries);
        }
        else
        {
            node.Bounds = CalculateBounds(node.Children.SelectMany(c => c.Entries).ToList());
        }
    }

    public EdgeDto? FindNearest(double x, double y, double limitDistance)
    {
        double minDistance = double.MaxValue;
        EdgeDto? nearestEdge = null;
        SearchNearest(Root, x, y, ref minDistance, ref nearestEdge);
        return minDistance <= limitDistance ? nearestEdge : null;
    }

    private void SearchNearest(RTreeNode node, double x, double y, ref double minDistance, ref EdgeDto? nearestEdge)
    {
        if (node.IsLeaf)
        {
            foreach (var (edge, rect) in node.Entries)
            {
                var startNode = Nodes.FirstOrDefault(n => n.Id == edge.StartNodeId);
                var endNode = Nodes.FirstOrDefault(n => n.Id == edge.EndNodeId);
                if (startNode == null || endNode == null) continue;

                var distanceResult = MapCompute.DistanceToEdge(x, y, edge, startNode, endNode);
                if (distanceResult.IsSuccess && distanceResult.Data < minDistance)
                {
                    minDistance = distanceResult.Data;
                    nearestEdge = edge;
                }
            }
        }
        else
        {
            var sortedChildren = node.Children.OrderBy(c => c.Bounds.DistanceToPoint(x, y));
            foreach (var child in sortedChildren)
            {
                if (child.Bounds.DistanceToPoint(x, y) < minDistance)
                {
                    SearchNearest(child, x, y, ref minDistance, ref nearestEdge);
                }
            }
        }
    }
}