nadi_core 0.8.1

use nadi_plugin::nadi_internal_plugin;

#[nadi_internal_plugin]
mod attrs {
    use crate::attrs::TableOrArray;
    use crate::functions::FunctionInput;
    use crate::prelude::*;
    use abi_stable::std_types::Tuple2;
    use abi_stable::std_types::{RHashMap, RString};
    use nadi_plugin::{env_func, network_func, node_func};
    use std::collections::HashMap;
    use std::str::FromStr;

    /// None function that returns None
    #[env_func]
    fn none() {}

    /// check for none value
    ///
    /// ```task
    /// assert_eq(is_none(none()), true)
    /// ```
    #[env_func]
    fn is_none(inp: Option<FunctionInput>) -> bool {
        inp.is_none()
    }

    /// Set node attributes
    ///
    /// Use this function to set the node attributes of all nodes, or
    /// a select few nodes using the node selection methods (path or
    /// list of nodes)
    ///
    /// # Error
    /// The function should not error.
    ///
    /// # Example
    /// Following will set the attribute `a2d` to `true` for all nodes
    /// from `A` to `D`
    ///
    /// ```task
    /// network load_str("A -> B\n B -> D");
    /// nodes[A -> D] set_attrs(a2d = true)
    /// ```
    /// This is equivalent to the following:
    /// ```task
    /// nodes[A->D].a2d = true;
    /// ```
    #[node_func]
    fn set_attrs(
        node: &mut NodeInner,
        /// Key value pairs of the attributes to set
        #[kwargs]
        attrs: RHashMap<&RString, Attribute>,
    ) -> Result<(), String> {
        for Tuple2(k, v) in attrs {
            node.set_attr(k.as_str(), v);
        }
        Ok(())
    }

    /// Delete attributes from the given node
    ///
    /// ```task
    /// network load_str("a -> b");
    /// nodes set_attrs(val = true);
    /// node[a] del_attrs(["val"]);
    /// node[a] assert_eq(val?, false)
    /// node[b] assert_eq(val?, true)
    /// ```
    #[node_func]
    fn del_attrs(
        node: &mut NodeInner,
        /// the attributes to delete
        delete: Vec<String>,
    ) -> Result<(), String> {
        for attr in delete {
            node.attr_map_mut().remove(attr.as_str());
        }
        Ok(())
    }

    /// Retrive attribute
    ///
    /// ```task
    /// network load_str("A -> B\n B -> D");
    /// nodes assert_eq(get_attr("NAME"), NAME);
    /// ```
    #[node_func]
    fn get_attr(
        node: &NodeInner,
        /// Name of the attribute to get
        attr: &str,
        /// Default value if the attribute is not found
        default: Option<Attribute>,
    ) -> Option<Attribute> {
        node.attr(attr).cloned().or(default)
    }

    /// Retrive multiple attributes
    ///
    /// ```task
    /// network load_str("A -> B\n B -> D");
    /// nodes assert_eq(get_attrs("NAME"), array(NAME));
    /// nodes assert_eq(get_attrs("NAME", "ORDER"), array(NAME, ORDER));
    /// ```
    #[node_func]
    fn get_attrs(
        node: &NodeInner,
        /// Name of the attribute to get
        #[args]
        attr_names: Vec<String>,
    ) -> Option<Vec<Attribute>> {
        attr_names.iter().map(|a| node.attr(a).cloned()).collect()
    }

    /// Check if the attribute is present
    ///
    /// ```task
    /// network load_str("A -> B\n B -> D");
    /// nodes.x = 90;
    /// nodes assert(has_attr("x"))
    /// nodes assert(!has_attr("y"))
    /// ```
    #[node_func]
    fn has_attr(
        node: &NodeInner,
        /// Name of the attribute to check
        attr: &str,
    ) -> bool {
        node.attr(attr).is_some()
    }

