arimaa-engine 0.1.0

use super::bits::for_each_bit_idx;
use super::constants::MAX_NUM_STEPS;
use super::path_permutation::{PathPermutations, get_path_permutations, get_path_permutations_with_exact_steps};
use super::piece_board::ExcludedSquaresWithDirections;
use super::piece_board::get_exclusion_for_direction;
use super::piece_board::{BOTTOM_ROW_MASK, PieceBoard, TOP_ROW_MASK, TRAP_MASK};
use super::piece_board::{can_move_in_direction, get_excluded_direction, get_supported_square_exclusions, get_threatened_pieces_for_dir, get_threats_for_piece_type, shift_pieces_in_direction};
use super::repetition_counts::RepetitionCounts;
use super::value::Value;
use super::zobrist::Zobrist;
use super::{Action, Direction, Path, Piece, PushPullDirection, Square};

use itertools::Itertools;
use std::cell::OnceCell;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tinyvec::ArrayVec;

const P1_OBJECTIVE_MASK: u64 = TOP_ROW_MASK;
const P2_OBJECTIVE_MASK: u64 = BOTTOM_ROW_MASK;

#[derive(Clone, Debug)]
pub enum GameState {
    PlaceState(PlaceState),
    PlayState(PlayState),
}

#[derive(Clone, Debug)]
pub struct PlaceState {
    p1_turn_to_move: bool,
    piece_board: Box<PieceBoard>,
    hash: Zobrist,
}

#[derive(Clone, Debug)]
pub struct PlayState {
    turn_state: TurnState,
    step_state: StepState,
}

#[derive(Clone, Debug)]
pub struct TurnState {
    p1_turn_to_move: bool,
    move_number: u16,
    piece_board: PieceBoard,
    hash: Zobrist,
    hash_history: RepetitionCounts,
}

#[derive(Clone, Debug)]
pub struct StepState {
    steps: ArrayVec<[(Square, Direction); 3]>,
    actions: ArrayVec<[Action; 3]>,
    piece_board: PieceBoard,
    banned_piece_mask: Arc<AtomicU64>,
}

impl StepState {
    fn new(piece_board: PieceBoard) -> Self {
        Self {
            steps: ArrayVec::new(),
            actions: ArrayVec::new(),
            piece_board,
            banned_piece_mask: Arc::new(AtomicU64::new(u64::MAX)),
        }
    }
}

impl Hash for GameState {
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write_u64(self.hash().board_state_hash());
    }
}

impl PartialEq for GameState {
    fn eq(&self, other: &GameState) -> bool {
        self.hash().board_state_hash() == other.hash().board_state_hash()
    }
}

impl Eq for GameState {}

impl GameState {
    pub fn new(p1_turn_to_move: bool, move_number: usize, piece_board: PieceBoard) -> Self {
        let hash = Zobrist::from_piece_board(&piece_board, p1_turn_to_move, Some(0));

        let turn_state = TurnState {
            p1_turn_to_move,
            move_number: move_number as u16,
            piece_board: piece_board.clone(),
            hash,
            hash_history: RepetitionCounts::new().append(hash),
        };

        let step_state = StepState::new(piece_board);

        Self::PlayState(PlayState { turn_state, step_state })
    }

    pub fn initial() -> Self {
        Self::PlaceState(PlaceState {
            p1_turn_to_move: true,
            piece_board: Box::new(PieceBoard::initial()),
            hash: Zobrist::initial(),
        })
    }

    #[must_use]
    pub fn take_action(&self, action: &Action) -> Self {
        match self {
            Self::PlaceState(place_state) => place_state.take_action(action),
            Self::PlayState(play_state) => GameState::PlayState(play_state.take_action(action)),
        }
    }

    pub fn is_terminal(&self) -> Option<Value> {
        match self {
            Self::PlaceState(_) => None,
            Self::PlayState(play_state) => play_state.is_terminal(),
        }
    }

    pub fn valid_actions(&self) -> Vec<Action> {
        self.valid_actions_(true)
    }

    pub fn valid_actions_no_exclusions(&self) -> Vec<Action> {
        self.valid_actions_(false)
    }

    pub fn get_piece_board(&self) -> &PieceBoard {
        match self {
            Self::PlaceState(place_state) => place_state.get_piece_board(),
            Self::PlayState(play_state) => play_state.get_piece_board(),
        }
    }

    pub fn is_p1_turn_to_move(&self) -> bool {
        match self {
            Self::PlaceState(place_state) => place_state.p1_turn_to_move,
            Self::PlayState(play_state) => play_state.turn_state.p1_turn_to_move,
        }
    }

    pub fn get_move_number(&self) -> usize {
        match self {
            Self::PlaceState(_) => 1,
            Self::PlayState(play_state) => play_state.turn_state.move_number,
        }
        .into()
    }

    pub fn valid_actions_(&self, check_repetitions: bool) -> Vec<Action> {
        match self {
            Self::PlaceState(place_state) => place_state.valid_actions(),
            Self::PlayState(play_state) => play_state.valid_actions(check_repetitions),
        }
    }

    pub fn hash(&self) -> Zobrist {
        match self {
            Self::PlaceState(place_state) => place_state.hash,
            Self::PlayState(play_state) => play_state.hash(),
        }
    }

    /// Gets the current step number during the play phase. Note that the step number is 0 based. So values are always in the range 0..=3
    pub fn step(&self) -> usize {
        match self {
            Self::PlaceState(_) => 0,
            Self::PlayState(play_state) => play_state.get_step(),
        }
    }

