use std::collections::HashMap;
use std::collections::VecDeque;
use std::hash::Hash;
use std::usize;
use std::vec;

use itertools::Itertools;
use itertools::kmerge;

use crate::voxel;

#[derive(Clone, Copy, PartialEq, Debug)]
pub struct Color(pub f32, pub f32, pub f32, pub f32);

pub struct NTree<const N: usize>
{
    structure: HashMap<NTreeLocator<N>, NTreeNode<N>>,
    depth: u32,
}

#[derive(Clone, Copy)]
struct NTreeNode<const N: usize>
{
    color: Color,
    subdivided: bool,
}

impl<const N: usize> NTree<N>
where
    [(); N * N * N]:,
{
    pub fn constant(color: Color, depth: u32) -> Self
    {
        let mut structure = HashMap::new();
        structure.insert(
            NTreeLocator::root(),
            NTreeNode {
                color,
                subdivided: false,
            },
        );
        Self { structure, depth }
    }

    pub fn get(&self, x: usize, y: usize, z: usize) -> Color
    {
        let mut local_x = x;
        let mut local_y = y;
        let mut local_z = z;

        let mut size = N.pow(self.depth);
        assert!(x < size && y < size && z < size);

        let mut current_loc = NTreeLocator::<N>::root();
        let mut current_node = *self.structure.get(&current_loc).unwrap();

        loop
        {
            if !current_node.subdivided
            {
                return current_node.color;
            }

            size /= N;

            // Descend
            let child_x = local_x / size;
            let child_y = local_y / size;
            let child_z = local_z / size;

            // Node is subdivided, descend
            current_loc = current_loc.get_child(child_x, child_y, child_z);
            current_node = *self.structure.get(&current_loc).unwrap();

            local_x -= child_x * size;
            local_y -= child_y * size;
            local_z -= child_z * size;
        }
    }

    pub fn from_arrays_old(chunk: &[Color], depth: u32) -> Self
    {
        let chunk_width = N.pow(depth);
        let mut structure = HashMap::with_capacity(chunk_width * chunk_width * chunk_width);
        println!("Inserting voxels");
        for ((x, y), z) in (0..chunk_width)
            .cartesian_product(0..chunk_width)
            .cartesian_product(0..chunk_width)
        {
            structure.insert(
                NTreeLocator::<N>::from_depth_coords(x, y, z, depth as usize),
                NTreeNode::<N> {
                    color: chunk[x + y * chunk_width + z * chunk_width * chunk_width],
                    subdivided: false,
                },
            );
        }

        println!("Starting bottom up");
        let mut current_size = 1;
        for d in (0..depth).rev()
        {
            println!("depth: {d}");
            current_size *= N;
            for ((x, y), z) in (0..(chunk_width / current_size))
                .cartesian_product(0..(chunk_width / current_size))
                .cartesian_product(0..(chunk_width / current_size))
            {
                let loc = NTreeLocator::<N>::from_depth_coords(x, y, z, d as usize);
                let children = (0..N)
                    .cartesian_product(0..N)
                    .cartesian_product(0..N)
                    .map(|((cx, cy), cz)| structure.get(&loc.get_child(cx, cy, cz)).unwrap())
                    .collect::<Vec<_>>();

                let (can_merge, _) = children.iter().fold((true, None), |acc, child| match acc
                {
                    (_, None) => (true, Some(child.color)),
                    (b, Some(color)) =>
                    {
                        (b && !child.subdivided && color == child.color, Some(color))
                    }
                });

                let node;
                if can_merge
                {
                    node = NTreeNode::<N> {
                        color: children[0].color,
                        subdivided: false,
                    };

                    //Remove merged childrenn
                    (0..N)
                        .cartesian_product(0..N)
                        .cartesian_product(0..N)
                        .for_each(|((cx, cy), cz)| {
                            structure.remove(&loc.get_child(cx, cy, cz)).unwrap();
                        });
                }
                else
                {
                    node = NTreeNode::<N> {
                        color: Color::average(
                            children
                                .iter()
                                .map(|x| x.color)
                                .collect::<Vec<_>>()
                                .as_slice(),
                        ),
                        subdivided: true,
                    };
                }
                structure.insert(loc, node);
            }
        }

        Self { structure, depth }
    }

    pub fn from_arrays(chunk: &[Color], depth: usize) -> Self
    {
        let width = N.pow(depth as u32);
        let mut structure = HashMap::new();
        let mut taken = vec![0; chunk.len()]; // Whether or not the current voxel has been added

        // Voxel insertion/combination pass
        for (y, z) in (0..width).cartesian_product(0..width)
        {
            let mut x = 0;
            while x < width
            {
                if taken[x + y * width + z * width * width] != 0
                {
                    x += taken[x + y * width + z * width * width];
                    continue;
                }

                let max_combination_depth = trailing_zeroes::<N>(x)
                    .unwrap_or(depth)
                    .min(trailing_zeroes::<N>(y).unwrap_or(depth))
                    .min(trailing_zeroes::<N>(z).unwrap_or(depth));

                let prev_color = chunk[x + y * width + z * width * width];
                let mut locator = NTreeLocator::<N>::from_depth_coords(x, y, z, depth);
                //let mut insertion_depth = depth;
                let mut block_width = 1;
                'depth_loop: for depth in 1..=max_combination_depth
                {
                    let combination_width = N.pow(depth as u32);
                    for ((sx, sy), sz) in (0..combination_width)
                        .cartesian_product(0..combination_width)
                        .cartesian_product(0..combination_width)
                    {
                        let voxel_color =
                            chunk[(x + sx) + (y + sy) * width + (z + sz) * width * width];
                        let voxel_taken =
                            taken[(x + sx) + (y + sy) * width + (z + sz) * width * width];
                        if prev_color != voxel_color || voxel_taken != 0
                        {
                            // Cannot merge further
                            break 'depth_loop;
                        }
                    }

                    // At this point, voxel in combination_width^3 block can be merged
                    block_width = combination_width;
                    //insertion_depth -= 1;
                    locator = locator.get_parent();
                }

                structure.insert(
                    locator,
                    NTreeNode {
                        color: prev_color,
                        subdivided: false,
                    },
                );

                // Mark as taken
                for ((sx, sy), sz) in (0..block_width)
                    .cartesian_product(0..block_width)
                    .cartesian_product(0..block_width)
                {
                    taken[(x + sx) + (y + sy) * width + (z + sz) * width * width] = block_width;
                }

                x += block_width;
            }
        }

        // Increase pass
        for d in (0..=(depth - 1)).rev()
        {
            // Iterate on blocks of depth d
            let block_count = N.pow(d as u32); // Number of such blocks along axis
            let block_width = width / block_count; // Number of such blocks along axis
            for ((bx, by), bz) in (0..(block_count))
                .cartesian_product(0..(block_count))
                .cartesian_product(0..(block_count))
            {
                // Get how was the origin voxel merged
                let merged_width = taken[(bx * block_width)
                    + (by * block_width) * width
                    + (bz * block_width) * width * width];

                if merged_width >= block_width
                {
                    // Current voxels have been merged in bigger or equal block
                    continue;
                }

                let locator = NTreeLocator::<N>::from_depth_coords(bx, by, bz, d);

                // Otherwise, merge

                // Children CANNOT be merged, Otherwise they would have been already
                // Compute average color

                let colors = locator
                    .iter_children()
                    .map(|loc| structure.get(&loc).unwrap().color)
                    .collect::<Vec<_>>();
                structure.insert(
                    locator,
                    NTreeNode::<N> {
                        color: Color::average(&colors),
                        subdivided: true,
                    },
                );
            }
        }

        Self {
            structure,
            depth: depth as u32,
        }
    }

    pub fn set(&mut self, x: usize, y: usize, z: usize, color: Color)
    {
        let mut local_x = x;
        let mut local_y = y;
        let mut local_z = z;

        let mut size = N.pow(self.depth);
        assert!(x < size && y < size && z < size);

        let mut current_loc = NTreeLocator::<N>::root();
        let mut current_node = *self.structure.get(&current_loc).unwrap();
        let mut current_depth = 0;

        loop
        {
            // Subnodes alread with correct color
            if current_node.color == color
            {
                return;
            }

            if current_depth == self.depth
            {
                current_node.color = color;
                self.structure.get_mut(&current_loc).unwrap().color = color;
                break;
            }

            if !current_node.subdivided
            {
                // Have to subdivide
                let sub_child = NTreeNode::<N> {
                    color: current_node.color,
                    subdivided: false,
                };
                // Set node as subdivided
                self.structure.get_mut(&current_loc).unwrap().subdivided = true;
                current_node.subdivided = true;

                // Insert new children
                current_loc.iter_children().for_each(|x| {
                    self.structure.insert(x, sub_child);
                });
            }

            if current_node.subdivided
            {
                size /= N;
                let child_x = local_x / size;
                let child_y = local_y / size;
                let child_z = local_z / size;

                // Node is subdivided, descend
                current_loc = current_loc.get_child(child_x, child_y, child_z);
                current_node = *self.structure.get(&current_loc).unwrap();
                current_depth += 1;

                local_x -= child_x * size;
                local_y -= child_y * size;
                local_z -= child_z * size;
            }
        }

        // Insertion has been done, travel back up to optimise
        loop
        {
            // Try to simplify, compute average
            if current_node.subdivided
            {
                let (compressable, color) = current_loc
                    .iter_children()
                    .map(|x| self.structure.get(&x).unwrap())
                    .fold((true, None), |state, child| match state
                    {
                        (_, None) => (!child.subdivided, Some(child.color)),
                        (b, Some(c)) => (b && !child.subdivided && c == child.color, Some(c)),
                    });

                if compressable
                {
                    // Compress
                    *self.structure.get_mut(&current_loc).unwrap() = NTreeNode {
                        color: color.unwrap(),
                        subdivided: false,
                    };

                    // Remove children
                    current_loc.iter_children().for_each(|x| {
                        self.structure.remove(&x).unwrap();
                    });
                }
                else
                {
                    // Update average color
                    let colors = current_loc
                        .iter_children()
                        .map(|x| self.structure.get(&x).unwrap().color)
                        .collect::<Vec<_>>();
                    self.structure.get_mut(&current_loc).unwrap().color =
                        Color::average(colors.as_slice());
                }
            }

            // Travel up
            if current_depth == 0
            {
                break; // Finished
            }

            current_loc = current_loc.get_parent();
            current_node = *self.structure.get(&current_loc).unwrap();
            current_depth -= 1;
        }
    }

    pub fn to_gpu_rep(&self) -> (Color, Vec<GPUStructureTile<N>>, Vec<[Color; N * N * N]>)
    {
        // No root + group by child group
        let tile_count = (self.structure.len() - 1) / (N * N * N);
        if tile_count == 0
        {
            return (
                self.structure
                    .get(&NTreeLocator::<N>::root())
                    .unwrap()
                    .color,
                vec![],
                vec![],
            );
        }

        let mut structure_tiles = vec![GPUStructureTile::<N>::zero(); tile_count];
        let mut color_tiles = vec![[Color(0., 0., 0., 0.); N * N * N]; tile_count];

        let mut current_tile = 0usize;
        let mut queue = VecDeque::new();
        queue.push_back((NTreeLocator::<N>::root(), current_tile));
        current_tile += 1;
        while !queue.is_empty()
        {
            let (loc, dest_tile) = queue.pop_front().unwrap();

            let children = loc
                .iter_children()
                .map(|x| (x, self.structure.get(&x).unwrap()))
                .collect::<Vec<_>>();

            let mut structure_tile = GPUStructureTile::<N>::zero();
            let mut color_tile = [Color(0., 0., 0., 0.); N * N * N];

            for (i, (child_loc, child)) in children.into_iter().enumerate()
            {
                structure_tile.children[i] = GPUStructureTileIndex::new(
                    child.subdivided,
                    child.subdivided,
                    current_tile as u32,
                );
                color_tile[i] = child.color;

                if child.subdivided
                {
                    queue.push_back((child_loc, current_tile));
                    current_tile += 1;
                }
            }

            structure_tiles[dest_tile] = structure_tile;
            color_tiles[dest_tile] = color_tile;
        }

        println!(
            "node count: {}, current_tile: {}",
            self.structure.len(),
            current_tile
        );
        (
            self.structure
                .get(&NTreeLocator::<N>::root())
                .unwrap()
                .color,
            structure_tiles,
            color_tiles,
        )
    }
}

