tyler 0.4.1

//! Spatial data structures for indexing the features.
// Copyright 2023 Balázs Dukai, Ravi Peters
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::parser::FeatureSet;
use log::{debug, info, warn};
use std::collections::VecDeque;
use std::fmt::{Display, Formatter};
use std::fs::{self, File};
use std::io::prelude::*;
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};

use morton_encoding::{morton_decode, morton_encode};

/// Quadtree
///
/// We don't expect that the quadtree has more than 65535 levels (u16).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct QuadTree {
    pub id: QuadTreeNodeId,
    side_length: u64,
    pub children: Vec<QuadTree>,
    cells: Vec<CellId>,
    pub nr_items: usize,
}

impl QuadTree {
    pub fn from_world(world: &crate::parser::World, limit: QuadTreeCapacity) -> Self {
        Self::from_grid(&world.grid, limit)
    }

    fn from_grid(grid: &SquareGrid, limit: QuadTreeCapacity) -> Self {
        let mut merge_limit: usize = 0;
        let nr_cells = grid.length.pow(2) as f64;
        let max_level = (nr_cells.ln() / 4.0_f64.ln()).ceil() as u16;
        debug!("Calculated maximum level for quadtree: {}", &max_level);
        // Somehow, the morton_encoding::morton_encode needs [y,x] for creating a Z-curve.
        //  If [x,y] is used, it creates an N-curve.
        //  crate::spatial_structs::interleave(x,y) creates a Z-curve.
        let mut mortoncodes: Vec<u128> = grid
            .into_iter()
            .map(|(cellid, _)| morton_encode([cellid.row as u64, cellid.column as u64]))
            .collect();
        mortoncodes.sort();
        let tiles_morton: Vec<QuadTree> = mortoncodes
            .iter()
            .map(|mc| {
                let coords: [u64; 2] = morton_decode(*mc);
                coords
            })
            .map(|[x, y]| {
                let cellid = CellId {
                    row: y as usize,
                    column: x as usize,
                };
                let items: usize;
                match limit {
                    QuadTreeCapacity::Objects(l) => {
                        // Use the number of features as a limit
                        items = grid.cell(&cellid).feature_ids.len();
                        merge_limit = l;
                    }
                    QuadTreeCapacity::Vertices(l) => {
                        // Use the number of vertices as a limit
                        items = grid.cell(&cellid).nr_vertices;
                        merge_limit = l;
                    }
                }
                QuadTree {
                    id: QuadTreeNodeId::new(x as usize, y as usize, max_level),
                    side_length: grid.cellsize as u64,
                    children: Vec::new(),
                    cells: vec![cellid],
                    nr_items: items,
                }
            })
            .collect();
        Self::merge_tiles(0, tiles_morton, merge_limit)
    }

    fn merge_tiles(level: u16, tiles: Vec<QuadTree>, limit: usize) -> QuadTree {
        let len_tiles = tiles.len();
        if len_tiles > 4 {
            let q0: usize = len_tiles / 4;
            let q1: usize = q0 * 2;
            let q2: usize = q0 * 3;
            let next_level = level + 1;
            Self::merge_tiles(
                level,
                vec![
                    Self::merge_tiles(next_level, tiles[0..q0].to_vec(), limit),
                    Self::merge_tiles(next_level, tiles[q0..q1].to_vec(), limit),
                    Self::merge_tiles(next_level, tiles[q1..q2].to_vec(), limit),
                    Self::merge_tiles(next_level, tiles[q2..].to_vec(), limit),
                ],
                limit,
            )
        } else {
            let sum_items: usize = tiles.iter().map(|t| t.nr_items).sum();
            let mut cells: Vec<CellId> = Vec::new();
            for t in tiles.iter() {
                for c in t.cells.iter() {
                    cells.push(*c)
                }
            }
            let id = QuadTreeNodeId::new(tiles[0].id.x, tiles[0].id.y, level);
            let _id_string = id.to_string();
            // FIXME: this also adds the quadtree if sum_items == 0 so the parent will have 4
            //  children instead of 3. Probably should return Option<Quadtree>.
            //  Currently these empty tiles are removed in Tile.prune().
            if sum_items <= limit {
                QuadTree {
                    id,
                    side_length: tiles[0].side_length * 2,
                    children: vec![],
                    cells,
                    nr_items: sum_items,
                }
            } else {
                if tiles.len() % 4 != 0 {
                    warn!(
                        "number of children is not dividable by 4: {}, in level {}",
                        tiles.len(),
                        &level
                    );
                }
                QuadTree {
                    id,
                    side_length: tiles[0].side_length * 2,
                    children: tiles.clone(),
                    cells: vec![],
                    nr_items: sum_items,
                }
            }
        }
    }

