1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

pub mod linalg;
pub mod mvcalc;
#[cfg(test)]
mod tests;
use linalg::*;
use mvcalc::*;
use meval::{ Context, eval_str_with_context };
use std::collections::HashMap;

/// A result from the Newton-Raphson Solver that can be valid, valid with warnings, or erroneous.
pub enum SolverResult<T> {
    Ok(T),
    Warn(T),
    Err
}
impl <'a>SolverResult<HashMap<&'a str, Variable>> {
    /// Returns the contained value, consuming the `self` value in the process.
    pub fn unwrap(self) -> HashMap<&'a str, Variable> {
        match self {
            SolverResult::Ok(t) => t,
            SolverResult::Warn(t) => t,
            SolverResult::Err => panic!()
        }
    }
}

/// Returns a `HashMap` indicating where a given HashMap's keys will be placed in a Vec.
#[allow(dead_code)]
fn index_map<'a, V>(hm: &HashMap<&'a str, V>) -> HashMap<&'a str, usize> {
    let mut i: usize = 0;
    let mut res = HashMap::new();
    for k in hm.keys() {
        res.insert(*k, i);
        i += 1;
    }
    res
}

/// Returns a tuple of `Vec`s that contain the keys and values of the original HashMap. 
/// The index of the key will be the same as its corresponding value's index.
/// # Example
/// ```
/// use std::collections::HashMap;
/// use nexsys_core::split_hm;
/// 
/// let my_map = HashMap::from([
///     ("a", 1),
///     ("b", 2),
///     ("c", 3)
/// ]);
/// 
/// let vecs = split_hm(my_map.clone());
///  
/// assert_eq!(my_map[vecs.0[0]], vecs.1[0])
/// ```
pub fn split_hm<K, V>(hm: HashMap<K, V>) -> (Vec<K>, Vec<V>) {
    let mut keys = Vec::new();
    let mut vals = Vec::new();

    for i in hm {
        keys.push(i.0);
        vals.push(i.1);
    }

    (keys, vals)
}

/// Reverses the operation performed by `split_hm`
pub fn stitch_hm<K: std::hash::Hash + std::cmp::Eq, V>(mut keys: Vec<K>, mut vals: Vec<V>) -> HashMap<K, V> {
    let mut res = HashMap::new();
    for _ in 0..keys.len() {
        res.insert(
            keys.pop().unwrap(), 
            vals.pop().unwrap()
        );
    }
    res
}

/// Takes a mathematical expression given as a `&str` and returns a function that takes a `HashMap<&str, Variable>` and returns an `f64`.
fn functionify<'a>(text: &'a str) -> impl Fn(&HashMap<&str, Variable>) -> f64 + 'a {
    let func = move |v:&HashMap<&str, Variable>| -> f64 {
        let mut ctx = Context::new();
        
        for k in v {
            ctx.var(*k.0, k.1.as_f64());
        }

        eval_str_with_context(text, ctx)
            .expect(&format!("ERR: Failed to evaluate expression: {}", text))
    };
    func
}

fn newton_iteration<'a>(system: &Vec<&str>, mut guess: HashMap<&'a str, Variable>) -> HashMap<&'a str, Variable> {
    let mut j = jacobian(system, &guess);
    j.invert().expect("ERR: Jacobian matrix is non-invertible!");

    let fx = Vec::from_iter(
        system.iter().map(
            |i| functionify(i)(&guess)
        )
    );
    let x_n = stitch_hm(
        j.vars.clone().unwrap(),
        mat_vec_dot(j, fx)
    );
    for v in &mut guess {
        v.1.change(x_n[&v.0.to_string()])
    }

    guess
}

pub fn mv_newton_raphson<'a>(
    system: Vec<&str>, 
    guess: HashMap<&'a str, Variable>,
    threshold: f64,
    limit: i32
) -> SolverResult<HashMap<&'a str, Variable>> {
    
    let error = |guess: &HashMap<&str, Variable>| -> f64 {
        system.iter().map(
            |i| {
                let mut ctx = Context::new();
                
                for j in guess {
                    ctx.var(*j.0, j.1.as_f64()); 
                }
                
                let exp = i.replace("=", "-");
                let error_msg = format!("Correctness function failed to evaluate the system string: {}", &exp);
                
                eval_str_with_context(&exp, ctx)
                .expect(&error_msg)
                .abs()
            }
        ).sum()
    };
    
    let mut count = 0;
    let mut new_guess = guess;    

    loop {
        new_guess = newton_iteration(&system, new_guess);
        let e = error(&new_guess);

        if e < threshold { // Solution is valid and acceptable
            return SolverResult::Ok(new_guess)
        } else if count > limit { // Solution is valid, but timed out
            new_guess.insert("__error__", Variable::new(e, None));
            return SolverResult::Warn(new_guess)
        } 
        count += 1;
    }
}