bad-apple/encoder/src/main.rs

use std::{fs::File, io::{stdout, Write}};

use enum_map::{Enum, EnumMap};
use num_enum::TryFromPrimitive;

mod dec;
mod enc;
mod util;
pub use util::*;

const INSPECT_ENC: bool = false;
const INSPECT_DEC: bool = false;
const MAX_ERROR: usize = 0; // max wrong pixels
const MAX_LOSS: usize = 0; // highest "loss" value tried for all lossy encodings

const LOSSLESS_ENCODINGS: &[FrameEncoder] = &[
	enc::rle_horizontal,
	enc::rle_horizontal_ext,
	enc::rle_vertical,
	enc::rle_vertical_ext,
	enc::rle_vertical_var,
	enc::rle_vertical_16,
	enc::rle_diff_horizontal,
	enc::rle_diff_vertical,
	// enc::bg_strips_horizontal_16, // only works for the tiny display
	enc::bg_strips_horizontal_24, // intended for the 240x320 display
	// enc::tree_16,// turns out to be useless
];
const LOSSY_ENCODINGS: &[FrameEncoderLossy] = &[
	enc::fill_white,
	enc::fill_black,
	// todo: adapt for big display
	// enc::cell_diff_8_vertical,
	enc::cell_diff_4_vertical,
	enc::cell_diff_8_vertical_big,
];

fn main() {
	let frames = get_all_frames("../video/frames/");
	let encoded = encode(&frames);

	let mut stats: EnumMap<Encoding, u32> = EnumMap::default();
	let mut reader = 0;
	let mut last_frame = FRAME_0;
	let mut frame_index = 0;
	while reader < encoded.len() {
		let frame_type: Encoding = encoded[reader].try_into().unwrap();
		stats[frame_type] += 1;
		reader += 1;
		let decoder = get_matching_decoder(frame_type);
		last_frame = decoder(&last_frame, &encoded[reader..], &mut reader);
		if INSPECT_DEC {
			println!(
				"\n{frame_type:?}, error: {}, index: {frame_index}",
				frame_error(&frames[frame_index], &last_frame)
			);
			render_images(&frames[frame_index], &last_frame);
			wait_for_input();
			frame_index += 1;
		}
	}
	for (encoding, &frames) in stats.iter() {
		if frames > 0 {
			println!("{encoding:?} - {frames} frames");
		}
	}
	println!();
	println!(
		"{} frames, total {} bytes (~{} bpf)",
		frames.len(),
		encoded.len(),
		encoded.len() / frames.len()
	);

	let mut export_string = String::from("// Generated by the `encoder` rust app\n");
	for (encoding, count) in stats {
		if count > 0 {
			export_string += &format!("#define USE_{encoding:?}\n");
		}
	}
	export_string += "\n\ntypedef enum Encoding {\n";
	for (encoding, count) in stats {
		if count > 0 {
			export_string += &format!("\tEncoding_{encoding:?} = {},\n", encoding as u8);
		}
	}
	export_string += "} Encoding_t;\n\n";
	export_string += "const unsigned char video[] = {";
	// export_string += &format!("const unsigned char video[{}] = {{", encoded.len());
	let mut i = 99;
	for byte in encoded {
		if i > 15 {
			export_string += "\n\t";
			i = 0;
		}
		i += 1;
		export_string += &format!("{byte},");
	}
	export_string += "\n};\n";
	let mut file = File::create(format!("../{OUTPUT_DIR}/data.h")).unwrap();
	file.write_all(export_string.as_bytes()).unwrap();
}

