tilemake_rs 0.1.2

use geozero::mvt::tile::{Feature, Layer};
use geozero::mvt::{Message, Tile};
use mlua::Lua;
use osmpbfreader::OsmObj::Node;
use osmpbfreader::{NodeId, OsmId, OsmObj, Tags};
use pmtiles::{PmTilesWriter, TileCoord, TileType};
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::f64::consts::PI;
use std::rc::Rc;

mod config;
mod input;
mod lua;

pub use config::Config;
pub use input::InputArguments;
use lua::{LuaProcess, LuaState, OsmNode, call_node_function, load_lua_script};

static EXTENT: u32 = 4096;

fn zigzag(n: i32) -> u32 {
    ((n << 1) ^ (n >> 31)) as u32
}

#[derive(Debug, PartialEq)]
struct TilePointInput {
    lat: f64,
    lon: f64,
    zoom: Zoom,
    extent: u32,
}

impl TilePointInput {
    fn new_from_decimicro(lat: i32, lon: i32, zoom: Zoom, extent: u32) -> Self {
        Self {
            lat: lat as f64 / 10_000_000.0,
            lon: lon as f64 / 10_000_000.0,
            zoom,
            extent,
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
struct TilePoint {
    pub tile_x: u32,
    pub tile_y: u32,
    pub x: u32,
    pub y: u32,
}

impl TilePoint {
    fn new(tile_x: u32, tile_y: u32, x: u32, y: u32) -> Self {
        Self {
            tile_x,
            tile_y,
            x,
            y,
        }
    }
}

fn lonlat_to_tile_and_pixel(input: TilePointInput) -> TilePoint {
    let n = 2f64.powi(input.zoom as i32);

    let x = (input.lon + 180.0) / 360.0 * n;

    let lat_rad = input.lat.to_radians();
    let y = (1.0 - ((lat_rad.tan() + 1.0 / lat_rad.cos()).ln() / PI)) / 2.0 * n;

    let tile_x = x.floor() as u32;
    let tile_y = y.floor() as u32;

    let pixel_x = ((x - tile_x as f64) * input.extent as f64) as u32;
    let pixel_y = ((y - tile_y as f64) * input.extent as f64) as u32;

    TilePoint::new(tile_x, tile_y, pixel_x, pixel_y)
}

type NodeMap = HashMap<NodeId, TilePoint>;

type TileLayers = HashMap<String, Layer>;

type ZoomTiles = HashMap<TileCoord, TileLayers>;

type Zoom = u8;

struct NodeZoomLevelMap(HashMap<Zoom, NodeMap>);

impl NodeZoomLevelMap {
    fn new() -> Self {
        Self(HashMap::new())
    }

    fn insert(&mut self, zoom: Zoom, node_id: NodeId, tile_point: TilePoint) {
        match self.0.get_mut(&zoom) {
            Some(node_map) => {
                node_map.insert(node_id, tile_point);
            }
            None => {
                let mut node_map = HashMap::new();
                node_map.insert(node_id, tile_point);
                self.0.insert(zoom, node_map);
            }
        };
    }

    fn find(&self, z: &Zoom, id: &NodeId) -> &TilePoint {
        let zoom_map = self.0.get(z).unwrap();
        zoom_map.get(id).unwrap()
    }
}

struct NodeZoomLayerMap(HashMap<Zoom, ZoomTiles /*NodeLayerMap*/>);

impl NodeZoomLayerMap {
    fn new() -> Self {
        Self(HashMap::new())
    }

    fn insert(&mut self, zoom: Zoom, layer_name: String, point: &TilePoint) {
        let layer = self
            .0
            .entry(zoom.clone())
            .or_default()
            .entry(TileCoord::new(zoom, point.tile_x, point.tile_y).unwrap())
            .or_default()
            .entry(layer_name.clone())
            .or_insert_with(|| {
                let mut layer = Layer::default();
                layer.name = layer_name;
                layer.version = 2;
                layer.extent = Some(4096);
                layer
            });
        let mut feature = Feature::default();
        feature.id = Some(1u64);
        feature.r#type = Some(1); // POINT
        let move_to = 1 | (1 << 3);

        // Convert u32 to i32 before calling zigzag
        let x_i32 = point.x as i32;
        let y_i32 = point.y as i32;

        feature.geometry = vec![move_to, zigzag(x_i32), zigzag(y_i32)];

        layer.features.push(feature);
    }
}

fn osm_obj_tags_to_hashmap(tags: &Tags) -> HashMap<String, String> {
    let mut tag_map = HashMap::new();
    for (key, value) in tags.iter() {
        tag_map.insert(key.to_string().clone(), value.to_string().clone());
    }
    tag_map
}

fn build_node_map(
    input_arguments: &InputArguments,
    objs: &BTreeMap<OsmId, OsmObj>,
) -> NodeZoomLevelMap {
    let mut node_map = NodeZoomLevelMap::new();
    for zoom in input_arguments.zoom_iter() {
        println!("z {:?}", &zoom);
        for (_id, obj) in objs {
            if let Node(node) = obj {
                // What is more efficient?
                // - insert decimicro and only convert if node is in layer
                // - convert all nodes here
                let tile_point = lonlat_to_tile_and_pixel(TilePointInput::new_from_decimicro(
                    node.decimicro_lat,
                    node.decimicro_lon,
                    zoom.clone(),
                    EXTENT,
                ));
                node_map.insert(zoom.clone(), node.id, tile_point);
            }
        }
    }
    node_map
}

fn process_nodes(
    input_arguments: &InputArguments,
    node_map: &NodeZoomLevelMap,
    lua_process: &LuaProcess,
    objs: &BTreeMap<OsmId, OsmObj>,
    layered_nodes: Rc<RefCell<NodeZoomLayerMap>>,
) {
    for zoom in input_arguments.zoom_iter() {
        for (_id, obj) in objs {
            if let Node(node) = obj {
                let tile_point = node_map.find(&zoom, &node.id);
                match lua_process {
                    LuaProcess::Some(lua_state) => {
                        let _res = call_node_function(
                            &lua_state.lua,
                            &lua_state.node_function,
                            OsmNode::new(
                                node.id.0,
                                zoom.clone(),
                                tile_point.clone(),
                                osm_obj_tags_to_hashmap(&obj.tags()),
                            ),
                            layered_nodes.clone(),
                        );
                    }
                    LuaProcess::None => {
                        layered_nodes.clone().borrow_mut().insert(
                            zoom,
                            "default".to_string(),
                            tile_point,
                        );
                    }
                }
            }
        }
    }
}

fn get_lua_process(input_arguments: &InputArguments) -> LuaProcess {
    match &input_arguments.process {
        Some(code) => {
            let lua = Lua::new();
            let node_function = load_lua_script(&lua, &code).unwrap();
            LuaProcess::Some(LuaState {
                lua: lua,
                node_function: node_function,
            })
        }
        None => LuaProcess::None,
    }
}

/// Runs the processing pipeline based on the input arguments
///
/// ```
/// # #[cfg(feature = "serde")]
/// # {
/// use tilemake_rs::{Config, InputArguments, process};
/// use std::path::Path;
/// use std::fs::File;
/// use std::fs;
///
/// let input_file = File::open(Path::new("tests/data/azores-260318-power.osm.pbf")).unwrap();
/// let output_file = File::create(Path::new("tests/data/output.pmtiles")).unwrap();
/// let mut input_arguments = InputArguments::new(input_file, output_file);
///
/// let config_string = fs::read_to_string(Path::new("tests/data/azores-config.json")).unwrap();
/// let config = Config::from_json(&config_string).unwrap();
/// input_arguments.set_config(Some(config));
///
/// let process_string = fs::read_to_string(Path::new("tests/data/azores.lua")).unwrap();
/// input_arguments.set_process(Some(process_string));
///
/// process(&input_arguments);
/// # }
/// ```
pub fn process(input_arguments: &InputArguments) {
    let mut writer = PmTilesWriter::new(TileType::Mvt)
        .create(&input_arguments.output_file)
        .unwrap();

    let mut pbf = osmpbfreader::OsmPbfReader::new(&input_arguments.input_file);
    let objs = pbf.get_objs_and_deps(|_obj| true).unwrap();
    let node_map = build_node_map(input_arguments, &objs);
    let lua_process = get_lua_process(input_arguments);

    let layered_nodes = Rc::new(RefCell::new(NodeZoomLayerMap::new()));
    process_nodes(
        input_arguments,
        &node_map,
        &lua_process,
        &objs,
        layered_nodes.clone(),
    );

    let mut set = HashSet::new();

    for (_zoom, zoom_tiles) in layered_nodes.clone().borrow().0.iter() {
        for (tile, layers) in zoom_tiles {
            for (layer_name, layer) in layers {
                set.insert(layer_name.clone());
                let mut tile_thing = Tile::default();
                tile_thing.layers.push(layer.clone());
                let tile_bytes = tile_thing.encode_to_vec();
                writer.add_tile(*tile, &tile_bytes).unwrap();
            }
        }
    }

    println!("Write layer: {:?}", set);
    writer.finalize().unwrap();
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;
    use std::fs;
    use std::fs::File;
    use std::path::Path;

    #[rstest]
    #[case(
        TilePointInput::new_from_decimicro(0, 0, 0, 4096),
        TilePoint::new(0, 0, 2048, 2048)
    )]
    #[case(
        TilePointInput::new_from_decimicro(0, 0, 1, 4096),
        TilePoint::new(1, 1, 0, 0)
    )]
    #[case(
        TilePointInput::new_from_decimicro(0, 0, 2, 4096),
        TilePoint::new(2, 2, 0, 0)
    )]
    #[case(
        TilePointInput::new_from_decimicro(0, 1_800_000_000, 0, 4096),
        TilePoint::new(1, 0, 0, 2048)
    )]
    #[case(
        TilePointInput::new_from_decimicro(437452160, 74283760, 0, 4096),
        TilePoint::new(0, 0, 2132, 1493)
    )]
    fn point_micro_test(#[case] input: TilePointInput, #[case] expected: TilePoint) {
        assert_eq!(expected, lonlat_to_tile_and_pixel(input))
    }

    #[rstest]
    #[case("azores.lua", 3000)]
    #[case("azores-generator.lua", 140)]
    #[cfg(feature = "serde")]
    fn test_with_lua_file(#[case] process_file: String, #[case] features: usize) {
        let input_file = File::open(Path::new("tests/data/azores-260318-power.osm.pbf")).unwrap();
        let output_file = File::create(Path::new("tests/data/output.pmtiles")).unwrap();
        let mut input_arguments = InputArguments::new(input_file, output_file);

        let config_string = fs::read_to_string(Path::new("tests/data/azores-config.json")).unwrap();
        let config = Config::from_json(&config_string).unwrap();
        input_arguments.set_config(Some(config));

        let process_string =
            fs::read_to_string(Path::new(&format!("tests/data/{}", process_file))).unwrap();
        input_arguments.set_process(Some(process_string));
        let mut pbf = osmpbfreader::OsmPbfReader::new(&input_arguments.input_file);
        let objs = pbf.get_objs_and_deps(|_obj| true).unwrap();
        let node_map = build_node_map(&input_arguments, &objs);
        assert_eq!(node_map.0.len(), 2);
        assert_eq!(node_map.0[&0].len(), 3_000);
        assert_eq!(node_map.0[&1].len(), 3_000);

        let lua_process = get_lua_process(&input_arguments);
        let layered_nodes = Rc::new(RefCell::new(NodeZoomLayerMap::new()));
        process_nodes(
            &input_arguments,
            &node_map,
            &lua_process,
            &objs,
            layered_nodes.clone(),
        );
        assert_eq!(layered_nodes.clone().borrow().0.len(), 2);
        let zoom_level_0 = layered_nodes.clone().borrow().0[&0].clone();
        assert_eq!(zoom_level_0.len(), 1);
        let keys: Vec<_> = zoom_level_0.keys().cloned().collect();
        assert_eq!(keys.len(), 1);
        assert_eq!(keys[0], TileCoord::new(0, 0, 0).unwrap());
        let power_layer_0 = &zoom_level_0[&TileCoord::new(0, 0, 0).unwrap()].clone();
        let power_layer_0_keys: Vec<_> = power_layer_0.keys().cloned().collect();
        assert_eq!(power_layer_0_keys.len(), 1);
        assert_eq!(power_layer_0_keys[0], "power");
        assert_eq!(power_layer_0["power"].keys.len(), 0);
        assert_eq!(power_layer_0["power"].features.len(), features);
        assert_eq!(power_layer_0["power"].name, "power");
        assert_eq!(power_layer_0["power"].values.len(), 0);
    }

    #[test]
    #[cfg(feature = "serde")]
    fn test_without_lua_file() {
        let input_file = File::open(Path::new("tests/data/azores-260318-power.osm.pbf")).unwrap();
        let output_file = File::create(Path::new("tests/data/output.pmtiles")).unwrap();
        let mut input_arguments = InputArguments::new(input_file, output_file);

        let config_string = fs::read_to_string(Path::new("tests/data/azores-config.json")).unwrap();
        let config = Config::from_json(&config_string).unwrap();
        input_arguments.set_config(Some(config));

        let mut pbf = osmpbfreader::OsmPbfReader::new(&input_arguments.input_file);
        let objs = pbf.get_objs_and_deps(|_obj| true).unwrap();
        let node_map = build_node_map(&input_arguments, &objs);
        assert_eq!(node_map.0.len(), 2);
        assert_eq!(node_map.0[&0].len(), 3_000);
        assert_eq!(node_map.0[&1].len(), 3_000);

        let lua_process = get_lua_process(&input_arguments);
        let layered_nodes = Rc::new(RefCell::new(NodeZoomLayerMap::new()));
        process_nodes(
            &input_arguments,
            &node_map,
            &lua_process,
            &objs,
            layered_nodes.clone(),
        );
        assert_eq!(layered_nodes.clone().borrow().0.len(), 2);
        let zoom_level_0 = layered_nodes.clone().borrow().0[&0].clone();
        assert_eq!(zoom_level_0.len(), 1);
        let keys: Vec<_> = zoom_level_0.keys().cloned().collect();
        assert_eq!(keys.len(), 1);
        assert_eq!(keys[0], TileCoord::new(0, 0, 0).unwrap());
        let power_layer_0 = &zoom_level_0[&TileCoord::new(0, 0, 0).unwrap()].clone();
        let power_layer_0_keys: Vec<_> = power_layer_0.keys().cloned().collect();
        assert_eq!(power_layer_0_keys.len(), 1);
        assert_eq!(power_layer_0_keys[0], "default");
        assert_eq!(power_layer_0["default"].keys.len(), 0);
        assert_eq!(power_layer_0["default"].features.len(), 3000);
        assert_eq!(power_layer_0["default"].name, "default");
        assert_eq!(power_layer_0["default"].values.len(), 0);
    }
}