    /// Return the first Attribute that exists
    ///
    /// This is useful when you have a bunch of attributes that might
    /// be equivalent but are using different names. Normally due to
    /// them being combined from different datasets.
    ///
    /// ```task
    /// network load_str("A -> B\n B -> D");
    /// nodes.x = 90;
    /// nodes assert_eq(first_attr(["y", "x"]), 90)
    /// nodes assert_eq(first_attr(["x", "NAME"]), 90)
    /// ```
    #[node_func]
    fn first_attr(
        node: &NodeInner,
        /// attribute names
        attrs: Vec<RString>,
        /// Default value if not found
        default: Option<Attribute>,
    ) -> Option<Attribute> {
        for attr in attrs {
            if let Ok(Some(v)) = node.attr_dot(attr.as_str()) {
                return Some(v.clone());
            }
        }
        default
    }

    /// map values from the attribute based on the given table
    ///
    /// ```task
    /// env.val = strmap("Joe", {Dave = 2, Joe = 20});
    /// env assert_eq(val, 20)
    /// env.val2 = strmap("Joe", {Dave=2}, default = 12);
    /// env assert_eq(val2, 12)
    /// ```
    #[env_func]
    fn strmap(
        /// Value to transform the attribute
        #[relaxed]
        attr: String,
        /// Dictionary of key=value to map the data to
        attrmap: &AttrMap,
        /// Default value if key not found in `attrmap`
        default: Option<Attribute>,
    ) -> Option<Attribute> {
        attrmap.get(attr.as_str()).cloned().or(default)
    }

    /// if else condition with multiple attributes
    ///
    /// ```task
    /// network load_str("a -> b");
    /// env.some_condition = true;
    /// nodes set_attrs_ifelse(
    ///     env.some_condition,
    ///     val1 = [1, 2],
    ///     val2 = ["a", "b"]
    /// );
    /// env assert_eq(nodes.val1, [1, 1])
    /// env assert_eq(nodes.val2, ["a", "a"])
    /// ```
    /// This is equivalent to using the if-else expression directly,
    ///
    /// ```task
    /// nodes.val1 = if (env.some_condition) {1} else {2};
    /// env assert_eq(nodes.val1, [1, 1])
    /// ```
    ///
    /// Furthermore if-else expression will give a lot more
    /// flexibility than this function in normal use cases. But this
    /// function is useful when you have to do something in a batch.
    #[node_func]
    fn set_attrs_ifelse(
        node: &mut NodeInner,
        /// Condition to check
        #[relaxed]
        cond: bool,
        /// key = [val1, val2] where key is set as first if `cond` is true else second
        #[kwargs]
        values: HashMap<&RString, (Attribute, Attribute)>,
    ) -> Result<(), String> {
        for (k, (t, f)) in values {
            let v = if cond { t } else { f };
            node.set_attr(k, v);
        }
        Ok(())
    }

    /// Set node attributes based on string templates
    ///
    /// This renders the template for each node, then it sets the
    /// values from the rendered results.
    ///
    /// ```task
    /// network load_str("a -> b");
    /// nodes set_attrs_render(val1 = "Node: {NAME}");
    /// node[a] assert_eq(val1, "Node: a")
    /// ```
    #[node_func]
    fn set_attrs_render(
        node: &mut NodeInner,
        /// key value pair of attribute to set and the Template to render
        #[kwargs]
        values: HashMap<&RString, Template>,
    ) -> Result<(), String> {
        for (k, templ) in values {
            let text = templ.render(node).map_err(|e| e.to_string())?;
            node.set_attr(k.as_str(), text.into());
        }
        Ok(())
    }

    /// Set node attributes by loading a toml from rendered template
    ///
    /// This function will render a string, and loads it as a toml
    /// string. This is useful when you need to make attributes based
    /// on some other variables that you can combine using the string
    /// template system.
    ///
    /// In most cases it is better to use the string manipulation
    /// functions and other environmental functions to get new
    /// attribute values to set.
    ///
    /// ```task
    /// network load_str("a -> b");
    /// nodes load_toml_render("label = \"Node: {NAME}\"")
    /// nodes assert_eq(label, render("Node: {NAME}"))
    /// ```
    #[node_func(echo = false)]
    fn load_toml_render(
        node: &mut NodeInner,
        /// String template to render and load as toml string
        toml: Template,
        /// Print the rendered toml or not
        echo: bool,
    ) -> anyhow::Result<()> {
        let toml = format!("{}\n", toml.render(node)?);
        if echo {
            println!("{toml}");
        }
        let tokens = crate::parser::tokenizer::get_tokens(&toml);
        let attrs = crate::parser::attrs::parse(tokens)?;
        node.attr_map_mut().extend(attrs);
        Ok(())
    }

