extern crate modelica_parser;

use std::clone::Clone;
use std::collections::HashMap;
use std::collections::HashSet;

use self::modelica_parser::ast::*;


//// Helpers

pub trait ModelicaModelExt {
    fn get_constant_vars(&self) -> HashMap<String,Option<Expr>>;
    fn get_free_vars(&self) -> Vec<String>;
    fn solve_for(&self, indep_vars: Vec<String>, dep_vars: Vec<String>) -> Result<ModelicaModel,String>;
}

impl ModelicaModelExt for ModelicaModel {

    fn get_constant_vars(&self) -> HashMap<String,Option<Expr>> {
        let mut binds = HashMap::new();
        // XXX: actually implement this...
        for c in &self.components {
            match c.prefix {
                Some(ComponentPrefix::Constant) => { binds.insert(c.name.clone(), Some(Expr::Integer(123))); },
                Some(ComponentPrefix::Parameter) => { binds.insert(c.name.clone(), Some(Expr::Float(4.56))); },
                _ => (),
            }
        }
        binds
    }

    // This crude function finds "unbound" variables: those which are not constants, parameters, or
    // the sole element on the LHS of an equation.
    // Bugs:
    //   if a var is on LHS and RHS of same equation
    fn get_free_vars(&self) -> Vec<String> {
        // Start with components, and remove constants and parameters
        let vars = self.components.iter().filter(|v| match v.prefix {
                Some(ComponentPrefix::Constant) | Some(ComponentPrefix::Parameter) => false,
                _ => true,
            });

        // Remove LHS (bound) vars
        let mut outputs = vec![];
        for eq in self.equations.iter() {
            // TODO:
            if let Expr::Ident(ref symb) = eq.lhs {
                outputs.push(symb.to_string());
            }
        }
        let vars = vars.filter(|v| !outputs.contains(&v.name));

        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
    fn solve_for(&self, indep_vars: Vec<String>, dep_vars: Vec<String>) -> Result<ModelicaModel,String> {
        let constants = self.get_constant_vars();
        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(format!("Variable counts don't add up (V={} Q={} P={} M={})",
                all_vars.len(),
                constants.len(),
                indep_vars.len(),
                dep_vars.len()));
        }

        // 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());

        let mut unsolved_eqns = self.equations.clone();
        let mut solved: Vec<SimpleEquation> = vec![];
        let mut unsolved_vars = dep_vars.clone();
        while unsolved_eqns.len() > 0 {
            let next_i = unsolved_eqns
                .iter()
                .position(|ref x| intersect_strings(&unsolved_vars, &x.identifiers()).len() == 1);
            let eqn = match next_i {
                None => { return Err("NaiveImplementation (or poor equation selection?)".to_string()); },
                Some(i) => unsolved_eqns.remove(i),
            };
            let ref var = intersect_strings(&unsolved_vars, &eqn.identifiers())[0];
            let eqn = eqn.rebalance_for(var.to_string()).expect("rebalance for success");
            solved.push(eqn);
            let var_i = unsolved_vars.iter().position(|ref x| x == &var).unwrap();
            unsolved_vars.remove(var_i);
        };

        // TODO: sort output equations by LHS
        Ok(ModelicaModel{
            name: self.name.clone(),
            components: self.components.clone(),
            connections: vec![],
            equations: solved,
            extends: vec![],
        })
    }
}

fn intersect_strings(a: &Vec<String>, b: &Vec<String>) -> Vec<String> {
    let mut both = vec![];
    for e in a {
        if b.contains(&e) {
            both.push(e.clone());
        }
    }
    both
}

pub trait SimpleEquationExt {
    fn rebalance_for(&self, ident: String) -> Result<SimpleEquation,String>;
    fn simplify_lhs(&self, ident: &str) -> Result<SimpleEquation,String>;
}

impl SimpleEquationExt for SimpleEquation {

    fn rebalance_for(&self, ident: String) -> Result<SimpleEquation,String> {
        let lvars = self.lhs.identifiers();
        let rvars = self.rhs.identifiers();

        let ret = match (lvars.contains(&ident), rvars.contains(&ident)) {
            (true, true) => Err("SymbolicError: NaiveImplementation".to_string()),
            (false, false) => Err("SymbolicError: VariableNotFound".to_string()),
            (true, false) => self.simplify_lhs(&ident),
            (false, true) =>
                SimpleEquation{lhs: self.rhs.clone(),
                               rhs: self.lhs.clone()}.simplify_lhs(&ident),
        };
        match ret {
            Ok(eqn) => {
                if eqn.rhs.contains(&ident) {
                    Err("SymbolicError: NaiveImplementation".to_string())
                } else {
                    Ok(eqn)
                }},
            Err(_) => ret,
        }
    }

    fn simplify_lhs(&self, ident: &str) -> Result<SimpleEquation,String> {
        use modelica_parser::ast::Expr::*;
        use modelica_parser::ast::BinOperator::*;
        match self.lhs {
            Ident(ref s) if s == ident => Ok((*self).clone()),
            Ident(_) | Integer(_) | Float(_) =>
                Err("SymbolicError: InternalError: expected var on LHS".to_string()),
            Der(_) | Abs(_) =>
                Err("SymbolicError: NaiveImplementation: can't simplify der() or abs()".to_string()),
            // TODO: create a macro for the below...
            BinExpr(Multiply, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Divide, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Multiply, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Divide, Box::new(self.rhs.clone()), a.clone())}.simplify_lhs(&ident) },
            BinExpr(Divide, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Multiply, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Divide, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Divide, a.clone(), Box::new(self.rhs.clone()))}.simplify_lhs(&ident) },
            BinExpr(Add, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Subtract, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Add, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Subtract, Box::new(self.rhs.clone()), a.clone())}.simplify_lhs(&ident) },
            BinExpr(Subtract, ref a, ref b) if a.contains(ident) => {
                SimpleEquation{
                    lhs: *a.clone(),
                    rhs: BinExpr(Add, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
            BinExpr(Subtract, ref a, ref b) if b.contains(ident) => {
                SimpleEquation{
                    lhs: *b.clone(),
                    rhs: BinExpr(Subtract, a.clone(), Box::new(self.rhs.clone()))}.simplify_lhs(&ident) },
            BinExpr(_, _, _) => Err("SymbolicError: NotImplemented BinOperator (or else couldn't find var...)".to_string()),
            // in case we add opers: _ => Err("NotImplemented".to_string()),
        }
    }
}