iMessage Word Hunt solver weekend funtime

main.rs

use std::{
    collections::HashMap,
    fs::File,
    io::{BufRead, BufReader},
    process::Command,
    time::Duration,
};

use dbus::blocking::Connection;
use image::io::Reader as ImageReader;
use image::DynamicImage;
use std::time::Instant;

const INTER_CARDINAL_MOVEMENT_STEP: u16 = 78;
const CARDINAL_MOVEMENT_STEP: u16 = 82;

fn main() {
    let words = Words::read();

    let img = ImageReader::open("./frame.jpg").unwrap().decode().unwrap();

    let board = BoardA::build(img);

    dbg!("solving...");

    let start = Instant::now();

    let mut solves = board.solve(words);

    println!("solved in {:?}", start.elapsed());

    solves.retain(|x| x.path.len() > 2);
    solves.sort_by(|a, b| b.path.len().cmp(&a.path.len()));

    let mouse_client = MouseClient::new().unwrap();

    for solve in solves {
        println!("solving: {}", solve.as_string(&board));
        let mouse_moves = solve.into_mouse_movements();
        for mouse_move in mouse_moves {
            mouse_client.send_mouse_movement(mouse_move).unwrap();
            std::thread::sleep(Duration::from_millis(50));
        }
    }
}

#[derive(Default, Debug)]
pub struct WordNode {
    children: HashMap<char, WordNode>,
    is_tail: bool,
}

impl WordNode {
    pub fn new() -> Self {
        Self {
            children: Default::default(),
            is_tail: false,
        }
    }

    pub fn append(&mut self, mut chars: impl Iterator<Item = char>) {
        match chars.next() {
            Some(chr) => self.children.entry(chr).or_default().append(chars),
            None => self.is_tail = true,
        }
    }

    pub fn is_tail(&self) -> bool {
        self.is_tail
    }

    pub fn child(&self, chr: char) -> Option<&WordNode> {
        self.children.get(&chr)
    }

    pub fn test(&self, mut prefix: impl Iterator<Item = char>) -> Option<&WordNode> {
        let next = prefix.next();
        dbg!(next);
        if let Some(next_char) = next {
            if let Some(next_node) = self.child(next_char) {
                dbg!(next_node.is_tail, next_node.children.len());
                next_node.test(prefix)
            } else {
                None
            }
        } else {
            Some(self)
        }
    }
}

#[cfg(test)]
pub mod test {
    use super::*;

    #[test]
    fn test_word_node() {
        let words = Words::read();
        dbg!(words.root.test("DALANT".chars()));
    }
}

#[derive(Debug)]
pub struct Words {
    root: WordNode,
}

impl Words {
    pub fn read() -> Words {
        let file = File::open("./words.txt").unwrap();
        let lines = BufReader::new(file).lines();

        let mut this = Self {
            root: Default::default(),
        };

        for line in lines {
            let line = line.unwrap();
            let chars = line.chars();
            this.root.append(chars);
        }

        this
    }

    pub fn root(&self) -> &WordNode {
        &self.root
    }
}

#[derive(Eq, PartialEq, Clone, Copy, Debug, Hash)]
pub struct Pos {
    x: i8,
    y: i8,
}

impl Pos {
    pub const fn new(x: i8, y: i8) -> Self {
        Self { x, y }
    }

    pub fn translate(&self, other: Pos) -> Pos {
        Pos {
            x: self.x + other.x,
            y: self.y + other.y,
        }
    }
}

pub struct ImageRect {
    x: u32,
    y: u32,
    w: u32,
    h: u32,
}

impl ImageRect {
    pub fn new(x: u32, y: u32, w: u32, h: u32) -> Self {
        Self { x, y, w, h }
    }
}

const MOVES: [Pos; 8] = [
    Pos::new(-1, -1),
    Pos::new(0, -1),
    Pos::new(1, -1),
    Pos::new(-1, 0),
    Pos::new(1, 0),
    Pos::new(-1, 1),
    Pos::new(0, 1),
    Pos::new(1, 1),
];

#[derive(Debug)]
pub struct BoardA {
    grid: Grid,
}

impl BoardA {
    pub fn grid_mask() -> Vec<(Pos, ImageRect)> {
        let mut grid = vec![];
        for i in 0..4_u32 {
            for j in 0..4_u32 {
                grid.push((
                    Pos::new(j as i8, i as i8),
                    ImageRect::new(76 + (j * 91), 509 + (i * 91), 77, 77),
                ));
            }
        }
        grid
    }