    /// Parse attribute from string
    ///
    /// ```task
    /// env assert_eq(parse_attr("true"), true)
    /// env assert_eq(parse_attr("123"), 123)
    /// env assert_eq(parse_attr("12.34"), 12.34)
    /// env assert_eq(parse_attr("\"my value\""), "my value")
    /// env assert_eq(parse_attr("1234-12-12 00:00"), 1234-12-12T00:00)
    /// ```
    #[env_func]
    fn parse_attr(
        /// String to parse into attribute
        toml: &str,
    ) -> Result<Attribute, String> {
        Attribute::from_str(toml).map_err(|e| e.to_string())
    }

    /// Parse attribute map from string
    ///
    /// ```task
    /// env assert_eq(parse_attrmap("y = true"), {y = true})
    /// env assert_eq(parse_attrmap(
    ///     "x = [1234, true]"),
    ///     {x = [1234, true]}
    /// )
    /// ```
    #[env_func]
    fn parse_attrmap(
        /// String to parse into attribute
        toml: String,
    ) -> Result<AttrMap, String> {
        let tokens = crate::parser::tokenizer::get_tokens(&toml);
        let attrs = crate::parser::attrs::parse(tokens).map_err(|e| e.to_string())?;
        Ok(attrs)
    }

    /// keys of the attribute map
    #[env_func]
    fn keys(attrmap: AttrMap) -> Vec<String> {
        attrmap.keys().map(ToString::to_string).collect()
    }

    /// values of the attribute map
    #[env_func]
    fn values(attrmap: AttrMap) -> Vec<Attribute> {
        attrmap.values().cloned().collect()
    }

    /// get the choosen attribute from Array or AttrMap
    ///
    /// ```task
    /// env.some_ar = ["this", 12, true];
    /// env.some_am = {x = "this", y = [12, true]};
    /// env assert_eq(get(some_ar, 0), "this")
    /// env assert_eq(get(some_ar, 2), true)
    /// env assert_eq(get(some_am, "x"), "this")
    /// env assert_eq(get(some_am, "y"), [12, true])
    /// ```
    #[env_func]
    fn get(
        /// Array or AttrMap Attribute to index
        parent: Attribute,
        /// Index value (Integer for Array, String for AttrMap)
        index: Attribute,
        /// Default value if the index is not present
        default: Option<Attribute>,
    ) -> Result<Attribute, String> {
        match (parent, index) {
            (Attribute::Array(ar), Attribute::Integer(ind)) => ar
                .get(ind as usize)
                .cloned()
                .or(default)
                .ok_or(format!("Index {ind} not found")),
            (Attribute::Table(am), Attribute::String(key)) => am
                .get(&key)
                .cloned()
                .or(default)
                .ok_or(format!("Index {key} not found")),
            (Attribute::Array(_), b) => Err(format!(
                "Array index should be Integer not {}",
                b.type_name()
            )),
            (Attribute::Table(_), b) => Err(format!(
                "AttrMap index should be String not {}",
                b.type_name()
            )),
            (a, _) => Err(format!("{} cannot be indexed", a.type_name())),
        }
    }
    /// Float Division (same as / operator)
    ///
    /// ```task
    /// env assert_eq(float_div(10.0, 2), 10.0 / 2)
    /// ```
    #[env_func]
    fn float_div(
        /// numerator
        #[relaxed]
        value1: f64,
        /// denominator
        #[relaxed]
        value2: f64,
    ) -> f64 {
        value1 / value2
    }

