use image::{Rgb, RgbImage};

use crate::indexed::redmean_sq;

pub type Color = [u8; 3]; // TODO: Convert to indexed

#[derive(Debug)]
pub struct EncodeTree {
    pub error: f32,
    pub bytes: f32,
    pub strategy: EncodeStrategy,
}

#[derive(Debug)]
pub enum EncodeStrategy {
    Solid(Color),
    // TODO: add Bitmap([Color; 2]),
    QuadSplit(Box<[EncodeTree; 4]>),
}

impl EncodeTree {
    pub fn apply(&self, img: &mut RgbImage, region: Region) {
        match &self.strategy {
            EncodeStrategy::Solid(color) => {
                for y in region.y..region.y2() {
                    for x in region.x..region.x2() {
                        img.put_pixel(x, y, Rgb(*color))
                    }
                }
            }
            EncodeStrategy::QuadSplit(subtrees) => {
                let width = region.width / 2;
                let height = region.height / 2;
                let coords = [
                    [region.x, region.y],
                    [region.x + width, region.y],
                    [region.x, region.y + height],
                    [region.x + width, region.y + height],
                ];
                for (tree, &[x, y]) in subtrees.iter().zip(coords.iter()) {
                    tree.apply(
                        img,
                        Region {
                            x,
                            y,
                            width,
                            height,
                        },
                    );
                }
            }
        }
    }
}

#[derive(Clone, Copy)]
pub struct Region {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

impl Region {
    pub fn x2(&self) -> u32 {
        self.x + self.width
    }

    pub fn y2(&self) -> u32 {
        self.y + self.height
    }

    pub fn area(&self) -> u32 {
        self.width * self.height
    }
}

pub fn choose_encoding(img: &RgbImage, region: Region, max_error: f32) -> EncodeTree {
    // Attempt to quantize as solid color
    let info = analyze_solid(img, region);
    if info.error <= max_error {
        return EncodeTree {
            error: info.error,
            bytes: 1.0,
            strategy: EncodeStrategy::Solid(info.color),
        };
    }

    // Error is too high, so we have to subdivide
    let mut subinfo = vec![];
    let width = region.width / 2;
    let height = region.height / 2;
    for y in [region.y, region.y + height] {
        for x in [region.x, region.x + width] {
            subinfo.push(choose_encoding(
                img,
                Region {
                    x,
                    y,
                    width,
                    height,
                },
                max_error / 4.0,
            ));
        }
    }
    EncodeTree {
        error: subinfo.iter().fold(0.0, |acc, info| acc + info.error),
        bytes: 1.0 + subinfo.iter().fold(0.0, |acc, info| acc + info.bytes),
        strategy: EncodeStrategy::QuadSplit(Box::new(subinfo.try_into().unwrap())),
    }
}

struct SolidInfo {
    color: Color,
    error: f32,
}

fn analyze_solid(img: &RgbImage, region: Region) -> SolidInfo {
    // Calculate average color
    let mut avg = [0.0, 0.0, 0.0];
    for y in region.y..region.y2() {
        for x in region.x..region.x2() {
            let pixel = img.get_pixel(x, y);
            for (i, component) in avg.iter_mut().enumerate() {
                *component += (pixel.0[i] as f32).powf(2.2);
            }
        }
    }
    for component in &mut avg {
        *component /= region.area() as f32;
    }
    let avg = avg.map(|x| x.powf(0.4545).round() as u8);

    // Measure squared error
    let mut err = 0.0;
    for y in region.y..region.y2() {
        for x in region.x..region.x2() {
            let pixel = img.get_pixel(x, y);
            err += redmean_sq(avg, pixel.0);
        }
    }

    SolidInfo {
        color: avg,
        error: err,
    }
}
