ED_LRR/rust/src/lib.rs

// #![deny(warnings)]
#![allow(dead_code, clippy::needless_return, clippy::too_many_arguments)]
pub mod common;
pub mod edsm;
pub mod galaxy;
pub mod journal;
pub mod route;
pub mod ship;

extern crate derivative;
use crate::common::{find_matches, SysEntry};
use crate::route::{Router, SearchState};
use crate::ship::Ship;
use pyo3::exceptions::*;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyTuple};
use pyo3::PyObjectProtocol;
use std::path::PathBuf;

#[pyclass(dict)]
#[derive(Debug)]
#[text_signature = "(callback, /)"]
struct PyRouter {
    router: Router,
    primary_only: bool,
    stars_path: String,
}

#[pymethods]
impl PyRouter {
    #[new]
    #[args(callback = "None")]
    fn new(callback: Option<PyObject>, py: Python<'static>) -> PyResult<Self> {
        Ok(PyRouter {
            router: Router::new(Box::new(
                move |state: &SearchState| {
                    match callback.as_ref() {
                        Some(cb) => cb.call(py, (state.clone(),), None),
                        None => Ok(py.None()),
                    }
                }
            )),
            primary_only: false,
            stars_path: String::from(""),
        })
    }

    #[text_signature = "(ship, /)"]
    fn set_ship(&mut self, py: Python, ship: &PyShip) -> PyResult<PyObject> {
        self.router.set_ship(ship.ship.clone());
        Ok(py.None())
    }

    #[args(primary_only = "false")]
    #[text_signature = "(path, primary_only, /)"]
    fn load(
        &mut self,
        path: String,
        primary_only: bool,
        py: Python,
    ) -> PyResult<PyObject> {
        self.stars_path = path;
        self.primary_only = primary_only;
        Ok(py.None())
    }

    #[args(greedyness = "0.5", num_workers = "0", beam_width = "0")]
    #[text_signature = "(hops, range, greedyness, beam_width, num_workers, /)"]
    fn route(
        &mut self,
        hops: &PyList,
        range: Option<f32>,
        greedyness: f32,
        beam_width: usize,
        num_workers: usize,
        py: Python,
    ) -> PyResult<PyObject> {
        let route_res = self
            .router
            .load(&PathBuf::from(self.stars_path.clone()), self.primary_only);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<ValueError, _>(err_msg));
        };
        let mut sys_entries: Vec<SysEntry> = Vec::new();
        for hop in hops {
            if let Ok(id) = hop.extract() {
                sys_entries.push(SysEntry::ID(id));
            } else {
                sys_entries.push(SysEntry::parse(hop.extract()?));
            }
        }
        println!("Resolving systems...");
        let ids: Vec<u32> = match resolve(&sys_entries, &self.router.path) {
            Ok(ids) => ids,
            Err(err_msg) => {
                return Err(PyErr::new::<ValueError, _>(err_msg));
            }
        };
        match self
            .router
            .compute_route(&ids, range, greedyness, beam_width, num_workers)
        {
            Ok(route) => {
                let py_route: Vec<_> = route.iter().map(|hop| hop.to_object(py)).collect();
                Ok(py_route.to_object(py))
            }
            Err(err_msg) => Err(PyErr::new::<RuntimeError, _>(err_msg)),
        }
    }

    #[args(hops = "*")]
    #[text_signature = "(sys_1, sys_2, ..., /)"]
    fn resolve_systems(&self, hops: &PyTuple, py: Python) -> PyResult<PyObject> {
        let mut sys_entries: Vec<SysEntry> = Vec::new();
        for hop in hops {
            if let Ok(id) = hop.extract() {
                sys_entries.push(SysEntry::ID(id));
            } else {
                sys_entries.push(SysEntry::parse(hop.extract()?));
            }
        }
        println!("Resolving systems...");
        let stars_path = PathBuf::from(self.stars_path.clone());
        let ids: Vec<u32> = match resolve(&sys_entries, &stars_path) {
            Ok(ids) => ids,
            Err(err_msg) => {
                return Err(PyErr::new::<ValueError, _>(err_msg));
            }
        };
        let ret: Vec<(_, u32)> = hops.into_iter().zip(ids.into_iter()).collect();
        Ok(PyDict::from_sequence(py, ret.to_object(py))?.to_object(py))
    }

    #[staticmethod]
    fn preprocess_edsm() -> PyResult<()> {
        todo!("Implement EDSM Preprocessor")
    }

    #[staticmethod]
    fn preprocess_galaxy() -> PyResult<()> {
        todo!("Implement galaxy.json Preprocessor")
    }
}