    #[allow(dead_code)]
    fn collect_leaves_recurse<'collect>(&'collect self, leaves: &mut Vec<&'collect QuadTree>) {
        if !self.children.is_empty() {
            for child in self.children.iter() {
                child.collect_leaves_recurse(leaves);
            }
        } else {
            leaves.push(self);
        }
    }

    #[allow(dead_code)]
    pub fn collect_leaves(&self) -> Vec<&Self> {
        let mut leaves: Vec<&QuadTree> = Vec::new();
        self.collect_leaves_recurse(&mut leaves);
        leaves
    }

    pub fn bbox(&self, grid: &SquareGrid) -> Bbox {
        let minx = grid.origin[0] + (self.id.x * grid.cellsize as usize) as f64;
        let miny = grid.origin[1] + (self.id.y * grid.cellsize as usize) as f64;
        [
            minx,
            miny,
            grid.bbox[2],
            minx + self.side_length as f64,
            miny + self.side_length as f64,
            grid.bbox[5],
        ]
    }

    /// Compute the bounding box of all the features in the node
    pub fn node_content_bbox(
        &self,
        world: &crate::parser::World,
        arg_minz: Option<i32>,
        arg_maxz: Option<i32>,
    ) -> Bbox {
        let mut tile_content_bbox: Option<Bbox> = None;
        for cellid in self.cells() {
            let cell = world.grid.cell(cellid);
            if !cell.feature_ids.is_empty() {
                tile_content_bbox = Some(world.features[cell.feature_ids[0]].bbox);
                break;
            }
        }
        for cellid in self.cells() {
            let cell = world.grid.cell(cellid);
            for fi in cell.feature_ids.iter() {
                if let Some(bbox) = tile_content_bbox.as_mut() {
                    update_bbox(bbox, &world.features[*fi].bbox);
                }
            }
        }
        let mut bbox = tile_content_bbox.unwrap_or(world.grid.bbox);
        if let Some(minz) = arg_minz {
            if bbox[2] < minz as f64 {
                bbox[2] = minz as f64;
            }
        }
        if let Some(maxz) = arg_maxz {
            if bbox[5] > maxz as f64 {
                bbox[5] = maxz as f64;
            }
        }
        bbox
    }

    /// Breadth-first search for a node.
    pub fn node(&self, id: &QuadTreeNodeId) -> Option<&QuadTree> {
        let mut q = VecDeque::new();
        q.push_back(self);

        while let Some(n) = q.pop_front() {
            if &n.id == id {
                return Some(n);
            } else {
                for child in &n.children {
                    q.push_back(child);
                }
            }
        }
        // Did not find the node
        None
    }

    pub fn cells(&self) -> Vec<&CellId> {
        let mut cellids: Vec<&CellId> = Vec::new();
        let mut q = VecDeque::new();
        q.push_back(self);
        while let Some(node) = q.pop_front() {
            if node.children.is_empty() {
                for cell in node.cells.iter() {
                    cellids.push(cell);
                }
            } else {
                for child in &node.children {
                    q.push_back(child);
                }
            }
        }
        cellids
    }

    /// The WKT of the 2D boundary of the current node.
    pub fn to_wkt(&self, grid: &SquareGrid) -> String {
        bbox_to_wkt(&self.bbox(grid))
    }

    pub fn export(
        &self,
        world: &crate::parser::World,
        output_dir: Option<&Path>,
    ) -> std::io::Result<()> {
        let mut q = VecDeque::new();
        q.push_back(self);
        let mut quadtree_level: u16 = self.id.level;
        let [outdir_quadtree, outdir_quadtree_content] = match output_dir {
            None => [Path::new(""), Path::new("")],
            Some(outdir) => [outdir, outdir],
        };
        let mut file_quadtree =
            File::create(outdir_quadtree.join(format!("quadtree_level-{quadtree_level}.tsv")))?;
        let mut file_quadtree_content = File::create(
            outdir_quadtree_content.join(format!("quadtree_content_level-{quadtree_level}.tsv")),
        )?;

        file_quadtree
            .write_all("node_id\tnode_level\tnr_items\twkt\n".as_bytes())
            .expect("cannot write quadtree header");
        file_quadtree_content
            .write_all("node_id\tnode_level\tnr_items\twkt\n".as_bytes())
            .expect("cannot write quadtree content header");

        while let Some(node) = q.pop_front() {
            if node.id.level != quadtree_level {
                quadtree_level = node.id.level;
                file_quadtree = File::create(
                    outdir_quadtree.join(format!("quadtree_level-{quadtree_level}.tsv")),
                )?;
                file_quadtree_content = File::create(
                    outdir_quadtree_content
                        .join(format!("quadtree_content_level-{quadtree_level}.tsv")),
                )?;
                file_quadtree
                    .write_all("node_id\tnode_level\tnr_items\twkt\n".as_bytes())
                    .expect("cannot write quadtree header");
                file_quadtree_content
                    .write_all("node_id\tnode_level\tnr_items\twkt\n".as_bytes())
                    .expect("cannot write quadtree content header");
            }
            let wkt = node.to_wkt(&world.grid);
            file_quadtree
                .write_all(
                    format!(
                        "{}\t{}\t{}\t{}\n",
                        node.id, node.id.level, node.nr_items, wkt
                    )
                    .as_bytes(),
                )
                .expect("cannot write quadtree node");
            let node_content_bbox = node.node_content_bbox(world, None, None);
            let wkt_content_bbox = bbox_to_wkt(&node_content_bbox);
            file_quadtree_content
                .write_all(
                    format!(
                        "{}\t{}\t{}\t{}\n",
                        node.id, node.id.level, node.nr_items, wkt_content_bbox
                    )
                    .as_bytes(),
                )
                .expect("cannot write quadtree node content");

            for child in &node.children {
                q.push_back(child);
            }
        }
        Ok(())
    }

    pub fn export_bincode(
        &self,
        name: Option<&str>,
        output_dir: Option<&Path>,
    ) -> bincode::Result<()> {
        let file_name: &str = name.unwrap_or("quadtree");
        let file = match output_dir {
            None => File::create(format!("{file_name}.bincode"))?,
            Some(outdir) => File::create(outdir.join(format!("{file_name}.bincode")))?,
        };
        bincode::serialize_into(file, self)
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct QuadTreeNodeId {
    pub x: usize,
    pub y: usize,
    pub level: u16,
}

impl QuadTreeNodeId {
    pub fn new(x: usize, y: usize, level: u16) -> Self {
        Self { x, y, level }
    }
}

impl Display for QuadTreeNodeId {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}/{}/{}", self.level, self.x, self.y)
    }
}