fn encode(frames: &[Frame]) -> Vec<u8> {
	let mut out = Vec::new();
	let mut last_frame = FRAME_0;
	print!("encoding frames");
	let frame_count = frames.len();
	for (_i, frame) in frames.iter().enumerate() {
		print!("\rencoding frame {_i} of {frame_count}");
		stdout().flush().unwrap();
		let mut options = Vec::new();
		for encode in LOSSLESS_ENCODINGS.iter() {
			let encoded = encode(&last_frame, frame);
			let decode = get_matching_decoder(encoded.encoding);
			let decoded = decode(&last_frame, &encoded.data, &mut 0);
			let error = frame_error(frame, &decoded);
			if error == 0 {
				options.push(encoded);
			} else {
				dbg!(&encoded);
				eprintln!("{:?}, error: {error}, frame: {_i}", encoded.encoding);
				println!("{:?}, error: {error}, frame: {_i}", encoded.encoding);
				println!("original | decoded");
				render_images(&frame, &decoded);
				panic!("error in lossless compression");
			}
		}
		for encode in LOSSY_ENCODINGS.iter() {
			for loss in 0..MAX_LOSS {
				let encoded = encode(&last_frame, frame, loss);
				let decode = get_matching_decoder(encoded.encoding);
				let decoded = decode(&last_frame, &encoded.data, &mut 0);
				let error = frame_error(frame, &decoded);

				if error <= MAX_ERROR {
					options.push(encoded);
				} else {
					// higher loss value will mean more error so can be skipped
					break;
				}
			}
		}

		options.sort_by_key(|b| b.data.len());
		let best_encoding = options.into_iter().next().unwrap();
		if INSPECT_ENC {
			println!();
			println!(
				"{:?}, {} bytes",
				best_encoding.encoding,
				best_encoding.data.len() + 1
			);
			render_image(frame);
			wait_for_input();
		}
		let best_encoding = best_encoding.into_bytes();
		out.extend_from_slice(&best_encoding);
		last_frame = *frame;
	}
	println!();
	out
}

#[derive(Debug, TryFromPrimitive, Enum, Copy, Clone)]
#[repr(u8)]
enum Encoding {
	FillBlack,
	FillWhite,
	RLEHorizontal,
	RLEHorizontalExt,
	RLEVertical,
	RLEVerticalExt,
	RLEVerticalVar,
	RLEVertical16,
	RLEDiffHorizontal,
	RLEDiffVertical,
	BGStripsH16,
	BGStripsH24,
	// BGStripsV,
	// QuadTree,
	// DrawCommands,
	// CellDiff8H,
	CellDiff8V,
	CellDiff8VBig,
	// CellDiff4HH,
	// CellDiff4HV,
	// CellDiff4VH,
	CellDiff4VV,
	// CellDiff4VV_large,
	Tree16,
}

fn get_matching_decoder(encoding: Encoding) -> FrameDecoder {
	match encoding {
		Encoding::FillWhite => dec::fill_white,
		Encoding::FillBlack => dec::fill_black,
		Encoding::RLEHorizontal => dec::rle_horizontal,
		Encoding::RLEHorizontalExt => dec::rle_horizontal_ext,
		Encoding::RLEVertical => dec::rle_vertical,
		Encoding::RLEVerticalExt => dec::rle_vertical_ext,
		Encoding::RLEVerticalVar => dec::rle_vertical_var,
		Encoding::RLEVertical16 => dec::rle_vertical_16,
		Encoding::RLEDiffHorizontal => dec::rle_diff_horizontal,
		Encoding::RLEDiffVertical => dec::rle_diff_vertical,
		Encoding::BGStripsH16 => dec::bg_strips_horizontal_16,
		Encoding::BGStripsH24 => dec::bg_strips_horizontal_24,
		// Encoding::BGStripsV => todo!(),
		// Encoding::QuadTree => todo!(),
		// Encoding::DrawCommands => todo!(),
		// Encoding::CellDiff8H => todo!(),
		Encoding::CellDiff8V => dec::cell_diff_8_vertical,
		Encoding::CellDiff8VBig => dec::cell_diff_8_vertical_big,
		// Encoding::CellDiff4HH => todo!(),
		// Encoding::CellDiff4HV => todo!(),
		// Encoding::CellDiff4VH => todo!(),
		Encoding::CellDiff4VV => dec::cell_diff_4_vertical,
		Encoding::Tree16 => dec::tree_16,
	}
}

type FrameEncoder = fn(previous_frame: &Frame, new_frame: &Frame) -> EncodedFrame;
type FrameEncoderLossy = fn(previous_frame: &Frame, new_frame: &Frame, loss: usize) -> EncodedFrame;
type FrameDecoder = fn(previous_frame: &Frame, encoded_bytes: &[u8], reader: &mut usize) -> Frame;

#[derive(Debug)]
struct EncodedFrame {
	encoding: Encoding,
	data: Vec<u8>,
}

impl EncodedFrame {
	fn into_bytes(mut self) -> Vec<u8> {
		self.data.insert(0, self.encoding as u8);
		self.data
	}
}