shape tree

import numpy as np
from shapely import STRtree, Polygon

from tools import geoutil

"""
records = {"1": Point(0, 0), "2": Point(2, 2)}
"""


class StrTree(object):
    def __init__(self, fid_geom_dict):
        self.fid_geom_dict = fid_geom_dict
        self.fids = np.array([fid for fid in self.fid_geom_dict])
        self.tree = STRtree([fid_geom_dict[fid] for fid in self.fids])

    def query(self, box: Polygon):
        tree_query = self.tree.query(box)
        fids = self.fids.take(tree_query).tolist()
        return [[fid, self.fid_geom_dict[fid]] for fid in fids]

    def query_geoms(self, box: Polygon):
        tree_query = self.tree.query(box)
        fids = self.fids.take(tree_query).tolist()
        return [self.fid_geom_dict[fid] for fid in fids]


def df_build_srt_tree(df, id_col_name):
    if not df.is_empty():
        fid_geom_dict = {}
        for row in df.iter_rows(named=True):
            geom = geoutil.from_wkt_safe(row['geometry'])
            if geom is not None:
                fid_geom_dict[row[id_col_name]] = geom
        if fid_geom_dict:
            return StrTree(fid_geom_dict)