/// We have these double enum, QuadTreeCapacity and QuadTreeCapacityType, because of
/// how the CLI arguments are parsed. In the quadtree, we need QuadTreeCapacity, because
/// it can hold both the leaf capacity and the capacity type. But clap can only parse
/// into unit variants (I think), so we take the the capacity and capacity type as
/// separate arguments.
#[derive(Debug)]
pub enum QuadTreeCapacity {
    Objects(usize),
    Vertices(usize),
}

/// The type of items to count for the quadtree leaf capacity.
#[derive(Debug, Default, Clone, clap::ValueEnum)]
pub enum QuadTreeCriteria {
    Objects,
    #[default]
    Vertices,
}

/// 64-bit mask
#[allow(dead_code)]
fn part1by1_64(number: &u64) -> u64 {
    let mut n = *number;
    n &= 0x00000000ffffffff; // binary: 11111111111111111111111111111111,                                len: 32
    n = (n | (n << 16)) & 0x0000FFFF0000FFFF; // binary: 1111111111111111000000001111111111111111,                        len: 40
    n = (n | (n << 8)) & 0x00FF00FF00FF00FF; // binary: 11111111000000001111111100000000111111110000000011111111,        len: 56
    n = (n | (n << 4)) & 0x0F0F0F0F0F0F0F0F; // binary: 111100001111000011110000111100001111000011110000111100001111,    len: 60
    n = (n | (n << 2)) & 0x3333333333333333; // binary: 11001100110011001100110011001100110011001100110011001100110011,  len: 62
    n = (n | (n << 1)) & 0x5555555555555555; // binary: 101010101010101010101010101010101010101010101010101010101010101, len: 63
    n
}

/// Computing Morton-code from 64bit integers.
///
/// Reference: https://github.com/trevorprater/pymorton
#[allow(dead_code)]
pub fn interleave(x: &u64, y: &u64) -> u64 {
    part1by1_64(x) | (part1by1_64(y) << 1)
}

#[allow(dead_code)]
fn unpart1by1_64(mortoncode: &u64) -> u64 {
    let mut n = *mortoncode;
    n &= 0x5555555555555555; // binary: 101010101010101010101010101010101010101010101010101010101010101, len: 63
    n = (n ^ (n >> 1)) & 0x3333333333333333; // binary: 11001100110011001100110011001100110011001100110011001100110011,  len: 62
    n = (n ^ (n >> 2)) & 0x0f0f0f0f0f0f0f0f; // binary: 111100001111000011110000111100001111000011110000111100001111,    len: 60
    n = (n ^ (n >> 4)) & 0x00ff00ff00ff00ff; // binary: 11111111000000001111111100000000111111110000000011111111,        len: 56
    n = (n ^ (n >> 8)) & 0x0000ffff0000ffff; // binary: 1111111111111111000000001111111111111111,                        len: 40
    n = (n ^ (n >> 16)) & 0x00000000ffffffff; // binary: 11111111111111111111111111111111,                                len: 32
    n
}

/// Computing `[x, y]` from a Morton-code.
///
/// Reference: https://github.com/trevorprater/pymorton
#[allow(dead_code)]
pub fn deinterleave(mortoncode: &u64) -> [u64; 2] {
    [
        unpart1by1_64(mortoncode),
        unpart1by1_64(&(*mortoncode >> 1)),
    ]
}

/// Represents a square grid with square cells.
/// The grid stores the feature-indices in its cells.
/// The `length` of the grid is the number of cells of one dimension, thus the total
/// number of cells is obtained by `length * length`.
///
/// Note the a 'column' in the grid is represented by the X-axis, and a 'row' by the Y-axis.
/// See [CellId] for details.
///
/// The EPSG code of the input features is also stored in the grid, because the grid is initialized
/// directly from the feature coordinates without reprojection. Often we need to reproject the grid
/// to another CRS, for instance in order to convert it to 3D Tiles.
///
/// ```shell
///  (column)     (column)
///   +----+       +----+
///   |    | +---+ |    | +------------------+
///   |  --+-+-> | |    | |Vec<usize> (cell) |
///   |    | +---+ |    | +------------------+
///   |    |       |    |
///   |    | +---+ |    | +------------------+
///   |  --+-+-> | |    | |Vec<usize> (cell) |
///   |    | +---+ |    | +------------------+
///   |    |       |    |
///   |    | +---+ |    | +------------------+
///   |  --+-+-> | |    | |Vec<usize> (cell) |
///   | ^  | +---+ |    | +------------------+
///   +-+--+       +----+
///     |
/// +---+------------------------+
/// |   |                        |
/// | Vec<Vec<Vec<usize>>> (row) |
/// +----------------------------+
///
/// (created with https://asciiflow.com)
/// ```
///
/// ## Examples
///
/// ```ignore
/// use crate::spatial_structs::SquareGrid;
///
/// let grid = SquareGrid::new(&[0.0, 0.0, 0.0, 4.0, 4.0, 4.0], 1);
/// let grid_idx = grid.locate_point(&[2.5, 1.5]);
/// assert_eq!(grid_idx, [3_u64, 2_u64]);
/// ```
///
#[derive(Debug, Serialize, Deserialize)]
pub struct SquareGrid {
    origin: [f64; 3],
    pub bbox: Bbox,
    pub length: usize,
    cellsize: u32,
    pub data: Vec<Vec<Cell>>,
    pub epsg: u16,
}

