mod base83;
mod srgb_lookup;

use std::f32::consts::PI;

use srgb_lookup::srgb_to_linear;

pub struct Components {
    pub x: u32,
    pub y: u32,
}

#[derive(Clone, Copy, Debug)]
pub struct ImageBounds {
    pub width: u32,
    pub height: u32,
}

struct ComponentState {
    x: u32,
    y: u32,
    basis: f32,
}

#[derive(Debug)]
pub struct ComponentError;

impl std::fmt::Display for ComponentError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Components out of bounds")
    }
}

pub struct Encoder {
    index: usize,
    components: Components,
    factors: Box<[(ComponentState, [f32; 3])]>,
    bounds: ImageBounds,
}

pub fn encode(
    components: Components,
    bounds: ImageBounds,
    rgba8_image: &[u8],
) -> Result<String, ComponentError> {
    let mut encoder = Encoder::new(components, bounds)?;
    encoder.update(rgba8_image);
    Ok(encoder.finalize())
}

impl Encoder {
    pub fn new(
        Components { x, y }: Components,
        bounds: ImageBounds,
    ) -> Result<Self, ComponentError> {
        if !(1..=9).contains(&x) || !(1..=9).contains(&y) {
            return Err(ComponentError);
        }

        Ok(Self {
            index: 0,
            components: Components { x, y },
            factors: Box::from(
                (0..y)
                    .flat_map(|y| {
                        (0..x).map(move |x| (ComponentState { x, y, basis: 0. }, [0., 0., 0.]))
                    })
                    .collect::<Vec<_>>(),
            ),
            bounds,
        })
    }

    pub fn update(&mut self, rgba8_image: &[u8]) {
        const BYTES_PER_PIXEL: usize = 4;

        // get offset in terms of already-processed bytes
        let offset = self.index % BYTES_PER_PIXEL;
        // get offset in terms of remaining bytes on head of rgba8_image
        let offset = (BYTES_PER_PIXEL - offset) % BYTES_PER_PIXEL;

        for (ComponentState { basis, .. }, [_, g, b]) in self.factors.iter_mut() {
            for (byte, value) in rgba8_image[..offset].iter().zip(
                [&mut *b, &mut *g][..offset.saturating_sub(BYTES_PER_PIXEL - 2)]
                    .iter_mut()
                    .rev(),
            ) {
                **value += *basis * srgb_to_linear(*byte);
            }
        }

        let pixels = ((self.index + offset) / BYTES_PER_PIXEL) as u32;

        let mut chunks = rgba8_image[offset..].chunks_exact(BYTES_PER_PIXEL);

        for (i, chunk) in (&mut chunks).enumerate() {
            let px = pixels + i as u32;
            let px_x = px % self.bounds.width;
            let px_y = px / self.bounds.width;

            for (ComponentState { x, y, .. }, [r, g, b]) in self.factors.iter_mut() {
                let basis = compute_basis(
                    *x as _,
                    *y as _,
                    px_x as _,
                    px_y as _,
                    self.bounds.width as _,
                    self.bounds.height as _,
                );

                *r += basis * srgb_to_linear(chunk[0]);
                *g += basis * srgb_to_linear(chunk[1]);
                *b += basis * srgb_to_linear(chunk[2]);
            }
        }

        if !chunks.remainder().is_empty() {
            let px = pixels + (rgba8_image[offset..].len() / BYTES_PER_PIXEL) as u32;
            let px_x = px % self.bounds.width;
            let px_y = px / self.bounds.width;

            for (ComponentState { x, y, basis }, [r, g, b]) in self.factors.iter_mut() {
                *basis = compute_basis(
                    *x as _,
                    *y as _,
                    px_x as _,
                    px_y as _,
                    self.bounds.width as _,
                    self.bounds.height as _,
                );

                for (byte, value) in chunks.remainder().iter().zip([&mut *r, &mut *g, &mut *b]) {
                    *value += *basis * srgb_to_linear(*byte);
                }
            }
        }

        self.index += rgba8_image.len();
    }

    pub fn finalize(self) -> String {
        let mut blurhash = String::new();

        let (_, dc) = self.factors[0];
        let ac = &self.factors[1..];

        let size_flag = self.components.x - 1 + (self.components.y - 1) * 9;
        base83::encode(size_flag, 1, &mut blurhash);

        let maximum = ac.iter().fold(
            0.0_f32,
            |maximum, (ComponentState { x, y, .. }, [r, g, b])| {
                let normalisation = if *x == 0 && *y == 0 { 1. } else { 2. };

                let scale = normalisation / (self.bounds.width * self.bounds.height) as f32;

                let r = r * scale;
                let g = g * scale;
                let b = b * scale;

                maximum.max(r.abs()).max(g.abs()).max(b.abs())
            },
        );

        let quantized_maximum = (maximum * 166. - 0.5).floor().max(0.) as u32;

        base83::encode(quantized_maximum, 1, &mut blurhash);

        let maximum_value = (quantized_maximum + 1) as f32 / 166.;

        base83::encode(encode_dc(dc), 4, &mut blurhash);

        for (_, rgb) in ac {
            base83::encode(encode_ac(*rgb, maximum_value), 2, &mut blurhash);
        }

        blurhash
    }
}