    pub fn is_play_phase(&self) -> bool {
        matches!(self, Self::PlayState(_))
    }

    pub fn get_banned_piece_mask(&self) -> u64 {
        match self {
            Self::PlaceState(_) => 0,
            Self::PlayState(play_state) => play_state.get_banned_piece_mask(),
        }
    }

    pub fn get_vertical_symmetry(&self) -> Self {
        match self {
            Self::PlaceState(place_state) => Self::PlaceState(place_state.get_vertical_symmetry()),
            Self::PlayState(play_state) => Self::PlayState(play_state.get_vertical_symmetry()),
        }
    }

    pub fn get_trapped_animal_for_move(&self, square: Square, direction: Direction) -> Option<(Square, Piece, bool)> {
        match self {
            Self::PlaceState(_) => panic!("Get trapped animal is not valid in the place state"),
            Self::PlayState(play_state) => play_state.get_trapped_animal_for_move(square, direction),
        }
    }

    pub fn piece_to_place(&self) -> Piece {
        match self {
            Self::PlaceState(place_state) => place_state.piece_to_place(),
            Self::PlayState(_) => panic!("Piece to place is not valid in the play state"),
        }
    }
}

impl PlayState {
    fn initial(hash: Zobrist, piece_board: PieceBoard) -> Self {
        Self {
            turn_state: TurnState {
                piece_board: piece_board.clone(),
                p1_turn_to_move: true,
                move_number: 2,
                hash,
                hash_history: RepetitionCounts::new().append(hash),
            },
            step_state: StepState::new(piece_board),
        }
    }

    pub fn get_step(&self) -> usize {
        self.step_state.steps.len()
    }

    pub fn get_trapped_animal_for_move(&self, square: Square, direction: Direction) -> Option<(Square, Piece, bool)> {
        let mut piece_board = self.get_piece_board().clone();

        piece_board.move_piece_update_board(square, direction);

        let trapped_animal_bits = piece_board.get_trapped_piece_bits();
        if trapped_animal_bits != 0 {
            let square = Square::from_bit_board(trapped_animal_bits);
            let piece = piece_board.get_piece_type_at_square(&square).unwrap();
            let is_p1_piece = piece_board.get_bits_for_piece(piece, true) & square.as_bit_board() != 0;

            return Some((square, piece, is_p1_piece));
        }

        None
    }

    pub fn get_vertical_symmetry(&self) -> Self {
        Self {
            turn_state: self.turn_state.get_vertical_symmetry(),
            step_state: self.step_state.get_vertical_symmetry(),
        }
    }

    fn take_action(&self, action: &Action) -> Self {
        match action {
            Action::Pass => self.pass(),
            Action::Place(_) => panic!("Place action is not valid in the play state"),
            Action::Move(square, path) => self.move_piece(*square, *path),
            Action::PushPull(square, direction) => self.push_pull_piece(*square, *direction),
        }
    }

    fn is_terminal(&self) -> Option<Value> {
        let piece_board = &self.get_piece_board();

        // The order of checking for win/lose conditions is as follows assuming player A just made the move and player B now needs to move:
        if self.get_step() > 0 {
            return self.has_move(piece_board);
        }

        // Check if the position has occurred 3 times; the player who just moved loses.
        if self.turn_state.hash_history.num_repetitions(self.get_initial_hash_of_move()) >= 3 {
            return Some(if self.turn_state.p1_turn_to_move { Value([1.0, 0.0]) } else { Value([0.0, 1.0]) });
        }

        // Check if a rabbit of player A reached goal. If so player A wins.
        // Check if a rabbit of player B reached goal. If so player B wins.
        self.rabbit_at_goal(piece_board)
            // Check if player B lost all rabbits. If so player A wins.
            // Check if player A lost all rabbits. If so player B wins.
            .or_else(|| self.lost_all_rabbits(piece_board))
            // Check if player B has no possible move (all pieces are frozen or have no place to move). If so player A wins.
            // Check if the only moves player B has are 3rd time repetitions. If so player A wins.
            .or_else(|| self.has_move(piece_board))
    }

    fn valid_actions(&self, check_repetitions: bool) -> Vec<Action> {
        let piece_board = &self.get_piece_board();
        let hash = self.hash();
        let current_step = self.get_step();
        let p1_turn_to_move = self.turn_state.p1_turn_to_move;
        let banned_piece_mask = self.get_banned_piece_mask();

        let mut valid_actions = Vec::with_capacity(100);
        Self::each_with_push_pull_piece_actions(piece_board, p1_turn_to_move, current_step, |a| {
            if !self.is_action_passing(&a, piece_board, hash, check_repetitions) {
                valid_actions.push(a);
            }
            false
        });
        Self::each_valid_curr_player_piece_moves(piece_board, banned_piece_mask, p1_turn_to_move, current_step, |a| {
            if !self.is_action_passing(&a, piece_board, hash, check_repetitions) {
                valid_actions.push(a);
            }
            false
        });

        if !self.is_pass_passing_like_action(hash, check_repetitions) {
            valid_actions.push(Action::Pass);
        }

        valid_actions
    }

    fn has_move(&self, piece_board: &PieceBoard) -> Option<Value> {
        if self.can_move(piece_board) {
            None
        } else if self.turn_state.p1_turn_to_move {
            Some(Value([0.0, 1.0]))
        } else {
            Some(Value([1.0, 0.0]))
        }
    }

