arimaa-engine 0.1.0

use super::map_bit_board_to_squares;
use super::{Direction, Path, Piece, PushPullDirection, Square};

#[rustfmt::skip]
pub const TRAP_MASK: u64 = 0b__00000000__00000000__00100100__00000000__00000000__00100100__00000000__00000000;
#[rustfmt::skip]
const P1_PLACEMENT_MASK: u64 = 0b__11111111__11111111__00000000__00000000__00000000__00000000__00000000__00000000;
#[rustfmt::skip]
const P2_PLACEMENT_MASK: u64 = 0b__00000000__00000000__00000000__00000000__00000000__00000000__11111111__11111111;

#[rustfmt::skip]
pub const LEFT_COLUMN_MASK: u64 = 0b__00000001__00000001__00000001__00000001__00000001__00000001__00000001__00000001;
#[rustfmt::skip]
pub const RIGHT_COLUMN_MASK: u64 = 0b__10000000__10000000__10000000__10000000__10000000__10000000__10000000__10000000;
#[rustfmt::skip]
pub const TOP_ROW_MASK: u64 = 0b__00000000__00000000__00000000__00000000__00000000__00000000__00000000__11111111;
#[rustfmt::skip]
pub const BOTTOM_ROW_MASK: u64 = 0b__11111111__00000000__00000000__00000000__00000000__00000000__00000000__00000000;

#[derive(Clone, Debug)]

pub struct PieceBoard {
    pub p1_pieces: u64,
    pub all_pieces: u64,
    pub elephants: u64,
    pub camels: u64,
    pub horses: u64,
    pub dogs: u64,
    pub cats: u64,
    pub rabbits: u64,
}

impl PieceBoard {
    pub fn initial() -> Self {
        PieceBoard {
            p1_pieces: 0,
            all_pieces: 0,
            elephants: 0,
            camels: 0,
            horses: 0,
            dogs: 0,
            cats: 0,
            rabbits: 0,
        }
    }

    pub fn new(p1_pieces: u64, elephants: u64, camels: u64, horses: u64, dogs: u64, cats: u64, rabbits: u64) -> Self {
        Self {
            p1_pieces,
            elephants,
            camels,
            horses,
            dogs,
            cats,
            rabbits,
            all_pieces: rabbits | cats | dogs | horses | camels | elephants,
        }
    }

    pub fn get_bits_for_piece(&self, piece: Piece, p1_pieces: bool) -> u64 {
        let player_piece_mask = if p1_pieces { self.p1_pieces } else { !self.p1_pieces & self.all_pieces };
        self.get_bits_by_piece_type(piece) & player_piece_mask
    }

    pub fn get_player_piece_mask(&self, p1_pieces: bool) -> u64 {
        if p1_pieces { self.p1_pieces } else { !self.p1_pieces & self.all_pieces }
    }

    pub fn get_bits_by_piece_type(&self, piece: Piece) -> u64 {
        match piece {
            Piece::Elephant => self.elephants,
            Piece::Camel => self.camels,
            Piece::Horse => self.horses,
            Piece::Dog => self.dogs,
            Piece::Cat => self.cats,
            Piece::Rabbit => self.rabbits,
        }
    }

    pub fn bits_by_piece_type_mut(&mut self, piece: Piece) -> &mut u64 {
        match piece {
            Piece::Elephant => &mut self.elephants,
            Piece::Camel => &mut self.camels,
            Piece::Horse => &mut self.horses,
            Piece::Dog => &mut self.dogs,
            Piece::Cat => &mut self.cats,
            Piece::Rabbit => &mut self.rabbits,
        }
    }

    pub fn get_trapped_piece_bits(&self) -> u64 {
        let animal_is_on_trap = animal_is_on_trap(self);

        if animal_is_on_trap {
            let unsupported_piece_bits = get_both_player_unsupported_piece_bits(self);
            unsupported_piece_bits & TRAP_MASK
        } else {
            0
        }
    }

    pub fn get_piece_type_at_square(&self, square: &Square) -> Option<Piece> {
        let square_bit = square.as_bit_board();
        if square_bit & self.all_pieces != 0 { Some(get_piece_type_at_bit(square_bit, self)) } else { None }
    }

    pub fn get_curr_player_frozen_pieces(&self, p1_turn_to_move: bool) -> u64 {
        let opp_piece_mask = self.get_opponent_piece_mask(p1_turn_to_move);
        let curr_player_piece_mask = !opp_piece_mask & self.all_pieces;
        let threatened_pieces = get_threatened_pieces(opp_piece_mask, curr_player_piece_mask, self);

        curr_player_piece_mask & threatened_pieces & !get_supported_pieces(curr_player_piece_mask)
    }

