use crate::{LabelFormat,LabelU16};

use avl_tree::*;

#[derive(Debug,PartialEq)]
pub struct Data {
    pub label: u16,
    pub area: usize,
    pub perimeter: usize,
    pub neighboring: usize,
    pub non_neighboring: usize,
}

impl Data {
    const fn new(label: u16, area: usize, perimeter: usize, neighboring: usize, non_neighboring: usize) -> Self {
        Self {
            label,
            area,
            perimeter,
            neighboring,
            non_neighboring,
        }
    }
}

fn any(a: &[bool]) -> bool {
    for value in a.iter() {
        if *value {
            return true;
        }
    }
    false
}

impl LabelFormat<u16> {
    pub fn get_tree(&self) -> AVLTree<u16, Data> {
        let mut avt = avl_tree::AVLTree::<u16, Data>::new();
        for y in 0..self.height {
            for x in 0..self.width {
                let index = x + y * self.width;
                let key = self.buffer[index];
                if key.is_zero() {
                    continue;
                }
                let mut new_data = Data::new(
                    key,
                    0,
                    0,
                    0,
                    0);
                match avt.get(key) {
                    Some(data) => {
                        new_data.area = data.area;
                        new_data.perimeter = data.perimeter;
                        new_data.neighboring = data.neighboring;
                        new_data.non_neighboring = data.non_neighboring;
                    },
                    None => {
                    },
                }
                let neighbors: [Option<u16>;4] = [
                    self.get_up(x,y),
                    self.get_right(x,y),
                    self.get_down(x,y),
                    self.get_left(x,y),
                ];
                let perimeter_slice = neighbors.map(|x|
                    match x {
                        Some(y) => (y.is_zero()) || (y != new_data.label),
                        None => false,
                    });
                let add_perimeter = any(&perimeter_slice);
                let neighboring_slice = neighbors.map(|x|
                    match x {
                        Some(y) => (y != new_data.label) && (!y.is_zero()),
                        None => false,
                    });
                let add_neighboring = any(&neighboring_slice);
                let non_neighboring_slice = neighbors.map(|x|
                    match x {
                        Some(y) => y.is_zero(),
                        None => false,
                    });
                let add_non_neighboring = any(&non_neighboring_slice);
                // Always increase area
                new_data.area += 1;
                // Increase perimeter if there is a neighbor that is not the same cell
                if add_perimeter {
                    new_data.perimeter += 1;
                }
                if add_non_neighboring {
                    new_data.non_neighboring += 1;
                }
                if add_neighboring {
                    new_data.neighboring += 1;
                }
                avt.insert(key, new_data);
            }
        }
        avt
    }

    pub fn filter<T: Fn(&Data) -> bool>(&self, filter: T) -> LabelFormat<u16> {
        let mut output = LabelFormat::<u16> {
            buffer: vec![0u16; self.width*self.height],
            width: self.width,
            height: self.height,
        };
        let tree = self.get_tree();
        for y in 0..self.height {
            for x in 0..self.width {
                let index = x + y * self.width;
                let key = self.buffer[index];
                if key.is_zero() {
                    continue;
                }
                if let Some(data) = tree.get(key) {
                    if filter(data) {
                        output.buffer[index] = self.buffer[index];
                    }
                } else {
                    panic!("Tree should contain all non-zero values!");
                }
            }
        }
        output
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn get_areas_test() {
        const DIM: usize = 6;
        let mut test_data = LabelFormat::<u16> {
            buffer: vec![0u16; DIM*DIM],
            width: DIM,
            height: DIM,
        };
        test_data.buffer[2+3*DIM] = 1;
        test_data.buffer[3+3*DIM] = 1;
        test_data.buffer[4+3*DIM] = 2;
        test_data.idilate();

        let expected: [Vec<u16>;1] = [
            [
                0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0,
                0, 0, 1, 1, 2, 0,
                0, 1, 1, 1, 2, 2,
                0, 0, 1, 1, 2, 0,
                0, 0, 0, 0, 0, 0,
            ].to_vec(),
        ];

        assert_eq!(test_data.buffer, expected[0]);
        let tree = test_data.get_tree();
        let data = tree.get(0);
        assert_eq!(data, None);
        let data = tree.get(1).expect("Key should exist");
        assert_eq!(*data, Data::new(1, 7, 6, 3, 5));
        let data = tree.get(2).expect("Key should exist");
        assert_eq!(*data, Data::new(2, 4, 4, 3, 3));
        let data = tree.get(3);
        assert_eq!(data, None);
    }

    #[test]
    fn filter_test() {
        const DIM: usize = 6;
        let mut test_data = LabelFormat::<u16> {
            buffer: vec![0u16; DIM*DIM],
            width: DIM,
            height: DIM,
        };
        test_data.buffer[2+3*DIM] = 1;
        test_data.buffer[3+3*DIM] = 1;
        test_data.buffer[4+3*DIM] = 2;
        test_data.idilate();

        let expected: [Vec<u16>;2] = [
            [
                0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0,
                0, 0, 1, 1, 2, 0,
                0, 1, 1, 1, 2, 2,
                0, 0, 1, 1, 2, 0,
                0, 0, 0, 0, 0, 0,
            ].to_vec(),
            [
                0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0,
                0, 0, 1, 1, 0, 0,
                0, 1, 1, 1, 0, 0,
                0, 0, 1, 1, 0, 0,
                0, 0, 0, 0, 0, 0,
            ].to_vec(),
        ];

        let tree = test_data.get_tree();
        let data = tree.get(0);
        assert_eq!(data, None);
        let data = tree.get(1).expect("Key should exist");
        assert_eq!(*data, Data::new(1, 7, 6, 3, 5));
        let data = tree.get(2).expect("Key should exist");
        assert_eq!(*data, Data::new(2, 4, 4, 3, 3));
        let data = tree.get(3);
        assert_eq!(data, None);

        let filtered_data = test_data.filter(|x| x.area > 5);
        assert_eq!(filtered_data.buffer, expected[1]);
        let tree = filtered_data.get_tree();
        let data = tree.get(0);
        assert_eq!(data, None);
        let data = tree.get(1).expect("Key should exist");
        assert_eq!(*data, Data::new(1, 7, 6, 0, 6));
        let data = tree.get(2);
        assert_eq!(data, None);
    }
}