#[pyproto]
impl PyObjectProtocol for PyRouter {
    fn __str__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self))
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self))
    }
}

fn resolve(entries: &[SysEntry], path: &PathBuf) -> Result<Vec<u32>, String> {
    let mut names: Vec<String> = Vec::new();
    let mut ids: Vec<u32> = Vec::new();
    let mut ret: Vec<u32> = Vec::new();
    for ent in entries {
        match ent {
            SysEntry::Name(name) => names.push(name.to_owned()),
            SysEntry::ID(id) => ids.push(*id),
        }
    }
    if !path.exists() {
        return Err(format!(
            "Source file \"{:?}\" does not exist!",
            path.display()
        ));
    }

    let name_ids = find_matches(path, names, false)?;
    for ent in entries {
        match ent {
            SysEntry::Name(name) => {
                let ent_res = name_ids
                    .get(&name.to_owned())
                    .ok_or(format!("System {} not found", name))?;
                let sys = ent_res
                    .1
                    .as_ref()
                    .ok_or(format!("System {} not found", name))?;
                if ent_res.0 < 0.75 {
                    println!(
                        "WARNING: {} match to {} with low confidence ({:.2}%)",
                        name,
                        sys.system,
                        ent_res.0 * 100.0
                    );
                }
                ret.push(sys.id);
            }
            SysEntry::ID(id) => ret.push(*id),
        }
    }
    Ok(ret)
}

#[pyclass(dict)]
#[derive(Debug)]
struct PyShip {
    ship: Ship,
}

#[pyproto]
impl PyObjectProtocol for PyShip {
    fn __str__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self.ship))
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self.ship))
    }
}

#[pymethods]
impl PyShip {
    #[staticmethod]
    fn from_loadout(py: Python, loadout: &str) -> PyResult<PyObject> {
        match Ship::new_from_json(loadout) {
            Ok(ship) => Ok((PyShip { ship }).into_py(py)),
            Err(err_msg) => Err(PyErr::new::<ValueError, _>(err_msg)),
        }
    }
    #[staticmethod]
    fn from_journal(py: Python) -> PyResult<PyObject> {
        let mut ship = match Ship::new_from_journal() {
            Ok(ship) => ship,
            Err(err_msg) => {
                return Err(PyErr::new::<ValueError, _>(err_msg));
            }
        };
        let ships: Vec<(PyObject, PyObject)> = ship
            .drain()
            .map(|(k, v)| {
                let k_py = k.to_object(py);
                let v_py = (PyShip { ship: v }).into_py(py);
                (k_py, v_py)
            })
            .collect();
        Ok(PyDict::from_sequence(py, ships.to_object(py))?.to_object(py))
    }

    fn to_dict(&self, py: Python) -> PyResult<PyObject> {
        self.ship.to_object(py)
    }

    #[text_signature = "(dist, /)"]
    fn fuel_cost(&self, _py: Python, dist: f32) -> PyResult<f32> {
        Ok(self.ship.fuel_cost(dist))
    }

    #[text_signature = "(/)"]
    fn range(&self, _py: Python) -> PyResult<f32> {
        Ok(self.ship.range())
    }

    #[text_signature = "(/)"]
    fn max_range(&self, _py: Python) -> PyResult<f32> {
        Ok(self.ship.max_range())
    }

    #[text_signature = "(dist, /)"]
    fn make_jump(&mut self, dist: f32, _py: Python) -> PyResult<Option<f32>> {
        Ok(self.ship.make_jump(dist))
    }

    #[text_signature = "(dist, /)"]
    fn can_jump(&self, dist: f32, _py: Python) -> PyResult<bool> {
        Ok(self.ship.can_jump(dist))
    }