#[derive(Clone, Copy)]
pub struct GPUStructureTile<const N: usize>
where
    [(); N * N * N]:,
{
    children: [GPUStructureTileIndex; N * N * N],
}

impl<const N: usize> GPUStructureTile<N>
where
    [(); N * N * N]:,
{
    pub fn zero() -> Self
    {
        GPUStructureTile {
            children: [GPUStructureTileIndex::zero(); N * N * N],
        }
    }
}

#[derive(Clone, Copy)]
pub struct GPUStructureTileIndex(u32);

impl GPUStructureTileIndex
{
    pub fn zero() -> Self
    {
        GPUStructureTileIndex(0)
    }

    pub fn new(subdivided: bool, leaf: bool, ptr: u32) -> Self
    {
        assert_eq!((ptr >> 30), 0);
        GPUStructureTileIndex(ptr | ((leaf as u32) << 30) | ((subdivided as u32) << 31))
    }

    pub fn index(&self) -> u32
    {
        self.0 & 0x3FFFFFFFu32
    }

    pub fn subdivided(&self) -> bool
    {
        (self.0 >> 31) == 1
    }

    pub fn leaf(&self) -> bool
    {
        ((self.0 >> 30) & 1) == 1
    }
}

impl Color
{
    pub fn average(colors: &[Color]) -> Color
    {
        let sum = colors
            .iter()
            .copied()
            .reduce(|Color(r_a, g_a, b_a, a_a), Color(r_b, g_b, b_b, a_b)| {
                Color(r_a + r_b, g_a + g_b, b_a + b_b, a_a + a_b)
            })
            .unwrap_or(Color(0., 0., 0., 0.));
        //let len = colors.len().max(1) as f32;
        let len = colors.iter().map(|x| x.3).sum::<f32>();

        Color(sum.0 / len, sum.1 / len, sum.2 / len, sum.3 / len)
    }
}

