package com.xiaopeng.map.utils;

import com.google.flatbuffers.FlatBufferBuilder;
import com.xiaopeng.map.dto.Node;
import com.xiaopeng.map.dto.Point2D;
import com.xiaopeng.map.dto.RTree;
import com.xiaopeng.map.dto.Rect;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

/**
 * 等价于 C++ PbRTreeCreater<Node>，将叶节点列表构建成完整的 R-Tree 数组后写入 FlatBuffer。
 *
 * R-Tree 结构：完全二叉树，以数组形式存储（与 C++ 一致）
 *   - 左子节点下标：2 * parent + 1
 *   - 右子节点下标：2 * parent + 2
 *   - 总节点数：2^(ceil(log2(n))+1) - 1
 *
 * 节点类型（对应 map_area_rtree.proto 中 Node.type）：
 *   - 0: INVALID（空占位节点）
 *   - 1: LEAF（叶节点，存有实际区域数据）
 *   - 2: BRANCH（内部节点，存有子树的合并包围框）
 */
public final class PbRTreeBuilder {

    public static final int NODE_TYPE_INVALID = 0;
    public static final int NODE_TYPE_LEAF    = 1;
    public static final int NODE_TYPE_BRANCH  = 2;

    private PbRTreeBuilder() {
    }

    /**
     * R-Tree 中间节点结构，对应 C++ 中 map_area_rtree.proto 的 Node。
     */
    public static class RTreeNode {
        public int type = NODE_TYPE_INVALID;
        /** 包围框（range），单位与 FlatBuffers Point2D.x/y 一致（float） */
        public float minX, minY, maxX, maxY;
        /** 仅 LEAF 节点有效，每项 float[]{x, y} 对应一个多边形顶点 */
        public List<float[]> points;
        /** 业务 ID，对应 odd_id */
        public int id;
        /** 是否黑名单 */
        public boolean isBlacklist;
    }

    /**
     * 从叶节点列表构建完整 R-Tree 数组（等价于 PbRTreeCreater.CreatePbRTree）。
     *
     * @param leafNodes 原始叶节点，每个对应一个 OddArea（已算好 minX/Y, maxX/Y）
     * @return 按数组下标顺序排列的完整 R-Tree 节点列表（包含 INVALID 占位节点）
     */
    public static List<RTreeNode> createRTree(List<RTreeNode> leafNodes) {
        int n = leafNodes.size();
        // tree_height = ceil(log2(n))，n<=1 时 log2(1)=0
        int treeHeight = (n <= 1) ? 0 : (int) Math.ceil(Math.log(n) / Math.log(2));
        int nodeSize = (int) Math.pow(2, treeHeight + 1) - 1;

        List<RTreeNode> rtreeNodes = new ArrayList<>(nodeSize);
        for (int i = 0; i < nodeSize; i++) {
            rtreeNodes.add(new RTreeNode()); // type = INVALID
        }

        buildRecursive(new ArrayList<>(leafNodes), 0, leafNodes.size(), 0, rtreeNodes);
        return rtreeNodes;
    }

    private static void buildRecursive(List<RTreeNode> objects, int begin, int end,
                                       int index, List<RTreeNode> rtreeNodes) {
        int size = end - begin;
        if (size == 0) {
            return;
        }
        RTreeNode node = rtreeNodes.get(index);
        if (size == 1) {
            RTreeNode src = objects.get(begin);
            node.type        = NODE_TYPE_LEAF;
            node.minX        = src.minX;
            node.minY        = src.minY;
            node.maxX        = src.maxX;
            node.maxY        = src.maxY;
            node.points      = src.points;
            node.id          = src.id;
            node.isBlacklist = src.isBlacklist;
            return;
        }

        node.type = NODE_TYPE_BRANCH;
        combineRect(objects, begin, end, node);
        sort(objects, begin, end);

        int mid = begin + size / 2;
        buildRecursive(objects, begin, mid, 2 * index + 1, rtreeNodes);
        buildRecursive(objects, mid,   end, 2 * index + 2, rtreeNodes);
    }

    private static void combineRect(List<RTreeNode> objects, int begin, int end, RTreeNode out) {
        float minX = objects.get(begin).minX;
        float minY = objects.get(begin).minY;
        float maxX = objects.get(begin).maxX;
        float maxY = objects.get(begin).maxY;
        for (int i = begin + 1; i < end; i++) {
            RTreeNode n = objects.get(i);
            minX = Math.min(minX, n.minX);
            minY = Math.min(minY, n.minY);
            maxX = Math.max(maxX, n.maxX);
            maxY = Math.max(maxY, n.maxY);
        }
        out.minX = minX;
        out.minY = minY;
        out.maxX = maxX;
        out.maxY = maxY;
    }

    private static void sort(List<RTreeNode> objects, int begin, int end) {
        // 找当前子集的总包围框，以较长轴排序（与 C++ Sort 逻辑完全一致）
        float minX = objects.get(begin).minX, minY = objects.get(begin).minY;
        float maxX = objects.get(begin).maxX, maxY = objects.get(begin).maxY;
        for (int i = begin + 1; i < end; i++) {
            RTreeNode n = objects.get(i);
            minX = Math.min(minX, n.minX);
            minY = Math.min(minY, n.minY);
            maxX = Math.max(maxX, n.maxX);
            maxY = Math.max(maxY, n.maxY);
        }
        boolean sortByX = (maxX - minX) > (maxY - minY);

        List<RTreeNode> sub = objects.subList(begin, end);
        if (sortByX) {
            sub.sort(Comparator.comparingDouble(n -> n.maxX));
        } else {
            sub.sort(Comparator.comparingDouble(n -> n.maxY));
        }
    }

    /**
     * 将 R-Tree 节点列表序列化为 FlatBuffer，返回 RTree table 的 offset。
     * 与 C++ PbRTreeToFlatBuffer 的 FlatBuffer 写入逻辑对齐。
     */
    public static int writeFlatBuffer(FlatBufferBuilder builder, List<RTreeNode> rtreeNodes) {
        int[] nodeOffsets = new int[rtreeNodes.size()];
        for (int i = 0; i < rtreeNodes.size(); i++) {
            RTreeNode n = rtreeNodes.get(i);

            // points vector
            int pointsVector;
            if (n.points != null && !n.points.isEmpty()) {
                int[] ptOffsets = new int[n.points.size()];
                for (int j = 0; j < n.points.size(); j++) {
                    float[] pt = n.points.get(j);
                    ptOffsets[j] = Point2D.createPoint2D(builder, pt[0], pt[1]);
                }
                pointsVector = Node.createPointsVector(builder, ptOffsets);
            } else {
                pointsVector = Node.createPointsVector(builder, new int[0]);
            }

            // range
            int leftDown  = Point2D.createPoint2D(builder, n.minX, n.minY);
            int rightUp   = Point2D.createPoint2D(builder, n.maxX, n.maxY);
            int rangeOffset = Rect.createRect(builder, leftDown, rightUp);

            nodeOffsets[i] = Node.createNode(builder, n.type, rangeOffset, pointsVector, n.id, n.isBlacklist);
        }

        int nodesVector = RTree.createNodesVector(builder, nodeOffsets);
        return RTree.createRTree(builder, nodesVector);
    }
}