    pub fn get_curr_player_non_frozen_pieces(&self, p1_turn_to_move: bool) -> u64 {
        let opp_piece_mask = self.get_opponent_piece_mask(p1_turn_to_move);
        let curr_player_piece_mask = !opp_piece_mask & self.all_pieces;
        let threatened_pieces = get_threatened_pieces(opp_piece_mask, curr_player_piece_mask, self);

        curr_player_piece_mask & (!threatened_pieces | get_supported_pieces(curr_player_piece_mask))
    }

    pub fn get_curr_player_piece_mask(&self, p1_turn_to_move: bool) -> u64 {
        if p1_turn_to_move { self.p1_pieces } else { !self.p1_pieces & self.all_pieces }
    }

    pub fn get_opponent_piece_mask(&self, p1_turn_to_move: bool) -> u64 {
        if p1_turn_to_move { !self.p1_pieces & self.all_pieces } else { self.p1_pieces }
    }

    pub fn push_pull_piece(&mut self, square: Square, direction: PushPullDirection) -> bool {
        let mut piece_trapped = false;

        for (square, direction) in direction.get_steps(square) {
            piece_trapped |= self.move_piece(square, direction);
        }

        piece_trapped
    }

    pub fn move_piece_on_path(&mut self, square: Square, path: Path) -> bool {
        let mut piece_trapped = false;
        let mut square = square;

        for direction in path.iter() {
            piece_trapped |= self.move_piece(square, direction);
            square = square.shift_in_direction(direction);
        }

        piece_trapped
    }

    pub fn move_piece_update_board(&mut self, square: Square, direction: Direction) {
        let source_square_bit = square.as_bit_board();

        self.elephants = shift_piece_in_direction(self.elephants, source_square_bit, direction);
        self.camels = shift_piece_in_direction(self.camels, source_square_bit, direction);
        self.horses = shift_piece_in_direction(self.horses, source_square_bit, direction);
        self.dogs = shift_piece_in_direction(self.dogs, source_square_bit, direction);
        self.cats = shift_piece_in_direction(self.cats, source_square_bit, direction);
        self.rabbits = shift_piece_in_direction(self.rabbits, source_square_bit, direction);
        self.p1_pieces = shift_piece_in_direction(self.p1_pieces, source_square_bit, direction);
        self.all_pieces = shift_piece_in_direction(self.all_pieces, source_square_bit, direction);
    }

    pub fn move_piece(&mut self, square: Square, direction: Direction) -> bool {
        self.move_piece_update_board(square, direction);
        self.remove_trapped_pieces()
    }

    pub fn remove_trapped_pieces(&mut self) -> bool {
        let trapped_animal_bits = self.get_trapped_piece_bits();
        let piece_trapped = trapped_animal_bits != 0;

        if piece_trapped {
            let untrapped_animal_bits = !trapped_animal_bits;
            self.elephants &= untrapped_animal_bits;
            self.camels &= untrapped_animal_bits;
            self.horses &= untrapped_animal_bits;
            self.dogs &= untrapped_animal_bits;
            self.cats &= untrapped_animal_bits;
            self.rabbits &= untrapped_animal_bits;
            self.p1_pieces &= untrapped_animal_bits;
            self.all_pieces &= untrapped_animal_bits;
        }

        piece_trapped
    }

    pub fn get_vertical_symmetry(&self) -> Self {
        Self::new(
            get_bit_board_vertical_symmetry(self.p1_pieces),
            get_bit_board_vertical_symmetry(self.elephants),
            get_bit_board_vertical_symmetry(self.camels),
            get_bit_board_vertical_symmetry(self.horses),
            get_bit_board_vertical_symmetry(self.dogs),
            get_bit_board_vertical_symmetry(self.cats),
            get_bit_board_vertical_symmetry(self.rabbits),
        )
    }

    pub fn place(&mut self, piece: Piece, square: Square, p1_piece: bool) {
        let placement_bit = square.as_bit_board();

        assert!(placement_bit & self.empty_placements(p1_piece) != 0, "Piece was placed on a non-empty or non placement square");

        self.fill(piece, placement_bit, p1_piece);
    }

    pub fn fill(&mut self, piece: Piece, mask: u64, p1_piece: bool) {
        *self.bits_by_piece_type_mut(piece) |= mask;

        self.all_pieces |= mask;

        if p1_piece {
            self.p1_pieces |= mask;
        }
    }

    pub fn empty_placements(&self, p1_placements: bool) -> u64 {
        let placement_bits = if p1_placements { P1_PLACEMENT_MASK } else { P2_PLACEMENT_MASK };

        !self.all_pieces & placement_bits
    }
}