    fn can_move(&self, piece_board: &PieceBoard) -> bool {
        let piece_trapped = OnceCell::new();
        let piece_trapped = || piece_trapped.get_or_init(|| self.piece_trapped_this_turn());
        let banned_piece_mask = OnceCell::new();
        let banned_piece_mask = || banned_piece_mask.get_or_init(|| self.get_banned_piece_mask());
        let check_repetitions = true;
        let hash = self.hash();
        let current_step = self.get_step();
        let p1_turn_to_move = self.turn_state.p1_turn_to_move;
        let is_action_not_passing = |action| *piece_trapped() || !self.is_action_passing(&action, piece_board, hash, check_repetitions);
        let can_pass = || self.can_pass(hash, check_repetitions);
        let can_move_piece = || Self::each_valid_curr_player_piece_moves(piece_board, *banned_piece_mask(), p1_turn_to_move, current_step, is_action_not_passing);
        let can_push_pull_piece = || Self::each_with_push_pull_piece_actions(piece_board, p1_turn_to_move, current_step, is_action_not_passing);

        can_pass() || can_move_piece() || can_push_pull_piece()
    }

    fn can_pass(&self, hash: Zobrist, check_repetitions: bool) -> bool {
        self.get_step() > 0 && !self.is_pass_passing_like_action(hash, check_repetitions)
    }

    fn is_action_valid(piece_board: &PieceBoard, p1_turn_to_move: bool, action: &Action) -> bool {
        let is_matching_action = |a: Action| a == *action;

        match action {
            Action::Move(square, path) => {
                let piece_type = piece_board.get_piece_type_at_square(square);

                if piece_type.is_none() {
                    return false;
                }

                let curr_player_piece_mask = piece_board.get_curr_player_piece_mask(p1_turn_to_move);
                let opp_piece_mask = piece_board.get_opponent_piece_mask(p1_turn_to_move);
                let supported_square_exclusions = &get_supported_square_exclusions(curr_player_piece_mask);
                let path_permuations = get_path_permutations_with_exact_steps(path.len());
                let dangling_pieces_mask = OnceCell::new();
                let dangling_pieces_mask = || *dangling_pieces_mask.get_or_init(|| Self::get_dangling_pieces(curr_player_piece_mask));
                let banned_piece_mask = !square.as_bit_board();

                Self::each_valid_move_for_piece(
                    piece_board,
                    banned_piece_mask,
                    piece_type.unwrap(),
                    curr_player_piece_mask,
                    dangling_pieces_mask,
                    opp_piece_mask,
                    path_permuations,
                    supported_square_exclusions,
                    p1_turn_to_move,
                    is_matching_action,
                )
            }
            Action::PushPull(_, _) => Self::each_with_push_pull_piece_actions(piece_board, p1_turn_to_move, 2, is_matching_action),
            _ => panic!("Invalid sequential action type {}", action),
        }
    }

    fn get_piece_board(&self) -> &PieceBoard {
        &self.step_state.piece_board
    }

    fn get_initial_hash_of_move(&self) -> Zobrist {
        self.turn_state.hash
    }

    fn rabbit_at_goal(&self, piece_board: &PieceBoard) -> Option<Value> {
        let p1_objective_met = piece_board.p1_pieces & piece_board.rabbits & P1_OBJECTIVE_MASK != 0;
        let p2_objective_met = !piece_board.p1_pieces & piece_board.rabbits & P2_OBJECTIVE_MASK != 0;

        if p1_objective_met || p2_objective_met {
            // Objective is opposite of the player to move since we are checking if there is a winner after the turn is complete.
            // Logic should include the condition of if both players have a rabbit at the goal. In that case the player who was last to move wins.
            let last_to_move_is_p1 = !self.turn_state.p1_turn_to_move;
            let last_to_move_objective_met = if last_to_move_is_p1 { p1_objective_met } else { p2_objective_met };
            let p1_won = !(last_to_move_is_p1 ^ last_to_move_objective_met);
            Some(if p1_won { Value([1.0, 0.0]) } else { Value([0.0, 1.0]) })
        } else {
            None
        }
    }

    fn lost_all_rabbits(&self, piece_board: &PieceBoard) -> Option<Value> {
        let p1_lost_rabbits = piece_board.p1_pieces & piece_board.rabbits == 0;
        let p2_lost_rabbits = !piece_board.p1_pieces & piece_board.rabbits == 0;

        // Check if player B lost all rabbits. If so player A wins.
        // Check if player A lost all rabbits. If so player B wins.

        if p1_lost_rabbits || p2_lost_rabbits {
            // Objective is opposite of the player to move since we are checking if there is a winner after the turn is complete.
            // Logic should include the condition of if both players lost their rabbits. In that case the player who was last to move wins.
            let last_to_move_is_p1 = !self.turn_state.p1_turn_to_move;
            let last_to_move_objective_met = if last_to_move_is_p1 { p2_lost_rabbits } else { p1_lost_rabbits };
            let p1_won = !(last_to_move_is_p1 ^ last_to_move_objective_met);
            Some(if p1_won { Value([1.0, 0.0]) } else { Value([0.0, 1.0]) })
        } else {
            None
        }
    }

    fn is_action_passing(&self, action: &Action, piece_board: &PieceBoard, hash: Zobrist, check_repetitions: bool) -> bool {
        match action {
            Action::Move(square, path) => self.is_move_passing_like_action(*square, *path, piece_board, hash, check_repetitions),
            Action::PushPull(square, direction) => self.is_push_pull_passing_like_action(*square, *direction, piece_board, hash, check_repetitions),
            Action::Pass => self.is_pass_passing_like_action(hash, check_repetitions),
            Action::Place(_) => false,
        }
    }

