extern crate modelica_parser;

use std::clone::Clone;
use std::collections::HashSet;
use std::iter::FromIterator;

use self::modelica_parser::*;
use errors::Result;


/// Helpers

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

impl ModelicaModelExt for ModelicaModel {
    // 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> {

        let indep_vars: HashSet<String> = HashSet::from_iter(indep_vars);
        let dep_vars: HashSet<String> = HashSet::from_iter(dep_vars);
        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 bail!("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 bail!("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 => bail!("UnderSpecifiedConstant: {}", name),
                Some(Expr::Integer(_)) |
                Some(Expr::Float(_)) => (), // Ok,
                Some(_) => {
                    bail!("NaiveImplementation: can't handle constant: {}", name);
                }
            }
        }

        // check that there is a depdendent variable in each equation
        for eqn in &self.equations {
            if dep_vars.is_disjoint(&eqn.identifiers()) {
                bail!("NaiveImplementation/OverConstrained: at least one equation is \
                        missing a dependent variable")
            }
        }

        // check N >= M
        if self.equations.len() < dep_vars.len() {
            bail!("UnderConstrained: more dependent variables than equations");
        }

        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| unsolved_vars.intersection(&x.identifiers()).count() == 1);
            let eqn = match next_i {
                None => {
                    return bail!("NaiveImplementation (or poor equation selection?)");
                }
                Some(i) => unsolved_eqns.remove(i),
            };
            let eqn_tmp = eqn.identifiers();
            let ref var = unsolved_vars.intersection(&eqn_tmp).nth(0).unwrap().clone();
            let eqn = eqn.rebalance_for(var.to_string()).expect("rebalance for success");
            // Replace all other references to var with RHS of solved equation
            unsolved_eqns = unsolved_eqns.iter().map(|ref e|
                SimpleEquation{
                    lhs: substitute_with(&e.lhs, &Expr::Ident(var.to_string()), &eqn.rhs),
                    rhs: substitute_with(&e.rhs, &Expr::Ident(var.to_string()), &eqn.rhs),
                }).collect();
            solved.push(eqn);
            unsolved_vars.remove(var);
        }

        Ok(ModelicaModel {
            name: self.name.clone(),
            description: self.description.clone(),
            components: self.components.clone(),
            connections: vec![],
            equations: solved,
        })
    }
}

// Recurses through 'original', replacing all instances of 'a' with 'b'
pub fn substitute_with(original: &Expr, a: &Expr, b: &Expr) -> Expr {
    use modelica_parser::Expr::*;
    println!("original: {:?}  replacing: {:?}   with: {:?}", original, a, b);
    if *original == *a {
        return b.clone();
    }
    match *original {
        Integer(_) | Float(_) | Boolean(_) | StringLiteral(_) | Ident(_) => original.clone(),
        Der(ref e) => Der(Box::new(substitute_with(e, a, b))),
        Sign(ref e) => Sign(Box::new(substitute_with(e, a, b))),
        MathUnaryExpr(muf, ref e) =>
            MathUnaryExpr(muf, Box::new(substitute_with(e, a, b))),
        BinExpr(bo, ref e1, ref e2) =>
            BinExpr(bo,
                    Box::new(substitute_with(e1, a, b)),
                    Box::new(substitute_with(e2, a, b))),
        Array(ref l) => Array(l.iter().map(|ref e| substitute_with(e, a, b)).collect()),
    }
}

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

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

        let ret = match (lvars.contains(&ident), rvars.contains(&ident)) {
            (true, true) => bail!("SymbolicError: NaiveImplementation"),
            (false, false) => bail!("SymbolicError: VariableNotFound"),
            (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) {
                    bail!("SymbolicError: NaiveImplementation")
                } else {
                    Ok(eqn)
                }
            }
            Err(_) => ret,
        }
    }

    fn simplify_lhs(&self, ident: &str) -> Result<SimpleEquation> {
        use modelica_parser::Expr::*;
        use modelica_parser::BinOperator::*;
        match self.lhs {
            Ident(ref s) if s == ident => Ok((*self).clone()),
            Ident(_) | Integer(_) | Float(_) | Boolean(_) | StringLiteral(_) => {
                bail!("SymbolicError: InternalError: expected var on LHS")
            }
            Der(_) |
            MathUnaryExpr(_, _) |
            Sign(_) |
            Array(_) => bail!("SymbolicError: NaiveImplementation: can't simplify"),
            // 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(_, _, _) => {
                bail!("SymbolicError: NotImplemented BinOperator (or else couldn't find var...)")
            }
            // in case we add opers: _ => Err("NotImplemented".to_string()),
        }
    }
}