    pub fn get_grid_char(&self, pos: Pos) -> Option<char> {
        self.grid.get(&pos).copied()
    }

    pub fn build(img: DynamicImage) -> Self {
        let grid = img_to_grid(img, Self::grid_mask());
        dbg!(&grid);
        Self { grid }
    }

    pub fn solve(&self, word_tree: Words) -> Vec<FoundWord> {
        let mut words = Vec::new();

        //todo: rayon all these, concat results, and dedupe
        for (Pos { x, y }, chr) in self.grid.iter() {
            let pos = Pos::new(*x, *y);
            let root_node = word_tree
                .root()
                .child(*chr)
                .expect("every letter should have a starting word");
            let cons = Cons::new(pos);

            self.traverse(&mut words, root_node, cons, pos);
        }

        words
    }

    pub fn traverse<'a>(
        &self,
        words: &mut Vec<FoundWord>,
        node: &WordNode,
        cons: Cons<'a>,
        pos: Pos,
    ) {
        if node.is_tail {
            words.push(FoundWord::new(cons.to_vec()));
        }

        for mov in MOVES {
            let new_pos = pos.translate(mov);
            if let Some(chr) = self.get_grid_char(new_pos) {
                if let Some(new_cons) = cons.next(new_pos) {
                    if let Some(child) = node.child(chr) {
                        self.traverse(words, child, new_cons, new_pos);
                    }
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct FoundWord {
    path: Vec<Pos>,
}

impl FoundWord {
    fn new(path: Vec<Pos>) -> Self {
        Self { path }
    }

    fn as_string(&self, board: &BoardA) -> String {
        let mut output = String::with_capacity(self.path.len());
        for &pos in &self.path {
            let char = board.get_grid_char(pos).unwrap();
            output.push(char);
        }
        output
    }

    fn into_mouse_movements(self) -> Vec<MouseMovement> {
        let mut movements = Vec::new();
        let mut last_pos = Pos::new(0, 0);

        movements.push(MouseMovement::MoveToTopLeftBoard);

        let first_pos = self.path.first().unwrap();
        for _ in 0..first_pos.x {
            movements.push(MouseMovement::MoveCardinal(
                CARDINAL_MOVEMENT_STEP,
                Cardinal::E,
                ButtonState::Up,
            ));
        }

        for _ in 0..first_pos.y {
            movements.push(MouseMovement::MoveCardinal(
                CARDINAL_MOVEMENT_STEP,
                Cardinal::S,
                ButtonState::Up,
            ));
        }

        movements.push(MouseMovement::Button(ButtonState::Down));

        for (_, pos) in self.path.iter().enumerate() {
            let dx = pos.x - last_pos.x;
            let dy = pos.y - last_pos.y;

            last_pos = *pos;

            match (dx, dy) {
                (1, 0) => movements.push(MouseMovement::MoveCardinal(
                    CARDINAL_MOVEMENT_STEP,
                    Cardinal::E,
                    ButtonState::Down,
                )),
                (-1, 0) => movements.push(MouseMovement::MoveCardinal(
                    CARDINAL_MOVEMENT_STEP,
                    Cardinal::W,
                    ButtonState::Down,
                )),
                (0, 1) => movements.push(MouseMovement::MoveCardinal(
                    CARDINAL_MOVEMENT_STEP,
                    Cardinal::S,
                    ButtonState::Down,
                )),
                (0, -1) => movements.push(MouseMovement::MoveCardinal(
                    CARDINAL_MOVEMENT_STEP,
                    Cardinal::N,
                    ButtonState::Down,
                )),
                (-1, -1) => movements.push(MouseMovement::MoveInterCardinal(
                    INTER_CARDINAL_MOVEMENT_STEP,
                    InterCardinal::NW,
                    ButtonState::Down,
                )),
                (1, -1) => movements.push(MouseMovement::MoveInterCardinal(
                    INTER_CARDINAL_MOVEMENT_STEP,
                    InterCardinal::NE,
                    ButtonState::Down,
                )),
                (-1, 1) => movements.push(MouseMovement::MoveInterCardinal(
                    INTER_CARDINAL_MOVEMENT_STEP,
                    InterCardinal::SW,
                    ButtonState::Down,
                )),
                (1, 1) => movements.push(MouseMovement::MoveInterCardinal(
                    INTER_CARDINAL_MOVEMENT_STEP,
                    InterCardinal::SE,
                    ButtonState::Down,
                )),
                _ => println!("unexpected movement: {:?}", (dx, dy)),
            }
        }

        movements.push(MouseMovement::Button(ButtonState::Up));

        movements
    }
}

#[derive(Clone, Debug)]
pub struct Cons<'a> {
    pos: Pos,
    parent: Option<&'a Cons<'a>>,
}

impl<'a> Cons<'a> {
    pub fn new(pos: Pos) -> Self {
        Self { pos, parent: None }
    }

    pub fn can_traverse(&self, pos: Pos) -> bool {
        if self.pos == pos {
            return false;
        }

        match self.parent {
            Some(parent) => parent.can_traverse(pos),
            None => true,
        }
    }

    pub fn next<'b: 'a>(&'b self, pos: Pos) -> Option<Cons<'_>> {
        if self.can_traverse(pos) {
            Some(Self {
                pos,
                parent: Some(self),
            })
        } else {
            None
        }
    }

    pub fn to_vec(&self) -> Vec<Pos> {
        let mut ret = Vec::new();
        let mut cursor = self;
        loop {
            ret.push(cursor.pos);
            match cursor.parent {
                Some(parent) => cursor = parent,
                None => break,
            };
        }

        ret.reverse();
        ret
    }
}

pub type Grid = HashMap<Pos, char>;

// turn image and mask positions to hashmap grid
fn img_to_grid(img: DynamicImage, mask: Vec<(Pos, ImageRect)>) -> Grid {
    let mut grid: Grid = HashMap::new();

    for (pos, rect) in mask {
        let cropped_img = img.crop_imm(rect.x, rect.y, rect.w, rect.h);
        let mut letter_img = cropped_img.to_rgb8();
        for pixel in letter_img.pixels_mut() {
            let sum: u32 = pixel.0.iter().map(|i| *i as u32).sum();
            if sum > 50 {
                pixel.0 = [255, 255, 255];
            }
        }
        letter_img
            .save_with_format(
                format!("./cap/{}-{}.png", pos.x, pos.y),
                image::ImageFormat::Png,
            )
            .unwrap();
        letter_img
            .save_with_format("./t.png", image::ImageFormat::Png)
            .unwrap();
        let output = Command::new("/usr/bin/tesseract")
            .arg("t.png")
            .arg("stdout")
            .arg("--psm")
            .arg("10")
            .arg("--oem")
            .arg("0")
            .arg("-l")
            .arg("eng")
            .arg("--user-patterns")
            .arg("patterns.txt")
            .output()
            .unwrap();

        let output_char = String::from_utf8_lossy(&output.stdout)
            .to_string()
            .trim()
            .chars()
            .next()
            .unwrap()
            .to_ascii_uppercase();

        grid.insert(pos, output_char);
    }

    grid
}

pub struct MouseClient {
    connection: Connection,
}

impl MouseClient {
    fn new() -> Result<Self, dbus::Error> {
        let connection = Connection::new_system()?;
        Ok(Self { connection })
    }

    fn send_mouse(&self, mouse_command: MouseCommand) -> Result<(), dbus::Error> {
        let proxy = self.connection.with_proxy(
            "org.thanhle.btkbservice",
            "/org/thanhle/btkbservice",
            Duration::from_millis(50),
        );

        // let keys: Vec<u8> = vec![];

        println!("  -> executing mouse command: {:?}", mouse_command);
        proxy.method_call(
            "org.thanhle.btkbservice",
            "send_mouse",
            (0, &mouse_command.into_vec()),
        )?;

        Ok(())
    }

    fn send_mouse_movement(&self, mouse_movement: MouseMovement) -> Result<(), dbus::Error> {
        println!("executing mouse movement: {:?}", mouse_movement);
        for command in mouse_movement.into_commands() {
            self.send_mouse(command)?;
            std::thread::sleep(Duration::from_millis(100));
        }
        println!("mouse movement complete\n");
        Ok(())
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ButtonState {
    Down,
    Up,
}

#[derive(Debug)]
pub struct MouseCommand {
    button_state: ButtonState,
    x: i8,
    y: i8,
    z: i8,
}

impl MouseCommand {
    fn with_button_state(button_state: ButtonState) -> Self {
        Self {
            button_state,
            x: 0,
            y: 0,
            z: 0,
        }
    }

    fn into_bytes(self) -> [u8; 4] {
        [
            self.button_state.into(),
            self.x as u8,
            self.y as u8,
            self.z as u8,
        ]
    }
    fn into_vec(self) -> Vec<u8> {
        self.into_bytes().into()
    }
}

impl From<ButtonState> for u8 {
    fn from(val: ButtonState) -> Self {
        match val {
            ButtonState::Down => 1,
            ButtonState::Up => 0,
        }
    }
}

impl Default for ButtonState {
    fn default() -> Self {
        ButtonState::Up
    }
}

#[derive(Debug)]
pub enum MouseMovement {
    ResetTopLeft,
    ResetBottomRight,

    Button(ButtonState),

    MoveToTopLeftBoard,
    MoveY(i32, ButtonState),
    MoveInterCardinal(u16, InterCardinal, ButtonState),
    MoveCardinal(u16, Cardinal, ButtonState),
    MoveX(i32, ButtonState),
}

#[derive(Debug)]
pub enum InterCardinal {
    NW,
    NE,
    SW,
    SE,
}

#[derive(Debug)]
pub enum Cardinal {
    N,
    E,
    S,
    W,
}

const RESET_DISTANCE: u16 = 128 * 10;

impl MouseMovement {
    fn into_commands(self) -> Vec<MouseCommand> {
        match self {
            MouseMovement::Button(button_state) => {
                vec![MouseCommand::with_button_state(button_state)]
            }
            MouseMovement::ResetTopLeft => {
                MouseMovement::MoveInterCardinal(RESET_DISTANCE, InterCardinal::NW, ButtonState::Up)
                    .into_commands()
            }
            MouseMovement::ResetBottomRight => {
                MouseMovement::MoveInterCardinal(RESET_DISTANCE, InterCardinal::SE, ButtonState::Up)
                    .into_commands()
            }
            MouseMovement::MoveY(pixels, button_state) => {
                let mut result = vec![];
                let neg = pixels < 0;
                let mut pixels = pixels.abs();
                while pixels > 0 {
                    let step = pixels.min(i8::MAX as _) as i8;

                    result.push(MouseCommand {
                        button_state,
                        x: 0,
                        y: step * if neg { -1 } else { 1 },
                        z: 0,
                    });
                    pixels -= step as i32;
                }
                result
            }
            MouseMovement::MoveX(pixels, button_state) => {
                let mut result = vec![];
                let neg = pixels < 0;
                let mut pixels = pixels.abs();
                while pixels > 0 {
                    let step = pixels.min(i8::MAX as _) as i8;

                    result.push(MouseCommand {
                        button_state,
                        x: step * if neg { -1 } else { 1 },
                        y: 0,
                        z: 0,
                    });
                    pixels -= step as i32;
                }
                result
            }
            MouseMovement::MoveToTopLeftBoard => {
                let mut result = Vec::new();
                result.extend(MouseMovement::ResetTopLeft.into_commands());
                result.extend(MouseMovement::MoveY(463, ButtonState::Up).into_commands());
                result.extend(MouseMovement::MoveX(86, ButtonState::Up).into_commands());
                result
            }
            MouseMovement::MoveCardinal(pixels, direction, button_state) => match direction {
                Cardinal::N => MouseMovement::MoveY(-(pixels as i32), button_state).into_commands(),
                Cardinal::E => MouseMovement::MoveX(pixels as i32, button_state).into_commands(),
                Cardinal::S => MouseMovement::MoveY(pixels as i32, button_state).into_commands(),
                Cardinal::W => MouseMovement::MoveX(-(pixels as i32), button_state).into_commands(),
            },
            MouseMovement::MoveInterCardinal(mut pixels, direction, button_state) => {
                let mut result = Vec::new();
                while pixels > 0 {
                    let step = pixels.min(i8::MAX as _) as i8;

                    result.push(match direction {
                        InterCardinal::NW => MouseCommand {
                            button_state,
                            x: -step,
                            y: -step,
                            z: 0,
                        },
                        InterCardinal::NE => MouseCommand {
                            button_state,
                            x: step,
                            y: -step,
                            z: 0,
                        },
                        InterCardinal::SW => MouseCommand {
                            button_state,
                            x: -step,
                            y: step,
                            z: 0,
                        },
                        InterCardinal::SE => MouseCommand {
                            button_state,
                            x: step,
                            y: step,
                            z: 0,
                        },
                    });
                    pixels -= step as u16;
                }
                result
            }
        }
    }
}