    #[allow(clippy::let_and_return)]
    fn is_move_passing_like_action(&self, square: Square, path: Path, piece_board: &PieceBoard, hash: Zobrist, check_repetitions: bool) -> bool {
        let updated_piece_board_hash = || {
            let mut new_piece_board = piece_board.clone();
            new_piece_board.move_piece_on_path(square, path);

            let updated_piece_board_hash = hash.update_piece_board(piece_board, &new_piece_board);

            updated_piece_board_hash
        };

        self.is_passing_like_action(path.len(), updated_piece_board_hash, check_repetitions)
    }

    #[allow(clippy::let_and_return)]
    fn is_push_pull_passing_like_action(&self, square: Square, direction: PushPullDirection, piece_board: &PieceBoard, hash: Zobrist, check_repetitions: bool) -> bool {
        let update_piece_board_hash = || {
            let mut new_piece_board = piece_board.clone();
            new_piece_board.push_pull_piece(square, direction);

            let updated_piece_board_hash = hash.update_piece_board(piece_board, &new_piece_board);

            updated_piece_board_hash
        };

        self.is_passing_like_action(2, update_piece_board_hash, check_repetitions)
    }

    fn is_pass_passing_like_action(&self, hash: Zobrist, check_repetitions: bool) -> bool {
        self.is_passing_like_action(4, || hash, check_repetitions)
    }

    fn is_passing_like_action<F: FnMut() -> Zobrist>(&self, steps_used: usize, mut update_piece_board_hash: F, check_repetitions: bool) -> bool {
        let curr_step_num = self.get_step();

        if curr_step_num + steps_used < MAX_NUM_STEPS {
            return false;
        }

        let initial_hash_of_move = self.get_initial_hash_of_move();
        let updated_piece_board_hash = update_piece_board_hash();

        let curr_is_p1_turn_to_move = self.turn_state.p1_turn_to_move;
        let new_hash_no_player_switch = updated_piece_board_hash.update_step(curr_step_num, 0);
        let new_hash_switch_players = new_hash_no_player_switch.update_player_to_move(curr_is_p1_turn_to_move, !curr_is_p1_turn_to_move);

        new_hash_no_player_switch == initial_hash_of_move || (check_repetitions && self.turn_state.hash_history.num_repetitions(new_hash_switch_players) >= 2)
    }

    fn each_valid_curr_player_piece_moves<F: FnMut(Action) -> bool>(piece_board: &PieceBoard, banned_piece_mask: u64, p1_turn_to_move: bool, current_step: usize, mut cb: F) -> bool {
        let curr_player_piece_mask = piece_board.get_curr_player_piece_mask(p1_turn_to_move);
        let opp_piece_mask = piece_board.get_opponent_piece_mask(p1_turn_to_move);
        let num_steps_remaining = MAX_NUM_STEPS - current_step;
        let supported_square_exclusions = get_supported_square_exclusions(curr_player_piece_mask);
        let path_permuations = get_path_permutations(num_steps_remaining);
        let dangling_pieces_mask = OnceCell::new();
        let dangling_pieces_mask = || *dangling_pieces_mask.get_or_init(|| Self::get_dangling_pieces(curr_player_piece_mask));

        for piece_type in [Piece::Rabbit, Piece::Cat, Piece::Dog, Piece::Horse, Piece::Camel, Piece::Elephant] {
            let complete = Self::each_valid_move_for_piece(
                piece_board,
                banned_piece_mask,
                piece_type,
                curr_player_piece_mask,
                dangling_pieces_mask,
                opp_piece_mask,
                path_permuations,
                &supported_square_exclusions,
                p1_turn_to_move,
                &mut cb,
            );

            if complete {
                return true;
            }
        }

        false
    }

