arimaa-engine 0.1.0

use anyhow::anyhow;
use std::fmt;
use std::str::FromStr;

use super::{Direction, Square};

// Prescribed as the first being the direction of the piece which vacates the square and the second being the direction of piece to occupy the square.
#[derive(Hash, Clone, Copy, Debug, Eq, PartialEq)]
pub enum PushPullDirection {
    NN,
    NE,
    NW,
    EN,
    EE,
    ES,
    SE,
    SS,
    SW,
    WN,
    WS,
    WW,
}

impl PushPullDirection {
    pub fn get_steps(&self, square: Square) -> [(Square, Direction); 2] {
        match self {
            PushPullDirection::NN => [(square, Direction::Up), (square.shift_in_direction(Direction::Down), Direction::Up)],
            PushPullDirection::NE => [(square, Direction::Up), (square.shift_in_direction(Direction::Left), Direction::Right)],
            PushPullDirection::NW => [(square, Direction::Up), (square.shift_in_direction(Direction::Right), Direction::Left)],
            PushPullDirection::EN => [(square, Direction::Right), (square.shift_in_direction(Direction::Down), Direction::Up)],
            PushPullDirection::EE => [(square, Direction::Right), (square.shift_in_direction(Direction::Left), Direction::Right)],
            PushPullDirection::ES => [(square, Direction::Right), (square.shift_in_direction(Direction::Up), Direction::Down)],
            PushPullDirection::SE => [(square, Direction::Down), (square.shift_in_direction(Direction::Left), Direction::Right)],
            PushPullDirection::SS => [(square, Direction::Down), (square.shift_in_direction(Direction::Up), Direction::Down)],
            PushPullDirection::SW => [(square, Direction::Down), (square.shift_in_direction(Direction::Right), Direction::Left)],
            PushPullDirection::WN => [(square, Direction::Left), (square.shift_in_direction(Direction::Down), Direction::Up)],
            PushPullDirection::WS => [(square, Direction::Left), (square.shift_in_direction(Direction::Up), Direction::Down)],
            PushPullDirection::WW => [(square, Direction::Left), (square.shift_in_direction(Direction::Right), Direction::Left)],
        }
    }

    pub fn as_u8(&self) -> u8 {
        match self {
            PushPullDirection::NN => 0,
            PushPullDirection::NE => 1,
            PushPullDirection::NW => 2,
            PushPullDirection::EN => 3,
            PushPullDirection::EE => 4,
            PushPullDirection::ES => 5,
            PushPullDirection::SE => 6,
            PushPullDirection::SS => 7,
            PushPullDirection::SW => 8,
            PushPullDirection::WN => 9,
            PushPullDirection::WS => 10,
            PushPullDirection::WW => 11,
        }
    }

    pub fn directions() -> Vec<PushPullDirection> {
        vec![
            PushPullDirection::NN,
            PushPullDirection::NE,
            PushPullDirection::NW,
            PushPullDirection::EN,
            PushPullDirection::EE,
            PushPullDirection::ES,
            PushPullDirection::SE,
            PushPullDirection::SS,
            PushPullDirection::SW,
            PushPullDirection::WN,
            PushPullDirection::WS,
            PushPullDirection::WW,
        ]
    }

    pub fn is_valid_from_square(&self, square: Square) -> bool {
        match self {
            PushPullDirection::NN => square.can_shift_in_direction(Direction::Up) && square.can_shift_in_direction(Direction::Down),
            PushPullDirection::NE => square.can_shift_in_direction(Direction::Up) && square.can_shift_in_direction(Direction::Left),
            PushPullDirection::NW => square.can_shift_in_direction(Direction::Up) && square.can_shift_in_direction(Direction::Right),
            PushPullDirection::EN => square.can_shift_in_direction(Direction::Right) && square.can_shift_in_direction(Direction::Down),
            PushPullDirection::EE => square.can_shift_in_direction(Direction::Right) && square.can_shift_in_direction(Direction::Left),
            PushPullDirection::ES => square.can_shift_in_direction(Direction::Right) && square.can_shift_in_direction(Direction::Up),
            PushPullDirection::SE => square.can_shift_in_direction(Direction::Down) && square.can_shift_in_direction(Direction::Left),
            PushPullDirection::SS => square.can_shift_in_direction(Direction::Down) && square.can_shift_in_direction(Direction::Up),
            PushPullDirection::SW => square.can_shift_in_direction(Direction::Down) && square.can_shift_in_direction(Direction::Right),
            PushPullDirection::WN => square.can_shift_in_direction(Direction::Left) && square.can_shift_in_direction(Direction::Down),
            PushPullDirection::WS => square.can_shift_in_direction(Direction::Left) && square.can_shift_in_direction(Direction::Up),
            PushPullDirection::WW => square.can_shift_in_direction(Direction::Left) && square.can_shift_in_direction(Direction::Right),
        }
    }
}

impl fmt::Display for PushPullDirection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            PushPullDirection::NN => write!(f, "nn"),
            PushPullDirection::NE => write!(f, "ne"),
            PushPullDirection::NW => write!(f, "nw"),
            PushPullDirection::EN => write!(f, "en"),
            PushPullDirection::EE => write!(f, "ee"),
            PushPullDirection::ES => write!(f, "es"),
            PushPullDirection::SE => write!(f, "se"),
            PushPullDirection::SS => write!(f, "ss"),
            PushPullDirection::SW => write!(f, "sw"),
            PushPullDirection::WN => write!(f, "wn"),
            PushPullDirection::WS => write!(f, "ws"),
            PushPullDirection::WW => write!(f, "ww"),
        }
    }
}

