progress on solve_for function

1 files changed, 65 insertions, 0 deletions

diff --git a/src/modelica_ast.rs b/src/modelica_ast.rs
index f97cd0a..2958025 100644
--- a/src/modelica_ast.rs
+++ b/src/modelica_ast.rs
@@ -2,6 +2,7 @@
 use std::clone::Clone;
 use std::fmt::{Debug, Formatter, Error};
 use std::collections::HashMap;
+use std::collections::HashSet;
 
 #[derive(Clone, PartialEq)]
 pub struct ModelicaModel {
@@ -105,6 +106,70 @@ impl ModelicaModel {
         vars.map(|c| c.name.clone()).collect()
     }
 
+    // V variables
+    // Q constants (become kwargs)
+    // P bound vars (independent, inputs/passed, become like constants)
+    // M unknowns (dependent, outputs)
+    // N' total equations
+    // N equations with unknowns
+    //
+    // TODO: allow passing in Q
+    pub fn solve_for(&self, indep_vars: Vec<String>, dep_vars: Vec<String>) -> Result<ModelicaModel,String> {
+        let constants = self.get_constant_vars();
+        let param_objs = self.components.iter().filter(|v| match v.prefix {
+                Some(ComponentPrefix::Constant) | Some(ComponentPrefix::Parameter) => false,
+                _ => true,
+            });
+        let params: Vec<String> = param_objs.map(|c| c.name.clone()).collect();
+
+        let mut all_vars: HashSet<String> = HashSet::new();
+        for eqn in &self.equations {
+            all_vars.extend(eqn.identifiers());
+        }
+
+        // check that all dep and indep are in equations
+        let mut passed_vars = indep_vars.clone();
+        passed_vars.extend(dep_vars.clone());
+        for var in passed_vars {
+            if !all_vars.contains(&var) {
+                return Err(format!("Variable not found in equations: {}", var));
+            }
+        }
+
+        // check that  V = Q + P + M
+        if all_vars.len() != (constants.len() + indep_vars.len() + dep_vars.len()) {
+            return Err("Variable counts don't add up".to_string());
+        }
+
+        // check that all constants are bound and simple
+        for (name, value) in &constants {
+            match *value {
+                None => return Err(format!("UnderSpecifiedConstant: {}", name)),
+                Some(Expr::Integer(_)) | Some(Expr::Float(_)) => (), // Ok,
+                Some(_) => return Err(format!("NaiveImplementation: can't handle constant: {}", name)),
+            }
+        }
+
+        // check that there is a depdendent variable in each equation
+        for eqn in &self.equations {
+            if intersect_strings(&dep_vars, &eqn.identifiers()).len() == 0 {
+                return Err("NaiveImplementation/OverConstrained: at least one equation is missing a dependent variable".to_string());
+            }
+        }
+
+        // check N >= M
+        if self.equations.len() < dep_vars.len() {
+            return Err("UnderConstrained: more dependent variables than equations".to_string());
+        }
+
+        println!("Soliving for {:?} in terms of params {:?} and constants {:?}, with {} equations",
+            dep_vars, indep_vars, constants, self.equations.len());
+
+        Ok((*self).clone()) // XXX
+        // TODO: sort output equations by LHS
+    }
+}
+
 fn union_strings(a: &Vec<String>, b: &Vec<String>) -> Vec<String> {
     let mut u = a.clone();
     for e in b {