bad-apple/encoder/src/util.rs

use std::{
	fs::{self, File},
	io::{stdin, BufReader},
};

use image::{self, DynamicImage, GenericImageView, ImageFormat, Rgba};

// pub const OUTPUT_DIR: &str = "ch32_decoder";
pub const OUTPUT_DIR: &str = "pico_decoder/src";
pub const WIDTH: usize = 320;
pub const HEIGHT: usize = 240;
pub const FRAME_SIZE: usize = WIDTH * HEIGHT;

pub type Frame = [[u8; HEIGHT]; WIDTH];
pub const FRAME_0: Frame = [[0; HEIGHT]; WIDTH];
pub const FRAME_1: Frame = [[1; HEIGHT]; WIDTH];

pub fn wait_for_input() {
	stdin().read_line(&mut String::new()).unwrap();
}

fn convert_pixel(rgba: Rgba<u8>) -> u8 {
	(rgba.0[0] > 128) as u8
}

pub fn convert_image(image: &DynamicImage) -> Frame {
	let mut frame = FRAME_0;
	for x in 0..WIDTH {
		for y in 0..HEIGHT {
			frame[x][y] = convert_pixel(image.get_pixel(x as u32, y as u32));
		}
	}
	frame
}

pub fn get_all_frames(path: &str) -> Vec<Frame> {
	let frame_count = fs::read_dir(path).unwrap().count();
	let mut frames = Vec::new();
	for i in 0..frame_count {
		let path = format!("{}frame_{:04}.png", path, i + 1);
		let file = BufReader::new(File::open(path).unwrap());
		let image = image::load(file, ImageFormat::Png).unwrap();
		frames.push(convert_image(&image));
	}
	frames
}

pub fn frame_error(real: &Frame, decoded: &Frame) -> usize {
	let mut error = 0;
	for x in 0..WIDTH {
		for y in 0..HEIGHT {
			if real[x][y] != decoded[x][y] {
				error += 1;
			}
		}
	}
	error
}

pub fn most_common_pixel(frame: &Frame) -> u8 {
	let mut white_pixels = 0;
	let mut black_pixels = 0;
	for col in frame {
		for &p in col {
			let p = p as u32;
			white_pixels += p;
			black_pixels += 1 - p;
		}
	}
	if white_pixels > black_pixels {
		1
	} else {
		0
	}
}

fn render_pixel_pair(img: &Frame, x: usize, y: usize) {
	let char = match (img[x][y * 2], img[x][y * 2 + 1]) {
		(0, 0) => " ",
		(0, 1) => "▄",
		(1, 0) => "▀",
		(1, 1) => "█",
		_ => panic!("image contained nonbinary bytes"),
	};
	print!("{}", char);
}

pub fn render_image(img: &Frame) {
	for y in 0..(HEIGHT / 2) {
		for x in 0..WIDTH {
			render_pixel_pair(img, x, y);
		}
		println!();
	}
}

pub fn render_images(left: &Frame, right: &Frame) {
	for y in 0..(HEIGHT / 2) {
		print!("|");
		for x in 0..WIDTH {
			render_pixel_pair(left, x, y);
		}
		print!("|");
		for x in 0..WIDTH {
			render_pixel_pair(right, x, y);
		}
		println!("|");
	}
}

pub fn rle_255_encode(raw: &[u8]) -> Vec<u8> {
	rle_encode(raw, 255)
}

pub fn rle_encode(raw: &[u8], max: u8) -> Vec<u8> {
	let mut encoded = Vec::new();
	let mut last_val = 0;
	let mut run = 0;
	for &val in raw {
		if val != last_val || run == max {
			encoded.push(run);
			if run == max && val == last_val {
				encoded.push(0);
			}
			run = 1;
		} else {
			run += 1;
		}
		last_val = val;
	}
	encoded.push(run);
	encoded
}

pub fn pack_nybbles(mut nybbles: Vec<u8>) -> Vec<u8> {
	if nybbles.len() % 2 == 1 {
		nybbles.push(0);
	}
	assert!(nybbles.iter().all(|&n| n < 16));
	let mut packed = Vec::with_capacity(nybbles.len() / 2);
	for i in 0..(nybbles.len() / 2) {
		let upper = nybbles[i * 2] << 4;
		let lower = nybbles[i * 2 + 1] & 15;
		let byte = upper | lower;
		packed.push(byte);
	}
	packed
}

pub fn unpack_nybbles(packed: &[u8]) -> Vec<u8> {
	let mut nybbles = Vec::with_capacity(packed.len() * 2);
	for &p in packed {
		nybbles.push(p >> 4);
		nybbles.push(p & 15);
	}
	nybbles
}

pub fn rle_255_decode(encoded: &[u8], max_size: usize) -> (usize, Vec<u8>) {
	let mut raw = Vec::new();
	let mut val = 0;
	let mut consumed_bytes = 0;
	for &run in encoded {
		consumed_bytes += 1;
		for _ in 0..run {
			raw.push(val);
		}
		if raw.len() >= max_size {
			break;
		}
		val = 1 - val;
	}
	(consumed_bytes, raw)
}