use std::io::Write;

use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform};


struct PixelTile {
    // i32: representation that supports Walsh-Hadamard
    pixels: [i16; TILE_SZ2]
}
struct CoefTile {
    coefs: [i16; TILE_SZ2]
}

struct QuantTile {
    quants: [i16; TILE_SZ2]
}


pub fn compress<W: Write>(
    width: u32,
    height: u32,
    // n_components: u16,
    layers: &[&[u8]],
    writer: &mut protocol::ProtocolWriter<W>
) -> ProtocolWriterResult<()> {

    // validation
    for l in 0..layers.len() {
        assert!(layers[l].len() == width as usize * height as usize);
    }

    // write header
    writer.write_header(MAGIC)?;
    writer.write_u32_wide(width)?;
    writer.write_u32_wide(height)?;
    writer.write_u32_wide(layers.len() as u32)?;

    let mut tiles = vec![];
    for layer in layers.iter() {
        for x0 in (0..width).step_by(TILE_SZ) {
            for y0 in (0..height).step_by(TILE_SZ) {
                let pixel_tile = PixelTile::from_layer(
                    x0 as usize, y0 as usize, layer, 
                    width as usize, height as usize
                );
                let coef_tile = CoefTile::from_pixel_tile(&pixel_tile);
                let quant_tile = QuantTile::from_coef_tile(&coef_tile);
                tiles.push(quant_tile);
            }
        }
    }

    for t in tiles.iter() {
        t.write_zero(writer)?;
    }
    for t in tiles.iter() {
        t.write_rest(writer)?;
    }

    Ok(())
}

impl PixelTile {
    fn from_layer(
        x0: usize,
        y0: usize,
        layer: &[u8],
        width: usize,
        height: usize
    ) -> PixelTile {
        let mut pixels = [0; TILE_SZ2];
        for x in 0..TILE_SZ {
            for y in 0..TILE_SZ {
                let src_x = x0 + x;
                let src_y = y0 + y;
                pixels[y * TILE_SZ + x] = 
                    if src_x < width && src_y < height { 
                        layer[src_y * width + src_x] as i16
                    } else {
                        0
                    };
            }
        }

        return PixelTile { pixels };
    }
}

impl CoefTile {
    fn from_pixel_tile(pt: &PixelTile) -> CoefTile {
        let mut coefs = pt.pixels.clone();

        // rows
        for y in 0..TILE_SZ {
            transform::encode(&mut coefs, y * TILE_SZ, 1);
        }

        // columns
        for x in 0..TILE_SZ {
            transform::encode(&mut coefs, x, 8);
        }

        return CoefTile { coefs }
    }
}

impl QuantTile {
    fn from_coef_tile(pt: &CoefTile) -> QuantTile {
        QuantTile {
            quants: quantization::to_quantized(pt.coefs)
        }
    }

    fn write_zero<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
        writer.write_i16_wide(self.quants[0])?;
        Ok(())
    }

    fn write_rest<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
        for i in 1..self.quants.len() {
            writer.write_i16_packed(self.quants[i])?;
        }
        Ok(())
    }
}