use image::{imageops::ColorMap, Rgb, RgbImage};
use std::{array, iter};

const MAX_LEVEL: u8 = 8;

#[inline(always)]
fn get_color_index(color: &Rgb<u8>, level: u8) -> usize {
    let shift = MAX_LEVEL - level;
    let r = (color[0] >> shift) & 1;
    let g = (color[1] >> shift) & 1;
    let b = (color[2] >> shift) & 1;
    ((r << 2) | (g << 1) | b) as usize
}

struct Octree {
    levels: [Vec<OctreeNode>; MAX_LEVEL as usize],
}

impl Octree {
    fn new() -> Self {
        let mut octree = Self {
            levels: Default::default(),
        };
        octree.levels[0].push(OctreeNode::default());
        octree
    }

    fn insert(&mut self, color: &Rgb<u8>) {
        let mut current_node_id = OctreeNodeId { level: 1, index: 0 };
        while current_node_id.level < MAX_LEVEL {
            let child_index = get_color_index(color, current_node_id.level);
            if let Some(child_id) = self.get_node(current_node_id).children[child_index] {
                current_node_id = child_id;
            } else {
                current_node_id = self.insert_child(current_node_id, child_index);
            }
        }
        assert_eq!(current_node_id.level, MAX_LEVEL);
        self.get_node_mut(current_node_id).add_color(color);
    }

    fn get_index(&self, color: &Rgb<u8>) -> usize {
        let mut current_node_id = OctreeNodeId { level: 1, index: 0 };
        while !self.get_node(current_node_id).is_leaf() {
            let child_index = get_color_index(color, current_node_id.level);
            current_node_id = self.get_node(current_node_id).children[child_index]
                .expect("Proper tree can't have dangling nodes");
        }
        self.get_node(current_node_id).index
    }

    fn remove_leaves(&mut self, node_id: OctreeNodeId) -> usize {
        let mut cnt = 0;
        for child_index in 0..8 {
            if let Some(child) = self.delete_child(node_id, child_index) {
                self.get_node_mut(node_id).add_node(&child);
                cnt += 1;
            }
        }
        cnt
    }

    fn reduce_to(&mut self, color_count: usize) {
        let mut color_count_current = self.get_level(MAX_LEVEL).len();
        if color_count_current > color_count {
            for level in (1..MAX_LEVEL).rev() {
                for index in 0..self.get_level(level).len() {
                    let node_id = OctreeNodeId { level, index };
                    color_count_current -= self.remove_leaves(node_id).saturating_sub(1);
                    if color_count_current <= color_count {
                        break;
                    }
                }
            }
        }
    }

    fn finalize(&mut self) -> Vec<Rgb<u8>> {
        let mut palette = Vec::new();
        for level in 1..=MAX_LEVEL {
            for index in 0..self.get_level(level).len() {
                let node_id = OctreeNodeId { level, index };
                self.get_node_mut(node_id).index = palette.len();
                palette.push(Rgb::from(self.get_node(node_id)));
            }
        }
        palette
    }

    fn insert_child(&mut self, node_id: OctreeNodeId, child_index: usize) -> OctreeNodeId {
        assert!(child_index < 8);
        assert!(node_id.level < MAX_LEVEL);
        let level = node_id.level + 1;
        let child_id = OctreeNodeId {
            level,
            index: self.get_level(level).len(),
        };
        self.get_level_mut(level).push(OctreeNode::default());
        assert!(self.get_node(node_id).children[child_index].is_none());
        self.get_node_mut(node_id).children[child_index] = Some(child_id);
        child_id
    }

    fn delete_child(&mut self, node_id: OctreeNodeId, child_index: usize) -> Option<OctreeNode> {
        assert!(child_index < 8);
        let child_id = self.get_node_mut(node_id).children[child_index].take()?;
        assert!(!self.get_node(child_id).is_deleted);
        self.get_node_mut(child_id).is_deleted = true;
        Some(self.get_node(child_id).clone())
    }

    #[inline(always)]
    fn get_level(&self, level: u8) -> &Vec<OctreeNode> {
        &self.levels[(level - 1) as usize]
    }

    #[inline(always)]
    fn get_level_mut(&mut self, level: u8) -> &mut Vec<OctreeNode> {
        &mut self.levels[(level - 1) as usize]
    }

    #[inline(always)]
    fn get_node(&self, node_id: OctreeNodeId) -> &OctreeNode {
        &self.get_level(node_id.level)[node_id.index]
    }

    #[inline(always)]
    fn get_node_mut(&mut self, node_id: OctreeNodeId) -> &mut OctreeNode {
        &mut self.get_level_mut(node_id.level)[node_id.index]
    }
}

#[derive(Debug, Default, Clone, Copy)]
struct OctreeNodeId {
    level: u8,
    index: usize,
}

#[derive(Debug, Default, Clone)]
struct OctreeNode {
    color: [u64; 3],
    count: u64,
    children: [Option<OctreeNodeId>; 8],
    is_deleted: bool,
    index: usize,
}

impl From<&OctreeNode> for Rgb<u8> {
    #[inline(always)]
    fn from(node: &OctreeNode) -> Self {
        if node.count == 0 {
            Rgb::from([0, 0, 0])
        } else {
            Rgb::from(array::from_fn(|i| (node.color[i] / node.count) as u8))
        }
    }
}

impl OctreeNode {
    #[inline(always)]
    fn is_leaf(&self) -> bool {
        !self.is_deleted && self.children.iter().all(Option::is_none)
    }

    #[inline(always)]
    fn add_color(&mut self, color: &Rgb<u8>) {
        self.count += 1;
        iter::zip(&mut self.color, color.0).for_each(|(a, b)| *a += b as u64)
    }

    #[inline(always)]
    fn add_node(&mut self, node: &OctreeNode) {
        self.count += node.count;
        iter::zip(&mut self.color, node.color).for_each(|(a, b)| *a += b)
    }
}

pub struct ColorQuantizer {
    octree: Octree,
    colors: Vec<Rgb<u8>>,
}

impl ColorQuantizer {
    pub fn from(img: &RgbImage, palette_size: usize) -> Self {
        let palette_size = palette_size.min(256);
        let mut octree = Octree::new();
        for pixel in img.pixels() {
            octree.insert(pixel);
        }
        octree.reduce_to(palette_size);
        let colors = octree.finalize();
        assert!(
            colors.len() <= palette_size,
            "Color palette size exceeded {palette_size}"
        );
        println!("final color palette size: {}", colors.len());
        Self { octree, colors }
    }

    #[inline(always)]
    pub fn get_palette(&self) -> &[Rgb<u8>] {
        &self.colors
    }

    #[inline(always)]
    pub fn get_index(&self, color: &Rgb<u8>) -> usize {
        self.octree.get_index(color)
    }
}

impl ColorMap for ColorQuantizer {
    type Color = Rgb<u8>;

    #[inline(always)]
    fn index_of(&self, color: &Self::Color) -> usize {
        self.get_index(color)
    }

    #[inline(always)]
    fn map_color(&self, color: &mut Self::Color) {
        *color = self.get_palette()[self.get_index(color)]
    }
}