/// Lightweight, copyable view of a [`SquareGrid`]'s metadata that exposes the
/// read-only spatial queries (`locate_point`, `intersect_bbox_ranges`). Letting
/// workers hold a `GridLayout` by value lets the integrator hold `&mut SquareGrid`
/// concurrently without violating Rust's borrowing rules.
#[derive(Debug, Clone, Copy)]
pub struct GridLayout {
    origin: [f64; 3],
    pub bbox: Bbox,
    pub length: usize,
    cellsize: u32,
    pub epsg: u16,
}

impl GridLayout {
    pub fn locate_point(&self, point: &[f64; 2]) -> CellId {
        let dx = point[0] - self.origin[0];
        let dy = point[1] - self.origin[1];
        let max_index = self.length.saturating_sub(1);
        let col_i = ((dx / self.cellsize as f64).floor() as isize).clamp(0, max_index as isize);
        let row_i = ((dy / self.cellsize as f64).floor() as isize).clamp(0, max_index as isize);
        CellId {
            row: row_i as usize,
            column: col_i as usize,
        }
    }

    pub fn intersect_bbox_ranges(
        &self,
        bbox: &Bbox,
    ) -> (
        std::ops::RangeInclusive<usize>,
        std::ops::RangeInclusive<usize>,
    ) {
        let [minx, miny, _, maxx, maxy, _] = *bbox;
        let min_cellid = self.locate_point(&[minx, miny]);
        let max_cellid = self.locate_point(&[maxx, maxy]);
        (
            min_cellid.column..=max_cellid.column,
            min_cellid.row..=max_cellid.row,
        )
    }

    pub fn intersect_bbox(&self, bbox: &Bbox) -> Vec<CellId> {
        let mut cellids: Vec<CellId> = Vec::new();
        let (columns, rows) = self.intersect_bbox_ranges(bbox);
        for column in columns {
            for row in rows.clone() {
                cellids.push(CellId { row, column });
            }
        }
        cellids
    }
}

impl Display for SquareGrid {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "SquareGrid (origin: {:?}, bbox: {:?}, length: {}, cellsize: {}, data: not-displayed, epsg: {})",
            self.origin, self.bbox, self.length, self.cellsize, self.epsg
        )
    }
}

impl SquareGrid {
    /// Creates a grid with cells of `cellsize`, that covers the `extent`.
    /// The grid and the cells are square.
    /// The grid center is the `extent` center.
    /// The grid is returned as an origin coordinate and the number of cells.
    pub fn new(extent: &Bbox, cellsize: u32, epsg: u16) -> Self {
        // We only compute 2D. Z is constant for the grid.
        let extent_center = [
            extent[0] + (extent[3] - extent[0]) / 2.0,
            extent[1] + (extent[4] - extent[1]) / 2.0,
        ];
        let dx = extent[3] - extent[0];
        let dy = extent[4] - extent[1];
        // The grid starting dimension is the longest edge of the rectangle
        let mut d = if dx > dy { dx } else { dy };
        // The grid is constructed as a square, from the center of the extent.
        // We don't need to add any buffer to the extent, because we already round up when we
        // compute the number of cells.
        // We need a grid that is has 2^n cells in one dimension, so that we can build
        // a 4^n cells quadtree.
        // Adjust the cellsize so that we can get a tightly fit square on the extent
        let mut cellsize_new = d;
        let cellsize_f64 = cellsize as f64;
        loop {
            let cn = cellsize_new / 2.0;
            if cn < cellsize_f64 {
                break;
            } else {
                cellsize_new /= 2.0;
            }
        }
        let d_cells = (d / cellsize_new).ceil() as usize;
        let cellsize = cellsize_new.ceil() as u32;
        let estimated_cells = d_cells.saturating_mul(d_cells);
        let estimated_bytes = estimated_cells.saturating_mul(std::mem::size_of::<Cell>());
        let estimated_mib = estimated_bytes as f64 / (1024.0 * 1024.0);
        info!(
            "Allocating dense square grid: {}x{} cells (~{:.1} MiB base cell storage)",
            d_cells, d_cells, estimated_mib
        );
        if estimated_mib > 1024.0 {
            warn!(
                "Dense square grid allocation is large: {}x{} cells (~{:.1} MiB). This usually means the computed extent or cell size is off.",
                d_cells, d_cells, estimated_mib
            );
        }
        // Compute new dimension from the calculated length
        d = d_cells as f64 * cellsize as f64;
        let origin = [
            extent_center[0] - d / 2.0,
            extent_center[1] - d / 2.0,
            extent[2],
        ];
        let bbox = [
            origin[0],
            origin[1],
            origin[2],
            origin[0] + d,
            origin[1] + d,
            extent[5],
        ];
        // A row-vector (x-axis) to store the column-vectors (y-axis).
        let mut row: Vec<Vec<Cell>> = Vec::with_capacity(d_cells);
        // For each column create a column vector that stores the cells and for each row in the
        // column create a cell to store the feature IDs.
        row.resize_with(d_cells, || {
            let mut column: Vec<Cell> = Vec::with_capacity(d_cells);
            column.resize(
                d_cells,
                Cell {
                    feature_ids: Vec::new(),
                    nr_vertices: 0,
                },
            );
            column
        });
        Self {
            origin,
            bbox,
            length: d_cells,
            cellsize,
            data: row,
            epsg,
        }
    }

    /// A snapshot of the grid metadata needed for read-only spatial queries.
    /// Pass this to worker threads instead of `&self` so the integrator can keep
    /// exclusive access to the cell data.
    pub fn layout(&self) -> GridLayout {
        GridLayout {
            origin: self.origin,
            bbox: self.bbox,
            length: self.length,
            cellsize: self.cellsize,
            epsg: self.epsg,
        }
    }