    #[args(fuel_amount = "None")]
    #[text_signature = "(fuel_amount, /)"]
    fn refuel(&mut self, fuel_amount: Option<f32>, _py: Python) -> PyResult<()> {
        if let Some(fuel) = fuel_amount {
            self.ship.fuel_mass = (self.ship.fuel_mass + fuel).min(self.ship.fuel_capacity)
        } else {
            self.ship.fuel_mass = self.ship.fuel_capacity;
        }
        Ok(())
    }

    #[text_signature = "(factor, /)"]
    fn boost(&mut self, factor: f32, _py: Python) -> PyResult<()> {
        self.ship.boost(factor);
        Ok(())
    }
}

#[pymodule]
pub fn _ed_lrr(_py: Python, m: &PyModule) -> PyResult<()> {
    better_panic::install();

    m.add_class::<PyRouter>()?;
    m.add_class::<PyShip>()?;
    /*
    #[pyfn(m, "get_ships_from_journal")]
    fn get_ships_from_journal(py: Python) -> PyResult<PyObject> {
        let ship = match Ship::new_from_journal() {
            Ok(ship) => ship,
            Err(err_msg) => {
                return Err(PyErr::new::<ValueError, _>(err_msg));
            }
        };
        let ships: Vec<(_,_)> = ship.iter().map(|(k,v)| (k.to_object(py),v.to_object(py))).collect();
        Ok(PyDict::from_sequence(py, ships.to_object(py))?.to_object(py))
    }


    #[pyfn(m, "get_ships_from_loadout")]
    fn get_ship_from_loadout(py: Python, loadout: &str) -> PyResult<PyObject> {
        let ship = match Ship::new_from_json(loadout) {
            Ok(ship) => ship,
            Err(err_msg) => {
                return Err(PyErr::new::<ValueError, _>(err_msg));
            }
        };
        Ok(ship.to_object(py))
    }
    */
    Ok(())
}

/*

/// Preprocess bodies.json and systemsWithCoordinates.json into stars.csv
    #[pyfn(m, "preprocess")]
    #[text_signature = "(infile_systems, infile_bodies, outfile, callback, /)"]
    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 system by name
    #[pyfn(m, "find_sys")]
    #[text_signature = "(sys_names, sys_list_path, /)"]
    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, false) {
            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)),
        }
    }

    /// Compute a Route using the suplied parameters
    #[pyfn(m, "route")]
    #[text_signature = "(hops, range, mode, primary, permute, keep_first, keep_last, greedyness, precomp, path, num_workers, callback, /)"]
    #[allow(clippy::too_many_arguments)]
    fn py_route(
        py: Python<'static>,
        hops: Vec<&str>,
        range: f32,
        mode: String,
        primary: bool,
        permute: bool,
        keep_first: bool,
        keep_last: bool,
        greedyness: Option<f32>,
        precomp: Option<String>,
        path: String,
        num_workers: Option<usize>,
        callback: PyObject,
    ) -> PyResult<PyObject> {
        use route::*;
        let num_workers = num_workers.unwrap_or(1);
        let mode = match Mode::parse(&mode) {
            Ok(val) => val,
            Err(e) => {
                return Err(PyErr::new::<ValueError, _>(e));
            }
        };
        let state_dict = PyDict::new(py);
        {
            let cb_res = callback.call(py, (state_dict,), None);
            if cb_res.is_err() {
                println!("Error: {:?}", cb_res);
            }
        }
        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())?;
            let cb_res = callback.call(py, (state_dict,), None);
            if cb_res.is_err() {
                println!("Error: {:?}", cb_res);
            }
            cb_res
        };
        let hops: Vec<SysEntry> = (hops.iter().map(|v| SysEntry::from_str(&v)).collect::<Result<Vec<SysEntry>,_>>())?;
        println!("Resolving systems...");
        let hops: Vec<u32> = match resolve(&hops, &PathBuf::from(&path)) {
            Ok(ids) => ids,
            Err(err_msg) => {
                return Err(PyErr::new::<ValueError, _>(err_msg));
            }
        };
        let opts = RouteOpts {
            systems: hops,
            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,
            workers: num_workers,
        };
        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(py.None()),
            Err(e) => Err(PyErr::new::<ValueError, _>(e)),
        }
    }

*/