    #[allow(clippy::too_many_arguments)]
    fn each_valid_move_for_piece<F: FnMut(Action) -> bool, F2: Fn() -> u64>(
        piece_board: &PieceBoard,
        banned_piece_mask: u64,
        piece_type: Piece,
        curr_player_piece_mask: u64,
        dangling_pieces_mask: F2,
        opp_piece_mask: u64,
        path_permuations: &[PathPermutations],
        supported_square_exclusions: &ExcludedSquaresWithDirections,
        p1_turn_to_move: bool,
        mut cb: F,
    ) -> bool {
        let pieces_to_move_board = piece_board.get_bits_by_piece_type(piece_type) & curr_player_piece_mask & !banned_piece_mask;

        // If there are no pieces of this type, move to the next piece type.
        if pieces_to_move_board == 0 {
            return false;
        }

        // Exclude frozen pieces to allow for short circuiting since a frozen piece is unable to be moved.
        let threatened_squares = get_threats_for_piece_type(opp_piece_mask, piece_type, piece_board);
        let pieces_to_move_board = pieces_to_move_board & (supported_square_exclusions.none | !threatened_squares);

        // If there are pieces of this type but they are all frozen, move to the next piece type.
        if pieces_to_move_board == 0 {
            return false;
        }

        let empty_squares = !piece_board.all_pieces;
        let invalid_rabbit_direction = if p1_turn_to_move { Direction::Down } else { Direction::Up };
        let has_potential_dangling_piece = curr_player_piece_mask & TRAP_MASK != 0;

        'moves: for paths in path_permuations.iter() {
            let mut successful_moves = 0u64;

            'paths: for (path, packed_path) in paths {
                // Do not include moves that have been proven when calculating other paths.
                let mut moving_pieces_board = pieces_to_move_board & !successful_moves;

                for (i, direction) in path.iter().enumerate() {
                    if piece_type == Piece::Rabbit && *direction == invalid_rabbit_direction {
                        continue 'paths;
                    }
                    let mut supported_squares = get_excluded_direction(i, path, supported_square_exclusions);

                    // Check if the current piece was supporting a dangling piece on the trap.
                    // When the current piece moves along the outside of the trap, the trapped piece will no longer provide support.
                    if has_potential_dangling_piece && i == 2 && path.len() >= 3 && path[0] != path[1] {
                        let dangling_pieces_mask: u64 = dangling_pieces_mask();
                        if dangling_pieces_mask != 0 {
                            let no_longer_supported_squares = Self::get_no_longer_supported_pieces(moving_pieces_board, dangling_pieces_mask, supported_square_exclusions, path);
                            supported_squares &= !no_longer_supported_squares
                        }
                    }

                    let supported_unendangered_squares = supported_squares | !(TRAP_MASK | threatened_squares);
                    let unoccupied_directions = can_move_in_direction(*direction, empty_squares);
                    let valid_curr_piece_moves = moving_pieces_board & supported_unendangered_squares & unoccupied_directions;

                    // Move the pieces that are valid to be moved.
                    moving_pieces_board = shift_pieces_in_direction(valid_curr_piece_moves, *direction);

                    // If there are no more movable pieces. Continue from the other paths loop since we are done attempting this path.
                    if moving_pieces_board == 0 {
                        continue 'paths;
                    }
                }

                let successfully_moved_pieces = Self::unwind_piece_path(moving_pieces_board, path);

                let complete = Self::for_each_square(successfully_moved_pieces, |s| cb(Action::Move(s, *packed_path)));

                if complete {
                    return true;
                }

                // Set any of the moves that were completed to the list of successful moves.
                successful_moves |= successfully_moved_pieces;

                // If all pieces have found a path for this move, continue from the moves loop since we are done attempting this move for the current piece type.
                if pieces_to_move_board == successful_moves {
                    continue 'moves;
                }
            }
        }

        false
    }

    // Iterate over the directions in reverse to move the pieces back to their starting locations.
    fn unwind_piece_path(piece_board: u64, path: &[Direction]) -> u64 {
        path.iter().rev().fold(piece_board, |piece_board, dir| shift_pieces_in_direction(piece_board, dir.invert()))
    }

    fn for_each_square<F: FnMut(Square) -> bool>(piece_board: u64, mut cb: F) -> bool {
        let mut complete = false;
        for_each_bit_idx(piece_board, |idx| {
            if !complete {
                let square = Square::from_index(idx as u8);
                complete |= cb(square);
            }
        });
        complete
    }

    fn each_with_push_pull_piece_actions<F: FnMut(Action) -> bool>(piece_board: &PieceBoard, p1_turn_to_move: bool, current_step: usize, mut cb: F) -> bool {
        if current_step <= 2 {
            let predator_piece_mask = piece_board.get_curr_player_non_frozen_pieces(p1_turn_to_move);
            let opp_piece_mask = piece_board.get_opponent_piece_mask(p1_turn_to_move);
            let empty_squares = !piece_board.all_pieces;

            for (occupying_dir, vacating_dirs) in [
                (
                    Direction::Up,
                    [
                        (Direction::Up, PushPullDirection::NN),
                        (Direction::Right, PushPullDirection::EN),
                        (Direction::Left, PushPullDirection::WN),
                    ],
                ),
                (
                    Direction::Right,
                    [
                        (Direction::Up, PushPullDirection::NE),
                        (Direction::Right, PushPullDirection::EE),
                        (Direction::Down, PushPullDirection::SE),
                    ],
                ),
                (
                    Direction::Down,
                    [
                        (Direction::Right, PushPullDirection::ES),
                        (Direction::Down, PushPullDirection::SS),
                        (Direction::Left, PushPullDirection::WS),
                    ],
                ),
                (
                    Direction::Left,
                    [
                        (Direction::Up, PushPullDirection::NW),
                        (Direction::Down, PushPullDirection::SW),
                        (Direction::Left, PushPullDirection::WW),
                    ],
                ),
            ]
            .iter()
            {
                let opp_threatened_pieces_push = get_threatened_pieces_for_dir(predator_piece_mask, opp_piece_mask, *occupying_dir, piece_board);
                let opp_threatened_pieces_pull = get_threatened_pieces_for_dir(predator_piece_mask, opp_piece_mask, occupying_dir.invert(), piece_board);
                let opp_threatened_pieces_pull = shift_pieces_in_direction(opp_threatened_pieces_pull, *occupying_dir);

                for (vacating_dir, move_direction) in vacating_dirs {
                    let unoccupied_directions = can_move_in_direction(*vacating_dir, empty_squares);

                    let valid_push_moves = unoccupied_directions & opp_threatened_pieces_push;
                    let valid_pull_moves = unoccupied_directions & opp_threatened_pieces_pull;

                    debug_assert_eq!(valid_push_moves & valid_pull_moves, 0);

                    let valid_push_pull_moves = valid_push_moves | valid_pull_moves;

                    if valid_push_pull_moves != 0 {
                        let complete = Self::for_each_square(valid_push_pull_moves, |s| cb(Action::PushPull(s, *move_direction)));

                        if complete {
                            return true;
                        }
                    }
                }
            }
        }

        false
    }

    fn pass(&self) -> PlayState {
        let hash = self.hash().pass(self.get_step());
        let new_hash_history = if self.piece_trapped_this_turn() {
            RepetitionCounts::new()
        } else {
            self.turn_state.hash_history.clone()
        };
        let new_hash_history = new_hash_history.append(hash);

        let piece_board = self.get_piece_board();
        let turn_state = TurnState {
            p1_turn_to_move: !self.turn_state.p1_turn_to_move,
            move_number: self.turn_state.move_number + if self.turn_state.p1_turn_to_move { 0 } else { 1 },
            hash,
            hash_history: new_hash_history,
            piece_board: piece_board.clone(),
        };

        let step_state = StepState::new(piece_board.clone());

        PlayState { turn_state, step_state }
    }

    fn move_piece(&self, square: Square, path: Path) -> Self {
        let mut piece_board = self.get_piece_board().clone();
        piece_board.move_piece_on_path(square, path);
        let steps = path
            .iter()
            .scan(square, |square, direction| {
                let original = *square;
                *square = square.shift_in_direction(direction);
                Some((original, direction))
            })
            .collect::<ArrayVec<[_; 4]>>();

        self.get_state_after_move(piece_board, &steps, Action::Move(square, path))
    }

    fn push_pull_piece(&self, square: Square, direction: PushPullDirection) -> Self {
        let mut piece_board = self.get_piece_board().clone();
        piece_board.push_pull_piece(square, direction);
        let steps = direction.get_steps(square);

        self.get_state_after_move(piece_board, &steps, Action::PushPull(square, direction))
    }

    fn get_state_after_move(&self, new_piece_board: PieceBoard, steps: &[(Square, Direction)], action: Action) -> Self {
        let num_steps_taken = steps.len();
        let curr_step = self.get_step();
        let num_steps_this_turn = curr_step + num_steps_taken;
        debug_assert!(num_steps_this_turn <= MAX_NUM_STEPS, "Number of steps should not be greater than MAX_NUM_STEPS");
        let is_last_step = num_steps_this_turn >= MAX_NUM_STEPS;
        let new_p1_turn_to_move = if is_last_step { !self.turn_state.p1_turn_to_move } else { self.turn_state.p1_turn_to_move };
        let new_step = if is_last_step { 0 } else { num_steps_this_turn };

        let new_hash = self
            .get_initial_hash_of_move()
            .update_piece_board(&self.turn_state.piece_board, &new_piece_board)
            .update_step(0, new_step)
            .update_player_to_move(self.turn_state.p1_turn_to_move, new_p1_turn_to_move);

        let new_move_number = self.turn_state.move_number + if is_last_step && new_p1_turn_to_move { 1 } else { 0 };

        let new_turn_state = if is_last_step {
            let new_hash_history = if self.piece_trapped_this_turn() {
                RepetitionCounts::new()
            } else {
                self.turn_state.hash_history.clone()
            };

            TurnState {
                p1_turn_to_move: new_p1_turn_to_move,
                move_number: new_move_number,
                piece_board: new_piece_board.clone(),
                hash: new_hash,
                hash_history: new_hash_history.append(new_hash),
            }
        } else {
            self.turn_state.clone()
        };

        let new_step_state = if is_last_step {
            StepState::new(new_piece_board)
        } else {
            let mut new_steps = self.step_state.steps;
            new_steps.extend_from_slice(steps);

            let mut new_actions = self.step_state.actions.clone();
            new_actions.push(action);

            StepState {
                steps: new_steps,
                actions: new_actions,
                piece_board: new_piece_board,
                banned_piece_mask: Arc::new(AtomicU64::new(u64::MAX)),
            }
        };

        PlayState {
            turn_state: new_turn_state,
            step_state: new_step_state,
        }
    }

    fn piece_trapped_this_turn(&self) -> bool {
        let mut piece_board = self.turn_state.piece_board.clone();
        for (square, direction) in self.step_state.steps.iter() {
            if piece_board.move_piece(*square, *direction) {
                return true;
            }
        }

        false
    }

    fn own_piece_trapped(&self, piece_board: &PieceBoard, action: &Action) -> bool {
        let mut piece_board = piece_board.clone();
        for (square, direction) in action.get_steps() {
            piece_board.move_piece_update_board(square, direction);
            let trapped_piece_bit = piece_board.get_trapped_piece_bits();

            if trapped_piece_bit != 0 {
                let curr_player_piece_mask = piece_board.get_curr_player_piece_mask(self.turn_state.p1_turn_to_move);

                if trapped_piece_bit & curr_player_piece_mask != 0 {
                    return true;
                }
            }

            piece_board.remove_trapped_pieces();
        }

        false
    }

    fn is_valid_transposition(initial_turn_piece_board: &PieceBoard, p1_turn_to_move: bool, actions: &[&Action], candidate_action: Action, expected_piece_board: &PieceBoard) -> bool {
        let mut rearranged_order_piece_board = initial_turn_piece_board.clone();
        let candidate_action_steps = candidate_action.get_steps();
        let (candidate_action_departing_square, candidate_action_departing_direction) = candidate_action_steps.first().unwrap();
        let mut has_interposed = false;

        for action in actions {
            let is_valid = Self::is_action_valid(&rearranged_order_piece_board, p1_turn_to_move, action);

            if !is_valid {
                return false;
            }

            let steps = action.get_steps();
            for (square, direction) in steps {
                rearranged_order_piece_board.move_piece(square, direction);
            }

            if !has_interposed && matches!(action, Action::Move(_, _)) {
                let landing_square = action.landing_square();
                if landing_square == *candidate_action_departing_square {
                    let is_inverse_direction = candidate_action_departing_direction.invert() == action.get_steps().last().unwrap().1;

                    if is_inverse_direction {
                        return false;
                    }

                    let is_valid = Self::is_action_valid(&rearranged_order_piece_board, p1_turn_to_move, &candidate_action);

                    if !is_valid {
                        return false;
                    }

                    for (square, direction) in candidate_action_steps {
                        rearranged_order_piece_board.move_piece(square, direction);
                    }

                    has_interposed = true;
                }
            }
        }

        rearranged_order_piece_board == *expected_piece_board
    }

    fn update_and_verify_piece_board_with_actions(piece_board: &mut PieceBoard, p1_turn_to_move: bool, actions: &[Action]) -> bool {
        for action in actions {
            let is_valid = Self::is_action_valid(piece_board, p1_turn_to_move, action);

            if !is_valid {
                return false;
            }

            for (square, direction) in action.get_steps() {
                piece_board.move_piece(square, direction);
            }
        }

        true
    }

    fn get_banned_piece_mask(&self) -> u64 {
        let mut banned_piece_mask = self.step_state.banned_piece_mask.load(Ordering::Relaxed);

        if banned_piece_mask == u64::MAX {
            banned_piece_mask = self.get_banned_piece_mask_inner();
            self.step_state.banned_piece_mask.store(banned_piece_mask, Ordering::SeqCst);
        }

        banned_piece_mask
    }

    fn get_banned_piece_mask_inner(&self) -> u64 {
        if self.get_step() == 0 {
            return 0;
        }

        let p1_turn_to_move = self.turn_state.p1_turn_to_move;
        let initial_turn_piece_board = &self.turn_state.piece_board;

        let unban_pieces = |banned_piece_mask: u64, actions: &[Action], piece_board: &PieceBoard| -> u64 {
            if banned_piece_mask == 0 {
                return banned_piece_mask;
            }

            let mut banned_piece_mask = banned_piece_mask;

            // Force current_step to 3 since we only need to test for actions of 1 step.
            let current_step = 3;

            Self::each_valid_curr_player_piece_moves(piece_board, !banned_piece_mask, p1_turn_to_move, current_step, |candidate_action| {
                debug_assert!(candidate_action.get_steps().len() == 1, "Candidate actions should always be of 1 step.");

                let candidate_action_steps = candidate_action.get_steps();
                let candidate_action_departing_square = candidate_action_steps.first().unwrap().0;

                if candidate_action_departing_square.as_bit_board() & banned_piece_mask == 0 {
                    return false;
                }

                let expected_piece_board = {
                    let mut expected_piece_board = initial_turn_piece_board.clone();
                    let mut actions = Vec::from(actions);
                    actions.push(candidate_action.clone());
                    let is_valid = Self::update_and_verify_piece_board_with_actions(&mut expected_piece_board, p1_turn_to_move, &actions);

                    if !is_valid {
                        return false;
                    }

                    expected_piece_board
                };

                let min_actions = 1.max(actions.len() - 1);
                let has_valid_transpostion = (min_actions..=actions.len())
                    .flat_map(|c| actions.iter().permutations(c))
                    .any(|actions_inner| Self::is_valid_transposition(initial_turn_piece_board, p1_turn_to_move, &actions_inner, candidate_action.clone(), &expected_piece_board));

                if !has_valid_transpostion {
                    banned_piece_mask &= !candidate_action_departing_square.as_bit_board();
                }

                banned_piece_mask == 0
            });

            banned_piece_mask
        };

        let mut piece_board = initial_turn_piece_board.clone();
        let actions = &self.step_state.actions;
        let mut banned_piece_mask = 0;

        for (i, action) in actions.iter().enumerate() {
            let latest_action_steps = action.get_steps();
            let own_piece_trapped = self.own_piece_trapped(&piece_board, action);

            for (square, direction) in latest_action_steps {
                piece_board.move_piece(square, direction);
            }

            banned_piece_mask = unban_pieces(banned_piece_mask, &actions[..=i], &piece_board);
            let is_move_action = matches!(action, Action::Move(_, _));

            // Don't bother banning if action results in trapping of own piece. This allows for a shuffle of a piece as well as moving onto the trap square of the piece that was just trapped.
            if !is_move_action {
                continue;
            }

            // If a piece is trapped, clear it from the list of banned pieces since it no longer exists.
            if own_piece_trapped {
                let curr_player_piece_mask = piece_board.get_curr_player_piece_mask(p1_turn_to_move);
                banned_piece_mask &= curr_player_piece_mask;
            }

            // Shortcut since banning a piece here always results in an unban since the piece can be moved back to its previous square prior to the drop.
            if own_piece_trapped && latest_action_steps.len() <= 2 {
                continue;
            }

            // We can use steps here since we can only get to this point if the action is of type move.
            let landing_square_bit = action.landing_square().as_bit_board();

            // Skip banning the piece if it was frozen.
            let piece_is_frozen_after_move = piece_board.get_curr_player_frozen_pieces(p1_turn_to_move) & landing_square_bit != 0;

            if piece_is_frozen_after_move {
                continue;
            }

            banned_piece_mask |= landing_square_bit;
        }

        banned_piece_mask
    }

    fn get_dangling_pieces(curr_player_piece_mask: u64) -> u64 {
        let up = shift_up!(curr_player_piece_mask);
        let right = shift_right!(curr_player_piece_mask);
        let down = shift_down!(curr_player_piece_mask);
        let left = shift_left!(curr_player_piece_mask);
        let mut set_once = 0;
        let mut set_mult = 0;
        for board in [up, right, down, left] {
            set_mult |= set_once & board;
            set_once ^= board;
        }

        let dangling_piece_mask = curr_player_piece_mask & TRAP_MASK & !set_mult;
        shift_up!(dangling_piece_mask) | shift_right!(dangling_piece_mask) | shift_down!(dangling_piece_mask) | shift_left!(dangling_piece_mask)
    }

    fn get_no_longer_supported_pieces(moving_pieces_board: u64, dangling_pieces: u64, supported_square_exclusions: &ExcludedSquaresWithDirections, path: &ArrayVec<[Direction; 4]>) -> u64 {
        let started_from_dangling_pos = shift_pieces_in_direction(shift_pieces_in_direction(dangling_pieces, path[0]), path[1]);
        let dropped_dangling_piece = moving_pieces_board & started_from_dangling_pos & dangling_pieces;
        dropped_dangling_piece & !get_exclusion_for_direction(path[0], supported_square_exclusions)
    }

    fn hash(&self) -> Zobrist {
        let step_num = self.get_step();
        let hash = self.get_initial_hash_of_move();

        if step_num == 0 {
            return hash;
        }

        let turn_piece_board = &self.turn_state.piece_board;
        let curr_piece_board = self.get_piece_board();
        hash.update_piece_board(turn_piece_board, curr_piece_board).update_step(0, step_num)
    }
}