#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
pub struct NTreeLocator<const N: usize>(usize, usize, usize);

impl<const N: usize> NTreeLocator<N>
{
    pub fn root() -> Self
    {
        Self(1, 1, 1)
    }

    pub fn get_child(&self, child_x: usize, child_y: usize, child_z: usize) -> Self
    {
        assert!(child_x < N && child_y < N && child_z < N);
        let mut new_loc_x = self.0;
        let mut new_loc_y = self.1;
        let mut new_loc_z = self.2;

        // Shift to left three times
        new_loc_x *= N;
        new_loc_x += child_x;

        new_loc_y *= N;
        new_loc_y += child_y;

        new_loc_z *= N;
        new_loc_z += child_z;

        Self(new_loc_x, new_loc_y, new_loc_z)
    }

    pub fn from_depth_coords(x: usize, y: usize, z: usize, depth: usize) -> Self
    {
        if depth == 0
        {
            return Self::root();
        }
        let off = N.pow(depth as u32);
        Self(off + x, off + y, off + z)
    }

    pub fn iter_children(&self) -> impl Iterator<Item = NTreeLocator<N>>
    {
        (0..N)
            .cartesian_product(0..N)
            .cartesian_product(0..N)
            .map(|((x, y), z)| self.get_child(x, y, z))
    }