    pub fn from_debug_tsv(
        grid_path: &Path,
        features_path: Option<&Path>,
        epsg: u16,
        fallback_grid: Option<&Self>,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let grid_tsv = fs::read_to_string(grid_path)?;
        let mut lines = grid_tsv.lines();
        let _header = lines.next();
        let root = lines.next().ok_or_else(|| {
            format!(
                "debug grid file {} is missing its root row",
                grid_path.display()
            )
        })?;
        let mut root_parts = root.splitn(3, '\t');
        let root_id = root_parts.next().unwrap_or_default();
        if root_id != "x-x" {
            return Err(format!(
                "debug grid file {} has invalid root row id {root_id:?}",
                grid_path.display()
            )
            .into());
        }
        let _root_count = root_parts.next();
        let root_wkt = root_parts.next().ok_or_else(|| {
            format!(
                "debug grid file {} is missing root WKT",
                grid_path.display()
            )
        })?;
        let [minx, miny, maxx, maxy] = parse_grid_polygon_wkt(root_wkt)?;

        let mut cells = Vec::new();
        let mut max_column = 0usize;
        let mut max_row = 0usize;
        for line in lines {
            if line.trim().is_empty() {
                continue;
            }
            let mut parts = line.splitn(3, '\t');
            let cell_id = parse_cell_id(parts.next().unwrap_or_default())?;
            let nr_vertices = parts
                .next()
                .ok_or_else(|| {
                    format!(
                        "debug grid file {} is missing nr_items",
                        grid_path.display()
                    )
                })?
                .parse::<usize>()?;
            max_column = max_column.max(cell_id.column);
            max_row = max_row.max(cell_id.row);
            cells.push((cell_id, nr_vertices));
        }

        let length = max_column.max(max_row) + 1;
        let cellsize = ((maxx - minx) / length as f64).round() as u32;
        let fallback_bbox = fallback_grid.map_or([0.0; 6], |grid| grid.bbox);
        let bbox = [minx, miny, fallback_bbox[2], maxx, maxy, fallback_bbox[5]];

        let mut data = vec![
            vec![
                Cell {
                    feature_ids: Vec::new(),
                    nr_vertices: 0,
                };
                length
            ];
            length
        ];
        for (cell_id, nr_vertices) in cells {
            data[cell_id.column][cell_id.row].nr_vertices = nr_vertices;
        }

        if let Some(grid) = fallback_grid {
            let limit = length.min(grid.length);
            for (column_index, column) in data.iter_mut().enumerate().take(limit) {
                for (row_index, cell) in column.iter_mut().enumerate().take(limit) {
                    cell.feature_ids = grid.data[column_index][row_index].feature_ids.clone();
                }
            }
        }

        if let Some(features_path) = features_path.filter(|path| path.exists()) {
            for column in &mut data {
                for cell in column {
                    cell.feature_ids.clear();
                }
            }

            let features_tsv = fs::read_to_string(features_path)?;
            for line in features_tsv.lines().skip(1) {
                if line.trim().is_empty() {
                    continue;
                }
                let mut parts = line.splitn(3, '\t');
                let feature_id = parts
                    .next()
                    .ok_or_else(|| {
                        format!(
                            "debug features file {} is missing feature id",
                            features_path.display()
                        )
                    })?
                    .parse::<usize>()?;
                let cell_id = parse_cell_id(parts.next().unwrap_or_default())?;
                data[cell_id.column][cell_id.row]
                    .feature_ids
                    .push(feature_id);
            }
        }

        Ok(Self {
            origin: [minx, miny, bbox[2]],
            bbox,
            length,
            cellsize,
            data,
            epsg,
        })
    }

    /// Returns the cell index (x, y) where the point is located.
    pub fn locate_point(&self, point: &[f64; 2]) -> CellId {
        let dx = point[0] - self.origin[0];
        let dy = point[1] - self.origin[1];
        let max_index = self.length.saturating_sub(1);
        let col_i = ((dx / self.cellsize as f64).floor() as isize).clamp(0, max_index as isize);
        let row_i = ((dy / self.cellsize as f64).floor() as isize).clamp(0, max_index as isize);
        CellId {
            row: row_i as usize,
            column: col_i as usize,
        }
    }

    /// Insert a point into the grid, by assigning it to the Cell where it is located.
    pub fn insert(&mut self, point: &[f64; 2], feature_id: usize) -> CellId {
        let cell_id = self.locate_point(point);
        let cell = self.cell_mut(&cell_id);
        cell.feature_ids.push(feature_id);
        cell_id
    }

    /// Return the Cells that intersect the Bounding Box.
    pub fn intersect_bbox(&self, bbox: &Bbox) -> Vec<CellId> {
        let mut cellids: Vec<CellId> = Vec::new();
        let (columns, rows) = self.intersect_bbox_ranges(bbox);
        for column in columns {
            for row in rows.clone() {
                cellids.push(CellId { row, column });
            }
        }
        cellids
    }

    pub fn intersect_bbox_ranges(
        &self,
        bbox: &Bbox,
    ) -> (
        std::ops::RangeInclusive<usize>,
        std::ops::RangeInclusive<usize>,
    ) {
        let [minx, miny, _, maxx, maxy, _] = *bbox;
        let min_cellid = self.locate_point(&[minx, miny]);
        let max_cellid = self.locate_point(&[maxx, maxy]);
        (
            min_cellid.column..=max_cellid.column,
            min_cellid.row..=max_cellid.row,
        )
    }