    /// Float Multiplication (same as * operator)
    ///
    /// ```task
    /// env assert_eq(float_mult(5.0, 2), 5.0 * 2)
    /// ```
    #[env_func]
    fn float_mult(
        /// numerator
        #[relaxed]
        value1: f64,
        /// denominator
        #[relaxed]
        value2: f64,
    ) -> f64 {
        value1 * value2
    }

    /// Set network attributes
    ///
    /// # Arguments
    /// - `key=value` - Kwargs of attr = value
    ///
    /// ```task
    /// network set_attrs(val = 23.4)
    /// network assert_eq(val, 23.4)
    /// ```
    #[network_func]
    fn set_attrs(
        network: &mut Network,
        /// key value pair of attributes to set
        #[kwargs]
        values: HashMap<&RString, Attribute>,
    ) -> Result<(), String> {
        for (k, v) in values {
            network.set_attr(k.as_str(), v.clone());
        }
        Ok(())
    }

    /// Set node attributes in a network using a attrmap or array
    ///
    /// Currenly you can only set all nodes using array, if you want
    /// to set a subset of the nodes, use the attrmap option.
    ///
    /// ```task
    /// network load_str("a -> b")
    /// network set_node_attrs("val", {a = 23.4})
    /// network assert_eq(node[a].val, 23.4)
    /// network set_node_attrs("val", [2, 4])
    /// network assert_eq(nodes.val, [2, 4])
    /// ```
    #[network_func]
    fn set_node_attrs(
        network: &mut Network,
        /// Name of the attribute to set,
        attr_name: &str,
        /// array or a key value pair of attributes to set (key = node name)
        node_values: TableOrArray,
    ) -> Result<(), String> {
        match node_values {
            TableOrArray::Table(node_map) => {
                for Tuple2(k, v) in node_map.into_iter() {
                    network
                        .node_by_name(&k)
                        .ok_or(format!("Node {k} not found"))?
                        .lock()
                        .set_attr(attr_name, v);
                }
            }
            TableOrArray::Array(nodes_vals) => {
                network.nodes().zip(nodes_vals).for_each(|(n, v)| {
                    n.lock().set_attr(attr_name, v);
                });
            }
        }
        Ok(())
    }

    /// Set network attributes based on string templates
    ///
    /// It will set the attribute as a String
    ///
    /// ```task
    /// network.val = 23.4
    /// network set_attrs_render(val2 = "{val}05")
    /// network assert_eq(val2, "23.405")
    /// ```
    #[network_func]
    fn set_attrs_render(
        network: &mut Network,
        /// Kwargs of attr = String template to render
        #[kwargs]
        kwargs: HashMap<&RString, Template>,
    ) -> Result<(), String> {
        for (k, templ) in kwargs {
            let text = templ.render(network).map_err(|e| e.to_string())?;
            network.set_attr(k.as_str(), text.into());
        }
        Ok(())
    }

    fn node_attr(n: &Node, attr: &str) -> (RString, Option<Attribute>) {
        let n = n.lock();
        let a = n.attr(attr);
        (n.name().to_string().into(), a.cloned())
    }

    /// Generate attribute map for the given attribute for the nodes
    #[network_func(safe = false)]
    fn nodemap(
        net: &Network,
        /// attribute to be the value of the attrmap
        attr: String,
        /// Only include these nodes
        filter: Option<Vec<bool>>,
        /// Exclude nodes if they do not have the attribute
        safe: bool,
    ) -> Result<AttrMap, String> {
        let nodes = if let Some(filter) = filter {
            net.nodes()
                .zip(filter)
                .filter(|(_, f)| *f)
                .map(|(n, _)| n)
                .collect::<Vec<&Node>>()
        } else {
            net.nodes().collect::<Vec<&Node>>()
        };
        let mut amap = AttrMap::with_capacity(nodes.len());
        for n in nodes {
            let (name, at) = node_attr(n, &attr);
            match at {
                Some(a) => {
                    amap.insert(name, a);
                }
                None if safe => (),
                None => return Err(format!("Node {name:?} does not have {attr:?} attribute")),
            }
        }
        Ok(amap)
    }
}