    pub fn get_parent(&self) -> Self
    {
        if self.0 == 0
        {
            *self
        }
        else
        {
            Self(self.0 / N, self.1 / N, self.2 / N)
        }
    }

    pub fn get_local_location(&self) -> (usize, usize, usize)
    {
        let mut loc = self.0;
        let z = loc % N;
        loc /= N;

        let y = loc % N;
        loc /= N;

        let x = loc % N;

        (x, y, z)
    }

    pub fn is_root(&self) -> bool
    {
        self.0 == 1
    }
}

fn trailing_zeroes<const N: usize>(mut n: usize) -> Option<usize>
{
    if n == 0
    {
        return None;
    }
    let mut i = 0;
    while n.is_multiple_of(N)
    // n % N == 0
    {
        i += 1;
        n /= N;
    }
    Some(i)
}

#[cfg(test)]
mod test
{

    use crate::voxel::trailing_zeroes;
    use itertools::Itertools;
    use rand::Rng;

    use crate::voxel::Color;
    use crate::voxel::NTree;
    use crate::voxel::NTreeLocator;

    #[test]
    pub fn constant()
    {
        let color = Color(0.5, 0.3, 0.5, 1.);
        let depth = 5;
        const N: usize = 3;
        let width = N.pow(depth);
        let ntree = NTree::<N>::constant(color, depth);

        for ((x, y), z) in (0..width)
            .cartesian_product(0..width)
            .cartesian_product(0..width)
        {
            assert_eq!(ntree.get(x, y, z), color);
        }
    }