impl FromStr for PushPullDirection {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "nn" => Ok(PushPullDirection::NN),
            "ne" => Ok(PushPullDirection::NE),
            "nw" => Ok(PushPullDirection::NW),
            "en" => Ok(PushPullDirection::EN),
            "ee" => Ok(PushPullDirection::EE),
            "es" => Ok(PushPullDirection::ES),
            "se" => Ok(PushPullDirection::SE),
            "ss" => Ok(PushPullDirection::SS),
            "sw" => Ok(PushPullDirection::SW),
            "wn" => Ok(PushPullDirection::WN),
            "ws" => Ok(PushPullDirection::WS),
            "ww" => Ok(PushPullDirection::WW),
            _ => Err(anyhow!("Invalid value for direction")),
        }
    }
}

impl PushPullDirection {
    pub fn invert(&self) -> Self {
        match self {
            PushPullDirection::NN => PushPullDirection::SS,
            PushPullDirection::NE => PushPullDirection::SW,
            PushPullDirection::NW => PushPullDirection::SE,
            PushPullDirection::EN => PushPullDirection::WS,
            PushPullDirection::EE => PushPullDirection::WW,
            PushPullDirection::ES => PushPullDirection::WN,
            PushPullDirection::SE => PushPullDirection::NW,
            PushPullDirection::SS => PushPullDirection::NN,
            PushPullDirection::SW => PushPullDirection::NE,
            PushPullDirection::WN => PushPullDirection::ES,
            PushPullDirection::WS => PushPullDirection::EN,
            PushPullDirection::WW => PushPullDirection::EE,
        }
    }

    pub fn vertical_symmetry(&self) -> Self {
        match self {
            PushPullDirection::NN => PushPullDirection::NN,
            PushPullDirection::NE => PushPullDirection::NW,
            PushPullDirection::NW => PushPullDirection::NE,
            PushPullDirection::EN => PushPullDirection::WN,
            PushPullDirection::EE => PushPullDirection::WW,
            PushPullDirection::ES => PushPullDirection::WS,
            PushPullDirection::SE => PushPullDirection::SW,
            PushPullDirection::SS => PushPullDirection::SS,
            PushPullDirection::SW => PushPullDirection::SE,
            PushPullDirection::WN => PushPullDirection::EN,
            PushPullDirection::WS => PushPullDirection::ES,
            PushPullDirection::WW => PushPullDirection::EE,
        }
    }
}

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

    fn assert_can_shift_in_direction(invalid_squares: &[&str], direction: PushPullDirection) {
        for idx in 0..64 {
            let square = Square::from_index(idx);
            if invalid_squares.contains(&square.to_string().as_str()) {
                assert!(!direction.is_valid_from_square(square));
            } else {
                assert!(direction.is_valid_from_square(square));
            }
        }
    }

    #[test]
    fn test_is_valid_from_square_nn() {
        let invalid_squares = ["a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8", "a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1"];
        let direction = PushPullDirection::NN;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_ne() {
        let invalid_squares = ["a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8", "a8", "a7", "a6", "a5", "a4", "a3", "a2", "a1"];
        let direction = PushPullDirection::NE;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_nw() {
        let invalid_squares = ["a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8", "h8", "h7", "h6", "h5", "h4", "h3", "h2", "h1"];
        let direction = PushPullDirection::NW;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_en() {
        let invalid_squares = ["h8", "h7", "h6", "h5", "h4", "h3", "h2", "h1", "a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1"];
        let direction = PushPullDirection::EN;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_ee() {
        let invalid_squares = ["h8", "h7", "h6", "h5", "h4", "h3", "h2", "h1", "a8", "a7", "a6", "a5", "a4", "a3", "a2", "a1"];
        let direction = PushPullDirection::EE;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_es() {
        let invalid_squares = ["h8", "h7", "h6", "h5", "h4", "h3", "h2", "h1", "a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8"];
        let direction = PushPullDirection::ES;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_wn() {
        let invalid_squares = ["a8", "a7", "a6", "a5", "a4", "a3", "a2", "a1", "a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1"];
        let direction = PushPullDirection::WN;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_ws() {
        let invalid_squares = ["a8", "a7", "a6", "a5", "a4", "a3", "a2", "a1", "a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8"];
        let direction = PushPullDirection::WS;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_ww() {
        let invalid_squares = ["a8", "a7", "a6", "a5", "a4", "a3", "a2", "a1", "h8", "h7", "h6", "h5", "h4", "h3", "h2", "h1"];
        let direction = PushPullDirection::WW;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_se() {
        let invalid_squares = ["a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1", "a8", "a7", "a6", "a5", "a4", "a3", "a2", "a1"];
        let direction = PushPullDirection::SE;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_ss() {
        let invalid_squares = ["a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1", "a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8"];
        let direction = PushPullDirection::SS;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }

    #[test]
    fn test_can_shift_in_direction_sw() {
        let invalid_squares = ["a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1", "h8", "h7", "h6", "h5", "h4", "h3", "h2", "h1"];
        let direction = PushPullDirection::SW;

        assert_can_shift_in_direction(&invalid_squares, direction);
    }
}