use crate::data::{Party, Person, Value};
use std::collections::HashMap;
use struct_iterable::Iterable;
use rand::{ thread_rng, Rng };

pub fn party_probabilities(dataset: &Vec<Person>)
-> HashMap<Party, f64> {
    let count_republicans = dataset
        .into_iter()
        .fold(0.0, |acc, person| if person.party == Party::republican { acc + 1.0 } else { acc });
    let p_republicans = count_republicans / dataset.len() as f64;

    let mut ret = HashMap::new();
    ret.insert(Party::republican, p_republicans);
    ret.insert(Party::democrat, 1.0 - p_republicans);
    ret
}

pub type PMaps = (HashMap<Value, HashMap<String, f64>>, HashMap<Party, HashMap<Value, HashMap<String, f64>>>);

#[macro_export]
macro_rules! P {
    // P(Party=democrat|person)
    // P(Party=democrat|"crime"=y,...)
    ( Party = $party:tt | $person:ident , $prob:ident ) => {
        {
            use crate::data::Value;
            use struct_iterable::Iterable;
            use rand::{ thread_rng, Rng };

            let mut prob = if $party == crate::data::Party::democrat { 0.61 } else { 0.39 };
            for (name, val) in $person.iter() {
                if name == "party" { continue; }
                let mut val = val.downcast_ref::<Value>().unwrap().clone();
                if val == Value::q && ! $prob.0.contains_key(&Value::q) {
                    val = if thread_rng().gen_bool(0.5) { Value::y } else { Value::n };
                }
                prob *= P!(name = val | $party, $prob) / P!(name = val, $prob);
            }
            prob
        }
    };
    // P("crime"=y|democrat)
    ( $attr:tt = $value:tt | $party:tt , $prob:ident ) => {
        $prob.1[&$party][&$value][$attr]
    };
    // P("crime"=y)
    ( $attr:tt = $value:tt , $prob:ident ) => {
        $prob.0[&$value][$attr]
    };
}

/// Result: (P(Value[attr]), P(Value[attr]|Party))
/// Use   : .0[y]["crime"]
///         .1[democrat][y]["crime"]
pub fn value_3_probabilities(dataset: &Vec<Person>) -> PMaps {
    let mut len = dataset.len() as f64;

    // P(Value[attr])
    let mut ret_y = HashMap::new();
    let mut ret_n = HashMap::new();
    let mut ret_q = HashMap::new();

    for (name, _) in dataset[0].iter() {
        if name == "party" { continue; }

        let (count_y, count_n, count_q) = dataset
            .into_iter()
            .fold((0.0, 0.0, 0.0), |mut acc, person| {
                match person[name] {
                    Value::y => acc.0 += 1.0,
                    Value::n => acc.1 += 1.0,
                    Value::q => acc.2 += 1.0,
                }
                acc
            });
        ret_y.insert(name.to_string(), count_y / len);
        ret_n.insert(name.to_string(), count_n / len);
        ret_q.insert(name.to_string(), count_q / len);
        // assert_eq!(count_y + count_n + count_q, len);
    }

    let mut simple_p = HashMap::new();
    simple_p.insert(Value::y, ret_y);
    simple_p.insert(Value::n, ret_n);
    simple_p.insert(Value::q, ret_q);

    // assert_eq!(
    //     simple_p[&Value::y]["crime"] +
    //     simple_p[&Value::n]["crime"] +
    //     simple_p[&Value::q]["crime"],
    //     1.0,
    // );

    // P(Value[attr]|Party)
    let mut complex_p = HashMap::new();

    // P(Value[attr]|democrat)
    let mut ret_y = HashMap::new();
    let mut ret_n = HashMap::new();
    let mut ret_q = HashMap::new();
    len = 214.0;

    for (name, _) in dataset[0].iter() {
        if name == "party" { continue; }

        let (count_y, count_n, count_q) = dataset
            .into_iter()
            .fold((0.0, 0.0, 0.0), |mut acc, person| {
                if person.party == Party::democrat {
                    match person[name] {
                        Value::y => acc.0 += 1.0,
                        Value::n => acc.1 += 1.0,
                        Value::q => acc.2 += 1.0,
                    }
                }
                acc
            });
        ret_y.insert(name.to_string(), count_y / len);
        ret_n.insert(name.to_string(), count_n / len);
        ret_q.insert(name.to_string(), count_q / len);
        // assert_eq!(count_y + count_n + count_q, 214.0);
    }

    let mut party_p = HashMap::new();
    party_p.insert(Value::y, ret_y);
    party_p.insert(Value::n, ret_n);
    party_p.insert(Value::q, ret_q);

    complex_p.insert(Party::democrat, party_p);

    // P(Value[attr]|republican)
    let mut ret_y = HashMap::new();
    let mut ret_n = HashMap::new();
    let mut ret_q = HashMap::new();
    len = 134.0;

    for (name, _) in dataset[0].iter() {
        if name == "party" { continue; }

        let (count_y, count_n, count_q) = dataset
            .into_iter()
            .fold((0.0, 0.0, 0.0), |mut acc, person| {
                if person.party == Party::republican {
                    match person[name] {
                        Value::y => acc.0 += 1.0,
                        Value::n => acc.1 += 1.0,
                        Value::q => acc.2 += 1.0,
                    }
                }
                acc
            });
        ret_y.insert(name.to_string(), count_y / len);
        ret_n.insert(name.to_string(), count_n / len);
        ret_q.insert(name.to_string(), count_q / len);
        // assert_eq!(count_y + count_n + count_q, 134.0);
    }

    let mut party_p = HashMap::new();
    party_p.insert(Value::y, ret_y);
    party_p.insert(Value::n, ret_n);
    party_p.insert(Value::q, ret_q);

    complex_p.insert(Party::republican, party_p);

    // assert_eq!(
    //     complex_p[&Party::republican][&Value::y]["crime"] +
    //     complex_p[&Party::republican][&Value::n]["crime"] +
    //     complex_p[&Party::republican][&Value::q]["crime"],
    //     1.0,
    // );

    (simple_p, complex_p)
}

