use image::{imageops::ColorMap, Rgb, RgbImage};
use std::{
    array, cmp,
    collections::{BinaryHeap, VecDeque},
    iter,
};

use crate::queue::Queue;

const MAX_LEVEL: u8 = 8;
const MAX_NODES: usize = 1024;
const MAX_COLORS: usize = 256;

#[derive(Debug, Default)]
struct Node {
    rgb: [u32; 3],
    count: u32,
    index: u8,
    level: u8,
    is_leaf: bool,
    parent: Option<u32>,
    children: [Option<u32>; 8],
}

impl Node {
    fn merge_color(&mut self, color: Rgb<u8>) {
        self.count += 1;
        iter::zip(&mut self.rgb, color.0).for_each(|(a, b)| *a += b as u32)
    }

    fn merge_node(&mut self, node: Node) {
        self.count += node.count;
        iter::zip(&mut self.rgb, node.rgb).for_each(|(a, b)| *a += b)
    }
}

#[derive(Debug)]
struct Pool {
    nodes: [Node; MAX_NODES],
    ids: Queue<u32, MAX_NODES>,
}

impl Pool {
    fn new() -> Self {
        let mut ids = Queue::new();
        let nodes = array::from_fn(|i| {
            ids.push(i as u32);
            Node::default()
        });
        Self { nodes, ids }
    }

    fn create(&mut self) -> u32 {
        self.ids.pop().unwrap()
    }

    fn get(&self, id: u32) -> &Node {
        &self.nodes[id as usize]
    }

    fn get_mut(&mut self, id: u32) -> &mut Node {
        &mut self.nodes[id as usize]
    }

    fn delete(&mut self, id: u32) -> Node {
        self.ids.push(id);
        std::mem::take(&mut self.nodes[id as usize])
    }
}

fn get_color_index(color: Rgb<u8>, level: u8) -> usize {
    let shift = MAX_LEVEL - level;
    color
        .0
        .into_iter()
        .rev()
        .enumerate()
        .map(|(i, c)| (((c >> shift) & 1) << i) as usize)
        .fold(0, |s, c| s | c)
}

#[derive(Debug, Ord)]
struct Reducible {
    node_id: u32,
    level: u8,
    count: u32,
}

impl Reducible {
    fn new(node_id: u32, node: &Node) -> Self {
        Self {
            node_id,
            level: node.level,
            count: node.count,
        }
    }
}

impl Eq for Reducible {}

impl cmp::PartialEq for Reducible {
    fn eq(&self, other: &Self) -> bool {
        self.level.eq(&other.level) && self.count.eq(&other.count)
    }
}

impl cmp::PartialOrd for Reducible {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(match self.level.cmp(&other.level) {
            cmp::Ordering::Greater => cmp::Ordering::Greater,
            cmp::Ordering::Less => cmp::Ordering::Less,
            cmp::Ordering::Equal => self.count.cmp(&other.count),
        })
    }
}

struct Octree {
    pool: Pool,
    root: u32,
    reducible: BinaryHeap<Reducible>,
    color_count: usize,
    leaf_count: usize,
}

impl Octree {
    fn new(color_count: usize) -> Self {
        let mut pool = Pool::new();
        let root = pool.create();
        pool.get_mut(root).is_leaf = true;
        Self {
            pool,
            root,
            reducible: BinaryHeap::new(),
            color_count,
            leaf_count: 1,
        }
    }

    pub fn traverse<F>(&self, f: F)
    where
        F: FnMut(u32, &Node),
    {
        let mut f = f;
        let mut queue = VecDeque::new();
        queue.push_back(self.root);
        while let Some(node_id) = queue.pop_front() {
            let node = self.pool.get(node_id);
            f(node_id, node);
            for child_id in node.children.iter().flatten() {
                queue.push_back(*child_id);
            }
        }
    }

    pub fn traverse_mut<F>(&mut self, f: F)
    where
        F: FnMut(u32, &mut Node),
    {
        let mut f = f;
        let mut queue = VecDeque::new();
        queue.push_back(self.root);
        while let Some(node_id) = queue.pop_front() {
            let node = self.pool.get_mut(node_id);
            f(node_id, node);
            for child_id in node.children.iter().flatten() {
                queue.push_back(*child_id);
            }
        }
    }

    fn insert(&mut self, color: Rgb<u8>) {
        let mut node_id = self.root;
        for level in 1..=MAX_LEVEL {
            let child_index = get_color_index(color, level);
            node_id = match self.pool.get(node_id).children[child_index] {
                Some(child_id) => child_id,
                None => {
                    let child_id = self.pool.create();
                    {
                        let child = self.pool.get_mut(child_id);
                        child.level = level;
                        child.parent = Some(node_id);
                        child.is_leaf = true;
                        self.leaf_count += 1;
                    }
                    {
                        let parent = self.pool.get_mut(node_id);
                        parent.children[child_index] = Some(child_id);
                        if parent.is_leaf {
                            parent.is_leaf = false;
                            self.leaf_count -= 1;
                            self.reducible.push(Reducible::new(node_id, &parent));
                        }
                    }
                    child_id
                }
            }
        }
        self.pool.get_mut(node_id).merge_color(color);
        self.reduce();
    }

    fn get_index(&self, color: Rgb<u8>) -> usize {
        let mut node_id = self.root;
        for level in 1..=MAX_LEVEL {
            let node = self.pool.get(node_id);
            if node.is_leaf {
                break;
            }
            let child_index = get_color_index(color, level);
            node_id = match node.children[child_index] {
                Some(child_id) => child_id,
                None => {
                    match node
                        .children
                        .iter()
                        .enumerate()
                        .filter_map(|(i, c)| c.zip(Some(i)))
                        .map(|(c, i)| {
                            let d = child_index ^ i;
                            ([d >> 2, (d >> 1) & 1, d & 1].into_iter().sum::<usize>(), c)
                        })
                        .min_by(|a, b| a.0.cmp(&b.0))
                        .map(|(_, c)| c)
                    {
                        Some(child_id) => child_id,
                        None => break,
                    }
                }
            }
        }
        self.pool.get(node_id).index as usize
    }

    fn prune_node(&mut self, node_id: u32) {
        self.leaf_count -= {
            let node = self.pool.get_mut(node_id);
            node.is_leaf = true;
            self.leaf_count += 1;
            std::mem::take(&mut node.children)
        }
        .into_iter()
        .flatten()
        .map(|child_id| {
            let child = self.pool.delete(child_id);
            self.pool.get_mut(node_id).merge_node(child);
        })
        .count()
    }

    fn reduce(&mut self) {
        if self.leaf_count > self.color_count {
            while let Some(reducible) = self.reducible.pop() {
                self.prune_node(reducible.node_id);
                if let Some(parent_id) = self.pool.get(reducible.node_id).parent {
                    self.reducible
                        .push(Reducible::new(parent_id, self.pool.get(parent_id)));
                }
                if self.leaf_count <= self.color_count {
                    break;
                }
            }
        }
    }

    fn finalize(&mut self) -> Vec<Rgb<u8>> {
        let mut palette = Vec::new();
        self.traverse_mut(|_, node| {
            if node.is_leaf {
                node.index = palette.len() as u8;
                palette.push(Rgb::from(array::from_fn(|i| {
                    (node.rgb[i] / node.count) as u8
                })));
            }
        });
        palette
    }
}

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(MAX_COLORS);
        let mut octree = Octree::new(palette_size);
        for pixel in img.pixels() {
            octree.insert(*pixel);
        }
        let colors = octree.finalize();
        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)]
    }
}