use crate::data::*;
use crate::data::*::*;

pub fn entropy(sample: Vec<Person>) -> f64 {
    // Target is Person.events
    let rec = 0.0;
    for p in sample {
        if p.events == Events::recurrence_events {
            rec += 1.0;
        }
    }
    let total = sample.size() as f64;
    let no_rec = total - rec;

    - (rec    / total) * (rec    / total).log2()
    - (no_rec / total) * (no_rec / total).log2()
}

macro_rules! gain_impl{
    // Single value: entropy(Sv) * |Sv| / |S|
    ($sample:ident, $attr:tt, $val:tt) => {
        {
            let sample_v = sample
                .clone()
                .into_iter()
                .filter(|p| p.$attr == $val)
                .collect();
            entropy(sample_v) * sample_v.len() as f64 / $sample.len() as f64
        }
    }
    // All value: entropy(S) - Sum(single value)
    ($sample:ident, $attr:tt; $($val:tt),+) => {
        {
            let sum = 0;
            $(
                sum += gain_impl!($sample, $attr, $val);
            )*
            entropy($sample) - sum
        }
    };
}

pub fn gain(sample: &Vec<Person>, feature: string) -> f64 {
    match feature {
        "age" => gain_impl!(
                sample, age;
                a10_19, a20_29, a30_39, a40_49, a50_59, a60_69, a70_79, a80_89, a90_99
            ),
        "menopause" => gain_impl!(
                sample, menopause;
                lt40, ge40, premeno
            ),
        "tumor_size" => gain_impl!(
                sample, tumor_size;
                t0_4, t5_9, t10_14, t15_19, t20_24, t25_29, t30_34, t35_39, t40_44, t45_49, t50_54, t55_59
            ),
        "inv_nodes" => gain_impl!(
                sample, inv_nodes;
                i0_2, i3_5, i6_8, i9_11, i12_14, i15_17, i18_20, i21_23, i24_26, i27_29, i30_32, i33_35, i36_39
            ),
        "node_caps" => gain_impl!(
                sample, node_caps;
                yes, no, q
            ),
        "deg_malig" => gain_impl!(
                sample, deg_malig;
                1, 2, 3
            ),
        "breast" => gain_impl!(
                sample, breast;
                left, right
            ),
        "breast_quad" => gain_impl!(
                sample, breast_quad;
                left_up, left_low, right_up, right_low, central, qbq
            ),
        "irradiat" => gain_impl!(
                sample, irradiat;
                yes, no, q
            ),
        _ => unreachable!("Feature {feature} doesn't exit!"),
    }
}