    /// Exports the grid and the feature centroids into TSV files into the working directory.
    /// The grid is written to `grid.tsv`.
    /// If `feature_set` is provided, then the feature centroids are written to
    /// `features.tsv`.
    /// If `output_dir` is provided, the files are written there. Else they are written to the
    /// working directory.
    pub fn export(
        &self,
        feature_set: Option<&FeatureSet>,
        output_dir: Option<&Path>,
    ) -> std::io::Result<()> {
        let [file_grid_path, file_features_path] = match output_dir {
            None => [PathBuf::from("grid.tsv"), PathBuf::from("features.tsv")],
            Some(outdir) => [outdir.join("grid.tsv"), outdir.join("features.tsv")],
        };
        let mut file_grid = File::create(&file_grid_path)?;
        let mut file_features = File::create(&file_features_path)?;
        let root_wkt = format!(
            "POLYGON(({minx} {miny}, {maxx} {miny}, {maxx} {maxy}, {minx} {maxy}, {minx} {miny}))",
            minx = self.bbox[0],
            miny = self.bbox[1],
            maxx = self.bbox[3],
            maxy = self.bbox[4]
        );
        file_grid
            .write_all("cell_id\tnr_items\twkt\n".as_bytes())
            .expect("cannot write grid header");
        file_grid
            .write_all(format!("x-x\t0\t{}\n", root_wkt).as_bytes())
            .expect("cannot write grid line");
        file_features
            .write_all("fid\tcell_id\twkt\n".as_bytes())
            .expect("cannot write features header");
        for (cellid, cell) in self {
            let wkt = self.cell_to_wkt(&cellid);
            file_grid
                .write_all(format!("{}\t{}\t{}\n", &cellid, cell.nr_vertices, wkt).as_bytes())
                .expect("cannot write grid line");
            let mut cellbuffer = String::new();
            if let Some(fset) = feature_set {
                for fid in cell.feature_ids.iter() {
                    let f = &fset[*fid];
                    let centroid = f.centroid();
                    cellbuffer += format!(
                        "{}\t{}\tPOINT({} {})\n",
                        fid, &cellid, centroid[0], centroid[1]
                    )
                    .as_str();
                }
            }
            if feature_set.is_some() {
                file_features
                    .write_all(cellbuffer.as_bytes())
                    .expect("cannot write cell contents");
            }
        }
        if feature_set.is_none() {
            // Remove empty file
            std::fs::remove_file(file_features_path)?;
        }
        Ok(())
    }

    pub fn cell_to_wkt(&self, cellid: &CellId) -> String {
        let minx = self.origin[0] + (cellid.column * self.cellsize as usize) as f64;
        let miny = self.origin[1] + (cellid.row * self.cellsize as usize) as f64;
        bbox_to_wkt(&[
            minx,
            miny,
            self.bbox[2],
            minx + self.cellsize as f64,
            miny + self.cellsize as f64,
            self.bbox[5],
        ])
    }

    pub fn cell_bbox(&self, cellid: &CellId) -> Bbox {
        let minx = self.origin[0] + (cellid.column * self.cellsize as usize) as f64;
        let miny = self.origin[1] + (cellid.row * self.cellsize as usize) as f64;
        let minz = self.bbox[2];
        let maxx = minx + self.cellsize as f64;
        let maxy = miny + self.cellsize as f64;
        let maxz = self.bbox[5];
        [minx, miny, minz, maxx, maxy, maxz]
    }

    pub fn cell(&self, cell_id: &CellId) -> &Cell {
        &self.data[cell_id.column][cell_id.row]
    }

    pub fn cell_mut(&mut self, cell_id: &CellId) -> &mut Cell {
        &mut self.data[cell_id.column][cell_id.row]
    }

    /// Compute the vertex distribution in the cells of the grid.
    pub fn compute_statistics(&self) -> SquareGridStats {
        // nr. of vertices in the cells that are not empty
        let mut nr_vertices_not_empty: Vec<usize> = Vec::new();
        let mut nr_cells_not_empty: usize = 0;
        for (_, cell) in self {
            if cell.nr_vertices > 0 {
                nr_vertices_not_empty.push(cell.nr_vertices);
                nr_cells_not_empty += 1;
            }
        }
        // Handle empty case
        if nr_vertices_not_empty.is_empty() {
            return SquareGridStats {
                nr_vertices: 0,
                nr_cells_with_content: 0,
                nr_vertices_min: 0,
                nr_vertices_max: 0,
                nr_vertices_mean: 0.0,
                nr_vertices_median: 0.0,
            };
        }

        let sum = nr_vertices_not_empty.iter().sum();
        let mean = sum as f64 / nr_cells_not_empty as f64;

        // Sort for median calculation
        nr_vertices_not_empty.sort();

        // Calculate median
        let median = if nr_vertices_not_empty.len() % 2 == 0 {
            // Even number of elements: average of the two middle elements
            let mid = nr_vertices_not_empty.len() / 2;
            (nr_vertices_not_empty[mid - 1] as f64 + nr_vertices_not_empty[mid] as f64) / 2.0
        } else {
            // Odd number of elements: the middle element
            let mid = nr_vertices_not_empty.len() / 2;
            nr_vertices_not_empty[mid] as f64
        };

        SquareGridStats {
            nr_vertices: sum,
            nr_cells_with_content: nr_cells_not_empty,
            nr_vertices_min: *nr_vertices_not_empty.iter().min().unwrap(),
            nr_vertices_max: *nr_vertices_not_empty.iter().max().unwrap(),
            nr_vertices_mean: mean,
            nr_vertices_median: median,
        }
    }
}