pub fn value_2_probabilities(dataset: &Vec<Person>) -> PMaps {
    let mut len = dataset.len() as f64;

    // P(Value[attr])
    let mut ret_y = HashMap::new();
    let mut ret_n = HashMap::new();

    for (name, _) in dataset[0].iter() {
        if name == "party" { continue; }

        let (count_y, count_n) = dataset
            .into_iter()
            .fold((0.0, 0.0), |mut acc, person| {
                match person[name] {
                    Value::y => acc.0 += 1.0,
                    Value::n => acc.1 += 1.0,
                    Value::q => if thread_rng().gen_bool(0.5) { acc.0 += 1.0 } else { acc.1 += 1.0 },
                }
                acc
            });
        ret_y.insert(name.to_string(), count_y / len);
        ret_n.insert(name.to_string(), count_n / len);
    }

    let mut simple_p = HashMap::new();
    simple_p.insert(Value::y, ret_y);
    simple_p.insert(Value::n, ret_n);

    // P(Value[attr]|Party)
    let mut complex_p = HashMap::new();

    // P(Value[attr]|democrat)
    let mut ret_y = HashMap::new();
    let mut ret_n = HashMap::new();
    len = 214.0;

    for (name, _) in dataset[0].iter() {
        if name == "party" { continue; }

        let (count_y, count_n) = dataset
            .into_iter()
            .fold((0.0, 0.0), |mut acc, person| {
                if person.party == Party::democrat {
                    match person[name] {
                        Value::y => acc.0 += 1.0,
                        Value::n => acc.1 += 1.0,
                        Value::q => if thread_rng().gen_bool(0.5) { acc.0 += 1.0 } else { acc.1 += 1.0 },
                    }
                }
                acc
            });
        ret_y.insert(name.to_string(), count_y / len);
        ret_n.insert(name.to_string(), count_n / len);
    }

    let mut party_p = HashMap::new();
    party_p.insert(Value::y, ret_y);
    party_p.insert(Value::n, ret_n);

    complex_p.insert(Party::democrat, party_p);

    // P(Value[attr]|republican)
    let mut ret_y = HashMap::new();
    let mut ret_n = HashMap::new();
    len = 134.0;

    for (name, _) in dataset[0].iter() {
        if name == "party" { continue; }

        let (count_y, count_n) = dataset
            .into_iter()
            .fold((0.0, 0.0), |mut acc, person| {
                if person.party == Party::republican {
                    match person[name] {
                        Value::y => acc.0 += 1.0,
                        Value::n => acc.1 += 1.0,
                        Value::q => if thread_rng().gen_bool(0.5) { acc.0 += 1.0 } else { acc.1 += 1.0 },
                    }
                }
                acc
            });
        ret_y.insert(name.to_string(), count_y / len);
        ret_n.insert(name.to_string(), count_n / len);
    }

    let mut party_p = HashMap::new();
    party_p.insert(Value::y, ret_y);
    party_p.insert(Value::n, ret_n);

    complex_p.insert(Party::republican, party_p);

    (simple_p, complex_p)
}