impl PlaceState {
    fn take_action(&self, action: &Action) -> GameState {
        match action {
            Action::Place(square) => self.place(*square),
            _ => panic!("Provided action is not valid in the place state: {:?}", action),
        }
    }

    fn valid_actions(&self) -> Vec<Action> {
        let empty_placement_bits = self.get_piece_board().empty_placements(self.p1_turn_to_move);

        Square::from_multiple_bits(empty_placement_bits).into_iter().map(Action::Place).collect()
    }

    const PIECES_TO_PLACE: [(Piece, u32); 5] = [(Piece::Elephant, 1u32), (Piece::Camel, 1), (Piece::Horse, 2), (Piece::Dog, 2), (Piece::Cat, 2)];

    fn place(&self, square: Square) -> GameState {
        let mut piece_board = self.get_piece_board().clone();

        let p1_turn_to_move = self.p1_turn_to_move;
        let piece_to_place = self.piece_to_place();
        piece_board.place(piece_to_place, square, p1_turn_to_move);

        let last_piece_placed = piece_board.get_bits_for_piece(Piece::Cat, p1_turn_to_move).count_ones() == 2;

        let mut new_hash = self.hash;
        if last_piece_placed {
            Self::fill_remaining_rabbits(&mut piece_board, p1_turn_to_move);
            new_hash = new_hash.fill_rabbits(&piece_board, p1_turn_to_move);
        }

        let switch_players = last_piece_placed && p1_turn_to_move;
        let switch_phases = last_piece_placed && !p1_turn_to_move;
        let new_p1_turn_to_move = if switch_players {
            false
        } else if switch_phases {
            true
        } else {
            p1_turn_to_move
        };

        let new_hash = new_hash.place_piece(piece_to_place, square, p1_turn_to_move, switch_players, switch_phases);

        if switch_phases {
            GameState::PlayState(PlayState::initial(new_hash, piece_board))
        } else {
            GameState::PlaceState(PlaceState {
                p1_turn_to_move: new_p1_turn_to_move,
                piece_board: Box::new(piece_board),
                hash: new_hash,
            })
        }
    }