fn parse_cell_id(value: &str) -> Result<CellId, Box<dyn std::error::Error>> {
    let (column, row) = value
        .split_once('-')
        .ok_or_else(|| format!("invalid cell id {value:?}"))?;
    Ok(CellId {
        column: column.parse::<usize>()?,
        row: row.parse::<usize>()?,
    })
}

fn parse_grid_polygon_wkt(value: &str) -> Result<[f64; 4], Box<dyn std::error::Error>> {
    let coordinates = value
        .strip_prefix("POLYGON((")
        .and_then(|value| value.strip_suffix("))"))
        .ok_or_else(|| format!("invalid grid polygon WKT {value:?}"))?;
    let mut points = coordinates.split(',');
    let first = points
        .next()
        .ok_or_else(|| format!("grid polygon WKT is missing its first point: {value:?}"))?;
    let third = points
        .nth(1)
        .ok_or_else(|| format!("grid polygon WKT is missing its third point: {value:?}"))?;
    let [minx, miny] = parse_wkt_point(first)?;
    let [maxx, maxy] = parse_wkt_point(third)?;
    Ok([minx, miny, maxx, maxy])
}

fn parse_wkt_point(value: &str) -> Result<[f64; 2], Box<dyn std::error::Error>> {
    let mut parts = value.split_whitespace();
    let x = parts
        .next()
        .ok_or_else(|| format!("invalid WKT point {value:?}"))?
        .parse::<f64>()?;
    let y = parts
        .next()
        .ok_or_else(|| format!("invalid WKT point {value:?}"))?
        .parse::<f64>()?;
    Ok([x, y])
}

/// Returns a tuple of `(CellId, &Cell)` for each cell in column-major order.
impl<'squaregrid> IntoIterator for &'squaregrid SquareGrid {
    type Item = (CellId, &'squaregrid Cell);
    type IntoIter = SquareGridIterator<'squaregrid>;

    fn into_iter(self) -> Self::IntoIter {
        SquareGridIterator {
            row_index: 0,
            col_index: 0,
            items: &self.data,
        }
    }
}

pub struct SquareGridIterator<'squaregrid> {
    row_index: usize,
    col_index: usize,
    items: &'squaregrid Vec<Vec<Cell>>,
}

impl<'squaregrid> Iterator for SquareGridIterator<'squaregrid> {
    type Item = (CellId, &'squaregrid Cell);

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(column) = self.items.get(self.col_index) {
            if let Some(cell) = column.get(self.row_index) {
                let item = Some((
                    CellId {
                        row: self.row_index,
                        column: self.col_index,
                    },
                    cell,
                ));
                self.row_index += 1;
                item
            } else {
                // We are at the end of the current column, so jump to the next
                self.col_index += 1;
                self.row_index = 0;
                self.next()
            }
        } else {
            None
        }
    }
}

#[derive(Debug)]
pub struct SquareGridStats {
    nr_vertices: usize,
    nr_cells_with_content: usize,
    nr_vertices_min: usize,
    nr_vertices_max: usize,
    nr_vertices_mean: f64,
    nr_vertices_median: f64,
}

impl Display for SquareGridStats {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Nr. cells with vertices: {nr_cells}; Nr. vertices: {nr_vertices}, min.: {min}, max.: {max}, median: {median}, mean: {mean}",
            nr_vertices = self.nr_vertices,
            nr_cells = self.nr_cells_with_content,
            min = self.nr_vertices_min,
            max = self.nr_vertices_max,
            median = self.nr_vertices_median,
            mean = self.nr_vertices_mean
        )
    }
}

#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Serialize, Deserialize)]
pub struct Cell {
    pub feature_ids: Vec<usize>,
    pub nr_vertices: usize,
}

/// Grid cell identifier.
/// The identifier is (row, column).
///
/// ```shell
///   ^
/// Y |
///   | row-1     row-1
///   |
///   | row-0     row-0
///   | column-0  column-1
///   +--------------------->
///                       X
/// ```
///
#[derive(Copy, Clone, Hash, Debug, Ord, PartialOrd, PartialEq, Eq, Serialize, Deserialize)]
pub struct CellId {
    // A row is along the y-axis
    pub row: usize,
    // A column is along the x-axis
    pub column: usize,
}

impl Display for CellId {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}-{}", &self.column, &self.row)
    }
}

/// 3D bounding box.
///
/// [min x, min y, min z, max x, max y, max z]
/// TODO: this must become a struct and have a .to_wkt() method
pub type Bbox = [f64; 6];

/// Serialize a 3D bounding box as 2D WKT Polygon.
pub fn bbox_to_wkt(bbox: &Bbox) -> String {
    format!(
        "POLYGON(({minx} {miny}, {maxx} {miny}, {maxx} {maxy}, {minx} {maxy}, {minx} {miny}))",
        minx = bbox[0],
        miny = bbox[1],
        maxx = bbox[3],
        maxy = bbox[4]
    )
}

