ED_LRR/rust/src/lib.rs

extern crate strsim;
mod common;
mod preprocess;
mod route;
use common::{System, SystemSerde};
use pyo3::exceptions::*;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use std::collections::HashMap;
use std::path::PathBuf;

fn find_matches(
    path: &PathBuf,
    names: Vec<String>,
) -> Result<HashMap<String, (f64, Option<System>)>, String> {
    let mut best: HashMap<String, (f64, Option<System>)> = HashMap::new();
    for name in &names {
        best.insert(name.to_string(), (0.0, None));
    }
    let mut reader = match csv::ReaderBuilder::new().from_path(path) {
        Ok(rdr) => rdr,
        Err(e) => {
            return Err(format!("Error opening {}: {}", path.to_str().unwrap(), e).to_string());
        }
    };
    let systems = reader.deserialize::<SystemSerde>();
    for sys in systems {
        let sys = sys.unwrap();
        for name in &names {
            best.entry(name.clone()).and_modify(|ent| {
                let d1 = strsim::normalized_levenshtein(&sys.system, &name);
                if d1 > ent.0 {
                    *ent = (d1, Some(sys.build()))
                }
            });
        }
    }
    Ok(best)
}

#[pymodule]
pub fn _ed_lrr(_py: Python, m: &PyModule) -> PyResult<()> {
    /// preprocess(infile_systems, infile_bodies, outfile, callback)
    /// --
    ///
    /// Preprocess bodies.json and systemsWithCoordinates.json into stars.csv
    #[pyfn(m, "preprocess")]
    fn ed_lrr_preprocess(
        py: Python<'static>,
        infile_systems: String,
        infile_bodies: String,
        outfile: String,
        callback: PyObject,
    ) -> PyResult<PyObject> {
        use preprocess::*;
        let state = PyDict::new(py);
        let state_dict = PyDict::new(py);
        callback.call(py, (state_dict,), None).unwrap();
        let callback_wrapped = move |state: &PreprocessState| {
            // println!("SEND: {:?}",state);
            state_dict.set_item("file", state.file.clone())?;
            state_dict.set_item("total", state.total)?;
            state_dict.set_item("count", state.count)?;
            state_dict.set_item("done", state.done)?;
            state_dict.set_item("message", state.message.clone())?;
            callback.call(py, (state_dict,), None)
        };
        preprocess_files(
            &PathBuf::from(infile_bodies),
            &PathBuf::from(infile_systems),
            &PathBuf::from(outfile),
            &callback_wrapped,
        )
        .unwrap();
        Ok(state.to_object(py))
    }

    ///find_sys(sys_names,sys_list)
    /// --
    ///
    /// Find system by name
    #[pyfn(m, "find_sys")]
    fn find_sys(py: Python, sys_names: Vec<String>, sys_list: String) -> PyResult<PyObject> {
        let path = PathBuf::from(sys_list);
        match find_matches(&path, sys_names) {
            Ok(vals) => {
                let ret = PyDict::new(py);
                for (key, (diff, sys)) in vals {
                    let ret_dict = PyDict::new(py);
                    if let Some(val) = sys {
                        let pos = PyList::new(py, val.pos.iter());
                        ret_dict.set_item("star_type", val.star_type.clone())?;
                        ret_dict.set_item("system", val.system.clone())?;
                        ret_dict.set_item("body", val.body.clone())?;
                        ret_dict.set_item("distance", val.distance)?;
                        ret_dict.set_item("pos", pos)?;
                        ret_dict.set_item("id", val.id)?;
                        ret.set_item(key, (diff, ret_dict).to_object(py))?;
                    }
                }
                Ok(ret.to_object(py))
            }
            Err(e) => Err(PyErr::new::<ValueError, _>(e)),
        }
    }

    /// route(infile, hops, range, radius_mult, mode,primary, greedyness, precomp, callback)
    /// --
    ///
    /// Compute a Route using the suplied parameters
    #[pyfn(m, "route")]
    #[allow(clippy::too_many_arguments)]
    fn route(
        py: Python<'static>,
        hops: Vec<String>,
        range: f32,
        prune: Option<(usize, f32)>,
        mode: String,
        primary: bool,
        permute: bool,
        keep_first: bool,
        keep_last: bool,
        greedyness: Option<f32>,
        precomp: Option<String>,
        path: String,
        callback: PyObject,
    ) -> PyResult<PyObject> {
        use route::*;
        let mode = match Mode::parse(&mode) {
            Ok(val) => val,
            Err(e) => {
                return Err(PyErr::new::<ValueError, _>(e));
            }
        };
        let state_dict = PyDict::new(py);
        callback.call(py, (state_dict,), None).unwrap();
        let callback_wrapped = move |state: &SearchState| {
            state_dict.set_item("mode", state.mode.clone())?;
            state_dict.set_item("system", state.system.clone())?;
            state_dict.set_item("body", state.body.clone())?;
            state_dict.set_item("depth", state.depth)?;
            state_dict.set_item("queue_size", state.queue_size)?;
            state_dict.set_item("d_rem", state.d_rem)?;
            state_dict.set_item("d_total", state.d_total)?;
            state_dict.set_item("prc_done", state.prc_done)?;
            state_dict.set_item("n_seen", state.n_seen)?;
            state_dict.set_item("prc_seen", state.prc_seen)?;
            state_dict.set_item("from", state.from.clone())?;
            state_dict.set_item("to", state.to.clone())?;
            callback.call(py, (state_dict,), None)
        };
        let mut systems = Vec::new();
        for sys in hops {
            systems.push(route::SysEntry::parse(&sys))
        }
        let opts = RouteOpts {
            systems,
            range: Some(range),
            file_path: PathBuf::from(path),
            precomp_file: precomp.map(PathBuf::from),
            callback: Box::new(callback_wrapped),
            mode,
            factor: greedyness,
            precompute: false,
            permute,
            keep_first,
            keep_last,
            primary,
            prune,
        };
        let none = ().to_object(py);
        match route(opts) {
            Ok(Some(route)) => {
                let hops = route.iter().map(|hop| {
                    let pos = PyList::new(py, hop.pos.iter());
                    let elem = PyDict::new(py);
                    elem.set_item("star_type", hop.star_type.clone()).unwrap();
                    elem.set_item("system", hop.system.clone()).unwrap();
                    elem.set_item("body", hop.body.clone()).unwrap();
                    elem.set_item("distance", hop.distance).unwrap();
                    elem.set_item("pos", pos).unwrap();
                    elem
                });
                let lst = PyList::new(py, hops);
                Ok(lst.to_object(py))
            }
            Ok(None) => Ok(none),
            Err(e) => Err(PyErr::new::<ValueError, _>(e)),
        }
    }
    Ok(())
}