    #[test]
    pub fn full_insert()
    {
        const DEPTH: u32 = 4;
        const N: usize = 3;
        const WIDTH: usize = N.pow(DEPTH);
        let mut ntree = NTree::<N>::constant(Color(1., 0., 0., 0.), DEPTH);

        let mut rng = rand::rng();
        let mut storage = vec![vec![vec![Color(0., 0., 0., 0.); WIDTH]; WIDTH]; WIDTH];
        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let color = Color(rng.random(), rng.random(), rng.random(), rng.random());
            storage[x][y][z] = color;
            ntree.set(x, y, z, color);
        }

        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let color = ntree.get(x, y, z);
            assert_eq!(storage[x][y][z], color);
        }
        println!("Total nodes {}", ntree.structure.len());
    }

    #[test]
    pub fn full_insert_bottom_up()
    {
        const DEPTH: u32 = 4;
        const N: usize = 3;
        const WIDTH: usize = N.pow(DEPTH);

        let mut rng = rand::rng();
        let mut storage = vec![Color(0., 0., 0., 0.); WIDTH * WIDTH * WIDTH];
        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let color = Color(rng.random(), rng.random(), rng.random(), rng.random());
            storage[x + WIDTH * y + WIDTH * WIDTH * z] = color;
        }
        let ntree = NTree::<N>::from_arrays(&storage, DEPTH as usize);

        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let color = ntree.get(x, y, z);
            assert_eq!(storage[x + WIDTH * y + WIDTH * WIDTH * z], color);
        }
        println!("Total nodes {}", ntree.structure.len());
    }

    #[test]
    pub fn full_insert_const()
    {
        const DEPTH: u32 = 4;
        const N: usize = 3;
        const WIDTH: usize = N.pow(DEPTH);
        const NEW_COLOR: Color = Color(0., 1., 0., 1.);
        let mut ntree = NTree::<N>::constant(Color(1., 0., 1., 0.), DEPTH);

        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            ntree.set(x, y, z, NEW_COLOR);
        }

        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let color = ntree.get(x, y, z);
            assert_eq!(NEW_COLOR, color);
        }

        println!("Total nodes {}", ntree.structure.len());
    }

    #[test]
    pub fn full_insert_quantized()
    {
        const DEPTH: u32 = 4;
        const N: usize = 3;
        const WIDTH: usize = N.pow(DEPTH);
        let mut ntree = NTree::<N>::constant(Color(0., 0., 0., 0.), DEPTH);

        let mut rng = rand::rng();
        let mut storage = vec![vec![vec![Color(0., 0., 0., 0.); WIDTH]; WIDTH]; WIDTH];
        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let nbr = rng.random::<f32>();
            let mut color = Color(0., 1., 0., 0.);
            // Only 1 percent of different color
            if nbr > 0.01
            {
                color = Color(1., 0., 0., 0.);
            }
            storage[x][y][z] = color;
            ntree.set(x, y, z, color);
        }

        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let color = ntree.get(x, y, z);
            assert_eq!(storage[x][y][z], color);
        }
        println!("Total nodes {}", ntree.structure.len());
    }

    #[test]
    pub fn gpu_rep_quantized()
    {
        const DEPTH: u32 = 4;
        const N: usize = 3;
        const WIDTH: usize = N.pow(DEPTH);
        let mut ntree = NTree::<N>::constant(Color(0., 0., 0., 0.), DEPTH);

        let mut rng = rand::rng();
        let mut storage = vec![vec![vec![Color(0., 0., 0., 0.); WIDTH]; WIDTH]; WIDTH];
        for ((x, y), z) in (0..WIDTH)
            .cartesian_product(0..WIDTH)
            .cartesian_product(0..WIDTH)
        {
            let nbr = rng.random::<f32>();
            let mut color = Color(0., 1., 0., 0.);
            // Only 1 percent of different color
            if nbr > 0.01
            {
                color = Color(1., 0., 0., 0.);
            }
            storage[x][y][z] = color;
            ntree.set(x, y, z, color);
        }

        let (_color, _a, _b) = ntree.to_gpu_rep();

        drop(_a);
        drop(_b);
    }

    #[test]
    pub fn tree_locator_from_depth_coords()
    {
        const N: usize = 4;
        assert_eq!(
            NTreeLocator::<N>::root(),
            NTreeLocator::<N>::from_depth_coords(1, 2, 3, 0)
        );

        assert_eq!(
            NTreeLocator::<N>::root().get_child(1, 2, 3),
            NTreeLocator::<N>::from_depth_coords(1, 2, 3, 1)
        );

        assert_eq!(
            NTreeLocator::<N>::root()
                .get_child(1, 1, 1)
                .get_child(1, 2, 3),
            NTreeLocator::<N>::from_depth_coords(4 + 1, 4 + 2, 4 + 3, 2)
        );
    }

    #[test]
    fn trailing_zeroes_base10()
    {
        let n = 2139800000;
        assert_eq!(trailing_zeroes::<10>(n), Some(5));
        assert_eq!(trailing_zeroes::<10>(0), None);
    }
}
