use crate::{
    block::Block,
    frame::Frame,
    helpers::{pixel::Pixel, vector::Vec2},
    view::View,
};
use rustdct::DctPlanner;

pub fn compress_block(block: &mut Block, size: Vec2<usize>) {
    match block {
        Block::Literal(p) => {
            let comp = lit_compress(size.x, size.y, &p);
            *block = Block::CompressedLiteral(comp)
        }
        Block::Split(box [a, b]) => {
            let vert = size.x > size.y;
            compress_block(
                a,
                if vert {
                    Vec2 {
                        x: size.x / 2,
                        y: size.y,
                    }
                } else {
                    Vec2 {
                        x: size.x,
                        y: size.y / 2,
                    }
                },
            );
            compress_block(
                b,
                if vert {
                    Vec2 {
                        x: size.x - size.x / 2,
                        y: size.y,
                    }
                } else {
                    Vec2 {
                        x: size.x,
                        y: size.y - size.y / 2,
                    }
                },
            );
        }
        _ => (),
    }
}

pub fn lit_compress(w: usize, h: usize, pixels: &[Pixel]) -> Vec<u8> {
    let mut out = vec![];
    for ci in 0..3 {
        let mut ch = vec![0.; w * h];
        for y in 0..h {
            for x in 0..w {
                let p = pixels[x + y * w];
                ch[x + y * w] = match ci {
                    0 => p.r,
                    1 => p.g,
                    2 => p.b,
                    _ => unreachable!(),
                } as f32
            }
        }

        norm_dct_channel(w, h, &mut ch);
        for i in 0..w * h {
            out.push(unsafe { std::mem::transmute(ch[i] as i8) });
        }
    }
    out
}
pub fn lit_decompress(compressed: &[u8], mut target: View<&mut Frame>) {
    let (w, h) = (target.size.x as usize, target.size.y as usize);
    for ci in 0..3 {
        let mut ch = compressed[ci * w * h..(ci + 1) * w * h]
            .iter()
            .map(|v| unsafe { std::mem::transmute::<_, i8>(*v) } as f32)
            .collect::<Vec<_>>();
        norm_idct_channel(w, h, &mut ch);
        for y in 0..h {
            for x in 0..w {
                let p = &mut target[Vec2 {
                    x: x as isize,
                    y: y as isize,
                }];
                *(match ci {
                    0 => &mut p.r,
                    1 => &mut p.g,
                    2 => &mut p.b,
                    _ => unreachable!(),
                }) = ch[x + y * w] as u8
            }
        }
    }
}

fn norm_dct_channel(w: usize, h: usize, values: &mut [f32]) {
    let d = DctPlanner::new().plan_dct2(w);
    let mut temp = vec![0.; d.get_scratch_len()];
    for c in values.chunks_exact_mut(w) {
        d.process_dct2_with_scratch(c, &mut temp)
    }

    let mut values_trans = transpose(w, h, values);
    let d = DctPlanner::new().plan_dct2(h);
    let mut temp = vec![0.; d.get_scratch_len()];
    for c in values_trans.chunks_exact_mut(h) {
        d.process_dct2_with_scratch(c, &mut temp)
    }

    let scale = 2. / (w as f32 * h as f32).sqrt();
    for (i, v) in values_trans.iter().enumerate() {
        values[i] = *v * scale * 127.0
    }
}

fn norm_idct_channel(h: usize, w: usize, values: &mut [f32]) {
    let scale = 2. / (w as f32 * h as f32).sqrt();
    for v in values.iter_mut() {
        *v /= scale * 127.0
    }

    let d = DctPlanner::new().plan_dct2(w);
    let mut temp = vec![0.; d.get_scratch_len()];
    for c in values.chunks_exact_mut(w) {
        d.process_dct3_with_scratch(c, &mut temp)
    }

    let mut values_trans = transpose(w, h, values);
    let d = DctPlanner::new().plan_dct2(h);
    let mut temp = vec![0.; d.get_scratch_len()];
    for c in values_trans.chunks_exact_mut(h) {
        d.process_dct3_with_scratch(c, &mut temp)
    }
    values.copy_from_slice(&values_trans)
}

fn transpose(w: usize, h: usize, values: &[f32]) -> Vec<f32> {
    let mut io = 0;
    let mut it;
    let mut transposed = vec![0.; values.len()];
    for row in 0..h {
        it = row;
        for _ in 0..w {
            transposed[it] = values[io];
            io += 1;
            it += h;
        }
    }
    transposed
}