fn compute_basis(
    component_x: f32,
    component_y: f32,
    px_x: f32,
    px_y: f32,
    width: f32,
    height: f32,
) -> f32 {
    f32::cos(PI * component_x * px_x / width) * f32::cos(PI * component_y * px_y / height)
}

fn encode_dc([r, g, b]: [f32; 3]) -> u32 {
    let r = linear_to_srgb(r);
    let g = linear_to_srgb(g);
    let b = linear_to_srgb(b);

    (r << 16) + (g << 8) + b
}

fn encode_ac([r, g, b]: [f32; 3], maximum_value: f32) -> u32 {
    let r = encode_ac_digit(r, maximum_value);
    let g = encode_ac_digit(g, maximum_value);
    let b = encode_ac_digit(b, maximum_value);

    r * 19 * 19 + g * 19 + b
}

fn encode_ac_digit(d: f32, maximum_value: f32) -> u32 {
    ((sign_pow(d / maximum_value, 0.5) * 9. + 9.5) as i32).clamp(0, 18) as u32
}

pub fn linear_to_srgb(value: f32) -> u32 {
    let v = f32::max(0., f32::min(1., value));
    if v <= 0.003_130_8 {
        (v * 12.92 * 255. + 0.5).round() as u32
    } else {
        ((1.055 * f32::powf(v, 1. / 2.4) - 0.055) * 255. + 0.5).round() as u32
    }
}

fn sign(n: f32) -> f32 {
    if n < 0. {
        -1.
    } else {
        1.
    }
}

pub fn sign_pow(val: f32, exp: f32) -> f32 {
    sign(val) * val.abs().powf(exp)
}

#[cfg(test)]
mod tests {
    use image::{EncodableLayout, GenericImageView};

    #[test]
    fn contrived() {
        let input = [
            0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120,
            0, 0, 60, 120, 0, 0, 60, 120, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60,
            0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0,
        ];
        let width = 4;
        let height = 4;

        let b1 = super::encode(
            crate::Components { x: 4, y: 3 },
            crate::ImageBounds { width, height },
            &input,
        )
        .unwrap();

        let b2 = blurhash::encode(4, 3, width, height, &input).unwrap();

        assert_eq!(b1, b2);
    }

    #[test]
    fn one_component() {
        let inputs = [
            "data/19dd1c444d1c7939.png",
            "data/f73d2ee39133d871.jpg",
            "data/shenzi.png",
        ];

        for input in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();

            let b1 = super::encode(
                crate::Components { x: 1, y: 1 },
                crate::ImageBounds { width, height },
                img.to_rgba8().as_bytes(),
            )
            .unwrap();

            let b2 = blurhash::encode(1, 1, width, height, img.to_rgba8().as_bytes()).unwrap();

            assert_eq!(b1, b2);
        }
    }

    #[test]
    fn matches_blurhash() {
        let inputs = [
            "data/19dd1c444d1c7939.png",
            "data/f73d2ee39133d871.jpg",
            "data/shenzi.png",
        ];

        for input in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();

            let b1 = super::encode(
                crate::Components { x: 4, y: 3 },
                crate::ImageBounds { width, height },
                img.to_rgba8().as_bytes(),
            )
            .unwrap();

            let b2 = blurhash::encode(4, 3, width, height, img.to_rgba8().as_bytes()).unwrap();

            assert_eq!(b1, b2);
        }
    }

    #[test]
    fn matches_self_when_split() {
        let inputs = [
            "data/19dd1c444d1c7939.png",
            "data/f73d2ee39133d871.jpg",
            "data/shenzi.png",
        ];

        for input in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();
            let rgba8_img = img.to_rgba8();
            let bytes = rgba8_img.as_bytes();

            let b1 = super::encode(
                crate::Components { x: 4, y: 3 },
                crate::ImageBounds { width, height },
                bytes,
            )
            .unwrap();

            for chunk_count in 2..20 {
                let mut encoder = super::Encoder::new(
                    crate::Components { x: 4, y: 3 },
                    crate::ImageBounds { width, height },
                )
                .unwrap();

                let chunk_size = bytes.len() / chunk_count;

                for chunk in bytes.chunks(chunk_size) {
                    encoder.update(chunk);
                }

                let b2 = encoder.finalize();

                assert_eq!(b1, b2);
            }
        }
    }
}