fn update_bbox(target: &mut Bbox, other: &Bbox) {
    if other[0] < target[0] {
        target[0] = other[0];
    }
    if other[1] < target[1] {
        target[1] = other[1];
    }
    if other[2] < target[2] {
        target[2] = other[2];
    }
    if other[3] > target[3] {
        target[3] = other[3];
    }
    if other[4] > target[4] {
        target[4] = other[4];
    }
    if other[5] > target[5] {
        target[5] = other[5];
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_intersect_bbox() {
        let grid = SquareGrid::new(&[0.0, 0.0, 0.0, 4.0, 4.0, 0.0], 1, 7415);
        let bbox: Bbox = [1.2, 1.2, 0.0, 2.8, 2.8, 0.0];
        let expected = vec![
            CellId { row: 1, column: 1 },
            CellId { row: 2, column: 1 },
            CellId { row: 1, column: 2 },
            CellId { row: 2, column: 2 },
        ];
        let res = grid.intersect_bbox(&bbox);
        assert_eq!(expected.len(), res.len());
        assert!(res.iter().all(|res_cellid| expected.contains(res_cellid)))
    }

    #[test]
    fn test_intersect_bbox_ranges_matches_intersect_bbox() {
        let grid = SquareGrid::new(&[0.0, 0.0, 0.0, 4.0, 4.0, 0.0], 1, 7415);
        let bbox: Bbox = [1.2, 1.2, 0.0, 2.8, 2.8, 0.0];
        let (columns, rows) = grid.intersect_bbox_ranges(&bbox);
        let mut from_ranges = Vec::new();
        for column in columns {
            for row in rows.clone() {
                from_ranges.push(CellId { row, column });
            }
        }

        assert_eq!(grid.intersect_bbox(&bbox), from_ranges);
    }

    #[test]
    fn test_create_grid() {
        let extent = [13603.33, 314127.708, -15.0, 268943.608, 612658.036, 400.0];
        println!("extent: {}", bbox_to_wkt(&extent));
        let grid = SquareGrid::new(&extent, 500, 7415);
        println!("grid: {}", bbox_to_wkt(&grid.bbox));
    }

    #[test]
    fn test_locate_point() {
        let grid = SquareGrid::new(&[0.0, 0.0, 0.0, 4.0, 4.0, 4.0], 1, 0);
        let cellid = grid.locate_point(&[2.5, 1.5]);
        println!("{}", cellid);
        assert_eq!(
            cellid,
            CellId {
                row: 1_usize,
                column: 2_usize,
            }
        );
    }

    #[test]
    fn test_morton_encode() {
        let mut cells_rowwise: Vec<(u64, u64)> = Vec::new();
        for x in 0..4_u64 {
            for y in 0..4u64 {
                cells_rowwise.push((x, y));
            }
        }
        let mut mortoncodes: Vec<u64> = cells_rowwise
            .iter()
            .map(|cell| interleave(&cell.0, &cell.1))
            .collect();
        mortoncodes.sort();
        let cells_morton: Vec<[u64; 2]> = mortoncodes.iter().map(deinterleave).collect();

        let expected: Vec<[u64; 2]> = vec![
            [0, 0],
            [1, 0],
            [0, 1],
            [1, 1],
            [2, 0],
            [3, 0],
            [2, 1],
            [3, 1],
            [0, 2],
            [1, 2],
            [0, 3],
            [1, 3],
            [2, 2],
            [3, 2],
            [2, 3],
            [3, 3],
        ];

        assert_eq!(cells_morton, expected)
    }

    #[test]
    fn test_quadtree_construction() {
        let mut feature_set: FeatureSet = Vec::new();
        let mut grid = SquareGrid::new(&[0.0, 0.0, 0.0, 4.0, 4.0, 1.0], 1, 0);
        for x in 0..4_u64 {
            for y in 0..4u64 {
                for f in 0..5 {
                    feature_set.push(crate::parser::Feature {
                        centroid: [0.0, 0.0],
                        reference: Default::default(),
                        bbox: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                        needs_type_filter: false,
                    });
                    let xc: f64 = format!("{}.{}", &x, &f).parse().unwrap();
                    grid.insert(&[xc, y as f64], f as usize);
                }
            }
        }
        let _ = QuadTree::from_grid(&grid, QuadTreeCapacity::Objects(20));
    }

    #[test]
    fn test_quadtree_leaves() {
        let mut feature_set: FeatureSet = Vec::new();
        let mut grid = SquareGrid::new(&[0.0, 0.0, 0.0, 4.0, 4.0, 1.0], 1, 0);
        for x in 0..4_u64 {
            for y in 0..4u64 {
                for f in 0..5 {
                    feature_set.push(crate::parser::Feature {
                        centroid: [0.0, 0.0],
                        reference: Default::default(),
                        bbox: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                        needs_type_filter: false,
                    });
                    let xc: f64 = format!("{}.{}", &x, &f).parse().unwrap();
                    grid.insert(&[xc, y as f64], f as usize);
                }
            }
        }
        let qtree = QuadTree::from_grid(&grid, QuadTreeCapacity::Objects(20));
        let leaves: Vec<&QuadTree> = QuadTree::collect_leaves(&qtree);
        for tile in leaves {
            println!("{}", tile.id);
        }
    }

    #[test]
    fn test_quadtree_node() {
        let mut feature_set: FeatureSet = Vec::new();
        let mut grid = SquareGrid::new(&[0.0, 0.0, 0.0, 16.0, 16.0, 1.0], 1, 0);
        for x in 0..16_u64 {
            for y in 0..16u64 {
                for f in 0..5 {
                    feature_set.push(crate::parser::Feature {
                        centroid: [0.0, 0.0],
                        reference: Default::default(),
                        bbox: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                        needs_type_filter: false,
                    });
                    let xc: f64 = format!("{}.{}", &x, &f).parse().unwrap();
                    grid.insert(&[xc, y as f64], f as usize);
                }
            }
        }
        let qtree = QuadTree::from_grid(&grid, QuadTreeCapacity::Objects(20));
        let _leaves: Vec<&QuadTree> = QuadTree::collect_leaves(&qtree);
        let n = qtree.node(&QuadTreeNodeId::new(0, 0, 2));
        if n.is_some() {
            println!("{:?}", &n);
        } else {
            println!("did not find node");
        }
    }
}