fn get_bit_board_vertical_symmetry(bit_board: u64) -> u64 {
    map_bit_board_to_squares(bit_board).iter().map(|s| s.vertical_symmetry()).fold(0, |r, s| r | s.as_bit_board())
}

impl PartialEq for PieceBoard {
    fn eq(&self, other: &Self) -> bool {
        self.p1_pieces == other.p1_pieces
            && self.elephants == other.elephants
            && self.camels == other.camels
            && self.horses == other.horses
            && self.dogs == other.dogs
            && self.cats == other.cats
            && self.rabbits == other.rabbits
    }
}

pub struct ExcludedSquaresWithDirections {
    pub none: u64,
    pub north: u64,
    pub east: u64,
    pub south: u64,
    pub west: u64,
}

pub fn get_supported_square_exclusions(curr_player_piece_mask: u64) -> ExcludedSquaresWithDirections {
    let supported_squares_n = shift_pieces_up!(curr_player_piece_mask);
    let supported_squares_e = shift_pieces_right!(curr_player_piece_mask);
    let supported_squares_s = shift_pieces_down!(curr_player_piece_mask);
    let supported_squares_w = shift_pieces_left!(curr_player_piece_mask);
    let north = supported_squares_e | supported_squares_s | supported_squares_w;
    let east = supported_squares_n | supported_squares_s | supported_squares_w;
    let south = supported_squares_n | supported_squares_e | supported_squares_w;
    let west = supported_squares_n | supported_squares_e | supported_squares_s;
    let none = north | south;

    ExcludedSquaresWithDirections { north, east, south, west, none }
}

pub fn get_excluded_direction(i: usize, directions: &[Direction], exclusions: &ExcludedSquaresWithDirections) -> u64 {
    // Exclude the supported_squares for the cases where a piece would be able to support itself after the first move.
    // We check for i == 1 since a piece is only in a position to support itself after the first step.
    // We can exclude support depending on the direction since it is not possible to come from a square that has a piece which is supporting in the matching direction.
    if i == 1 { get_exclusion_for_direction(directions[0], exclusions) } else { exclusions.none }
}

pub fn get_exclusion_for_direction(direction: Direction, exclusions: &ExcludedSquaresWithDirections) -> u64 {
    match direction {
        Direction::Up => exclusions.north,
        Direction::Right => exclusions.east,
        Direction::Down => exclusions.south,
        Direction::Left => exclusions.west,
    }
}

pub fn get_threatened_pieces(predator_piece_mask: u64, prey_piece_mask: u64, piece_board: &PieceBoard) -> u64 {
    let predator_elephant_influence = get_influenced_squares(piece_board.elephants & predator_piece_mask);
    let predator_camel_influence = get_influenced_squares(piece_board.camels & predator_piece_mask);
    let predator_horse_influence = get_influenced_squares(piece_board.horses & predator_piece_mask);
    let predator_dog_influence = get_influenced_squares(piece_board.dogs & predator_piece_mask);
    let predator_cat_influence = get_influenced_squares(piece_board.cats & predator_piece_mask);

    let camel_threats = predator_elephant_influence;
    let horse_threats = camel_threats | predator_camel_influence;
    let dog_threats = horse_threats | predator_horse_influence;
    let cat_threats = dog_threats | predator_dog_influence;
    let rabbit_threats = cat_threats | predator_cat_influence;

    let threatened_pieces =
        (piece_board.camels & camel_threats) | (piece_board.horses & horse_threats) | (piece_board.dogs & dog_threats) | (piece_board.cats & cat_threats) | (piece_board.rabbits & rabbit_threats);

    threatened_pieces & prey_piece_mask
}

pub fn get_threatened_pieces_for_dir(predator_piece_mask: u64, prey_piece_mask: u64, direction: Direction, piece_board: &PieceBoard) -> u64 {
    let predator_elephant_influence = shift_pieces_in_direction(piece_board.elephants & predator_piece_mask, direction);
    let predator_camel_influence = shift_pieces_in_direction(piece_board.camels & predator_piece_mask, direction);
    let predator_horse_influence = shift_pieces_in_direction(piece_board.horses & predator_piece_mask, direction);
    let predator_dog_influence = shift_pieces_in_direction(piece_board.dogs & predator_piece_mask, direction);
    let predator_cat_influence = shift_pieces_in_direction(piece_board.cats & predator_piece_mask, direction);

    let camel_threats = predator_elephant_influence;
    let horse_threats = camel_threats | predator_camel_influence;
    let dog_threats = horse_threats | predator_horse_influence;
    let cat_threats = dog_threats | predator_dog_influence;
    let rabbit_threats = cat_threats | predator_cat_influence;

    let threatened_pieces =
        (piece_board.camels & camel_threats) | (piece_board.horses & horse_threats) | (piece_board.dogs & dog_threats) | (piece_board.cats & cat_threats) | (piece_board.rabbits & rabbit_threats);

    threatened_pieces & prey_piece_mask
}