    pub fn piece_to_place(&self) -> Piece {
        let piece_board = self.get_piece_board();

        Self::PIECES_TO_PLACE
            .iter()
            .find(|(piece, count)| piece_board.get_bits_for_piece(*piece, self.p1_turn_to_move).count_ones() < *count)
            .expect("Expected available piece to place.")
            .0
    }

    pub fn fill_remaining_rabbits(piece_board: &mut PieceBoard, p1_turn_to_move: bool) {
        let empty_squares = piece_board.empty_placements(p1_turn_to_move);

        assert_eq!(empty_squares.count_ones(), 8, "Expected 8 squares for rabbits to be placed.");

        piece_board.fill(Piece::Rabbit, empty_squares, p1_turn_to_move);
    }

    pub fn get_piece_board(&self) -> &PieceBoard {
        &self.piece_board
    }

    pub fn get_vertical_symmetry(&self) -> Self {
        let piece_board = Box::new(self.piece_board.get_vertical_symmetry());
        let hash = Zobrist::from_piece_board(&piece_board, self.p1_turn_to_move, Some(0));

        Self {
            p1_turn_to_move: self.p1_turn_to_move,
            piece_board,
            hash,
        }
    }
}

impl TurnState {
    fn get_vertical_symmetry(&self) -> Self {
        let piece_board = self.piece_board.get_vertical_symmetry();
        let hash = Zobrist::from_piece_board(&piece_board, self.p1_turn_to_move, Some(0));

        Self {
            p1_turn_to_move: self.p1_turn_to_move,
            move_number: self.move_number,
            piece_board,
            hash,
            hash_history: RepetitionCounts::new().append(hash),
        }
    }
}

impl StepState {
    fn get_vertical_symmetry(&self) -> Self {
        Self {
            steps: self.steps.iter().map(|(s, d)| (s.vertical_symmetry(), d.vertical_symmetry())).collect(),
            actions: self.actions.iter().map(|a| a.vertical_symmetry()).collect(),
            piece_board: self.piece_board.get_vertical_symmetry(),
            banned_piece_mask: Arc::new(AtomicU64::new(u64::MAX)),
        }
    }
}