pub fn get_threats_for_piece_type(predator_piece_mask: u64, piece_type: Piece, piece_board: &PieceBoard) -> u64 {
    match piece_type {
        Piece::Elephant => 0,
        Piece::Camel => get_influenced_squares(piece_board.elephants & predator_piece_mask),
        Piece::Horse => get_influenced_squares((piece_board.elephants | piece_board.camels) & predator_piece_mask),
        Piece::Dog => get_influenced_squares((piece_board.elephants | piece_board.camels | piece_board.horses) & predator_piece_mask),
        Piece::Cat => get_influenced_squares((piece_board.elephants | piece_board.camels | piece_board.horses | piece_board.dogs) & predator_piece_mask),
        Piece::Rabbit => get_influenced_squares((piece_board.elephants | piece_board.camels | piece_board.horses | piece_board.dogs | piece_board.cats) & predator_piece_mask),
    }
}

pub fn get_influenced_squares(piece_board: u64) -> u64 {
    shift_pieces_up!(piece_board) | shift_pieces_right!(piece_board) | shift_pieces_down!(piece_board) | shift_pieces_left!(piece_board)
}

pub fn can_move_in_direction(direction: Direction, empty_squares: u64) -> u64 {
    shift_pieces_in_opp_direction(direction, empty_squares)
}

fn shift_pieces_in_opp_direction(direction: Direction, bits: u64) -> u64 {
    match direction {
        Direction::Up => shift_pieces_down!(bits),
        Direction::Right => shift_pieces_left!(bits),
        Direction::Down => shift_pieces_up!(bits),
        Direction::Left => shift_pieces_right!(bits),
    }
}

pub fn shift_pieces_in_direction(bits: u64, direction: Direction) -> u64 {
    match direction {
        Direction::Up => shift_pieces_up!(bits),
        Direction::Right => shift_pieces_right!(bits),
        Direction::Down => shift_pieces_down!(bits),
        Direction::Left => shift_pieces_left!(bits),
    }
}

fn get_piece_type_at_bit(square_bit: u64, piece_board: &PieceBoard) -> Piece {
    if piece_board.rabbits & square_bit != 0 {
        Piece::Rabbit
    } else if piece_board.elephants & square_bit != 0 {
        Piece::Elephant
    } else if piece_board.camels & square_bit != 0 {
        Piece::Camel
    } else if piece_board.horses & square_bit != 0 {
        Piece::Horse
    } else if piece_board.dogs & square_bit != 0 {
        Piece::Dog
    } else {
        Piece::Cat
    }
}

fn animal_is_on_trap(piece_board: &PieceBoard) -> bool {
    (piece_board.all_pieces & TRAP_MASK) != 0
}

fn get_both_player_unsupported_piece_bits(piece_board: &PieceBoard) -> u64 {
    piece_board.all_pieces & !get_both_player_supported_pieces(piece_board)
}

fn get_both_player_supported_pieces(piece_board: &PieceBoard) -> u64 {
    let p1_pieces = piece_board.p1_pieces;
    let p2_pieces = piece_board.all_pieces & !p1_pieces;

    get_supported_pieces(p1_pieces) | get_supported_pieces(p2_pieces)
}

pub fn get_supported_pieces(piece_bits: u64) -> u64 {
    let up_supported_pieces = piece_bits & shift_pieces_up!(piece_bits);
    let right_supported_pieces = piece_bits & shift_pieces_right!(piece_bits);
    let down_supported_pieces = piece_bits & shift_pieces_down!(piece_bits);
    let left_supported_pieces = piece_bits & shift_pieces_left!(piece_bits);

    up_supported_pieces | right_supported_pieces | down_supported_pieces | left_supported_pieces
}

fn shift_piece_in_direction(piece_board: u64, source_square_bit: u64, direction: Direction) -> u64 {
    shift_in_direction(piece_board & source_square_bit, direction) | piece_board & !source_square_bit
}

fn shift_in_direction(bits: u64, direction: Direction) -> u64 {
    match direction {
        Direction::Up => shift_up!(bits),
        Direction::Right => shift_right!(bits),
        Direction::Down => shift_down!(bits),
        Direction::Left => shift_left!(bits),
    }
}