//
// jja: swiss army knife for chess file formats
// src/stockfish.rs: Stockfish compatible Zobrist hashing
//
// Copyright (c) 2023 Ali Polatel <alip@chesswob.org>
// Based in part upon Stockfish & Rustfish which is
//     Copyright (C) 2004-2023 The Stockfish developers (see AUTHORS file)
//
// SPDX-License-Identifier: GPL-3.0-or-later

/// Stockfish types
pub mod types {
    /// This type represents castling rights as implemented by Stockfish
    #[derive(Clone, Copy, PartialEq, Eq)]
    pub struct CastlingRight(pub u32);

    /// This type represents a position key as implemented by Stockfish
    #[derive(Clone, Copy, PartialEq, Eq, Hash)]
    pub struct Key(pub u64);

    impl std::ops::BitXor<Key> for Key {
        type Output = Self;
        fn bitxor(self, rhs: Self) -> Self {
            Key(self.0 ^ rhs.0)
        }
    }

    impl std::ops::BitXorAssign<Key> for Key {
        fn bitxor_assign(&mut self, rhs: Key) {
            *self = *self ^ rhs;
        }
    }

    impl std::fmt::Display for Key {
        fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
            write!(f, "{:X}", self.0)
        }
    }

    impl std::ops::BitAnd<CastlingRight> for CastlingRight {
        type Output = Self;
        fn bitand(self, rhs: Self) -> Self {
            CastlingRight(self.0 & rhs.0)
        }
    }

    impl std::ops::BitOr<CastlingRight> for CastlingRight {
        type Output = Self;
        fn bitor(self, rhs: Self) -> Self {
            CastlingRight(self.0 | rhs.0)
        }
    }

    impl std::ops::BitAndAssign<CastlingRight> for CastlingRight {
        fn bitand_assign(&mut self, rhs: Self) {
            *self = *self & rhs;
        }
    }

    impl std::ops::BitOrAssign<CastlingRight> for CastlingRight {
        fn bitor_assign(&mut self, rhs: Self) {
            *self = *self | rhs;
        }
    }

    impl std::ops::Not for CastlingRight {
        type Output = CastlingRight;
        fn not(self) -> Self {
            CastlingRight(!self.0)
        }
    }

    impl std::cmp::PartialEq<u32> for CastlingRight {
        fn eq(&self, rhs: &u32) -> bool {
            debug_assert!(*rhs == 0);
            self.0 == *rhs
        }
    }
}

/// Miscallenous utilities
pub mod misc {
    /// xorshift64star Pseudo-Random Number Generator
    /// This class is based on original code written and dedicated
    /// to the public domain by Sebastiano Vigna (2014).
    /// It has the following characteristics:
    ///
    ///  -  Outputs 64-bit numbers
    ///  -  Passes Dieharder and Smallcrush test batteries
    ///  -  Does not require warm-up, no zeroland to escape
    ///  -  Internal state is a single 64-bit integer
    ///  -  Period is 2^64 - 1
    ///  -  Speed: 1.60 ns/call (Core i7 @3.40GHz)
    ///
    /// For further analysis see
    ///   <http://vigna.di.unimi.it/ftp/papers/xorshift.pdf>
    #[derive(Clone, Copy)]
    pub struct Prng(u64);

    impl Prng {
        /// Returns a new Prng instance.
        pub fn new(seed: u64) -> Prng {
            Prng(seed)
        }

        /// Returns a random unsigned 64-bit number.
        pub fn rand64(&mut self) -> u64 {
            self.0 ^= self.0 >> 12;
            self.0 ^= self.0 << 25;
            self.0 ^= self.0 >> 27;
            u64::wrapping_mul(self.0, 2685821657736338717)
        }
    }
}

/// The zobrist module provides an implementation of the Zobrist hashing algorithm, which is
/// compatible with the Stockfish chess engine. This hashing method is commonly used in board games
/// such as chess to implement transposition tables, a form of caching that allows the engine to
/// avoid calculating the same position more than once.
pub mod zobrist {
    use crate::stockfish::misc;
    use crate::stockfish::types::*;

    static mut PSQ: [[Key; 64]; 16] = [[Key(0); 64]; 16];
    static mut ENPASSANT: [Key; 8] = [Key(0); 8];
    static mut CASTLING: [Key; 16] = [Key(0); 16];
    static mut SIDE: Key = Key(0);

    /// Returns the hash key for a given piece and square.
    pub fn psq(pc: u32, s: u32) -> Key {
        assert!(pc < 16, "piece index out of bounds");
        assert!(s < 64, "square index out of bounds");

        // SAFETY: The `pc` and `s` indices must be within bounds.
        unsafe { PSQ[pc as usize][s as usize] }
    }

    /// Returns the hash key for a given file for en passant.
    pub fn enpassant(f: u32) -> Key {
        assert!(f < 8, "en-passant index out of bounds");

        // SAFETY: The `f` index must be within bounds.
        unsafe { ENPASSANT[f as usize] }
    }

    /// Returns the hash key for a given set of castling rights.
    pub fn castling(cr: u32) -> Key {
        assert!(cr < 16, "castling right index out of bounds");

        // SAFETY: The `cr` index must be within bounds.
        unsafe { CASTLING[cr as usize] }
    }

    /// Returns the hash key for the side to move.
    pub fn side() -> Key {
        // SAFETY: No special safety measures are needed as this function just returns a single value.
        unsafe { SIDE }
    }

    /// Initializes the various arrays used to compute hash keys.
    /// SAFETY: This function is not thread safe and should only be called once during the
    /// program's startup.
    pub fn init() {
        let mut rng = misc::Prng::new(1070372);

        // SAFETY: Here we ensure we never try to write out of bounds.
        unsafe {
            #[allow(clippy::needless_range_loop)]
            for i in 1..15 {
                if i != 7 && i != 8 {
                    for s in 0..64 {
                        PSQ[i][s] = Key(rng.rand64());
                    }
                }
            }

            for item in &mut ENPASSANT {
                *item = Key(rng.rand64());
            }

            for item in &mut CASTLING {
                *item ^= Key(rng.rand64());
            }

            SIDE = Key(rng.rand64());
        }
    }
}

use shakmaty::{Castles, CastlingSide, Color, EnPassantMode, Piece, Position, Role};

use types::*;

/// Computes the Stockfish compatible Zobrist hash of a given chess position.
///
/// This function uses the Zobrist hashing algorithm to compute a unique hash for the given
/// position. Zobrist hashing is commonly used in board games such as chess to implement
/// transposition tables, a form of caching that allows the engine to avoid calculating the same
/// position more than once.
///
/// # Arguments
///
/// * `position`: A reference to the position to compute the hash for. The position is expected to
/// implement the `Position` trait.
///
/// # Returns
///
/// * A `u64` value representing the hash of the position.
pub fn stockfish_hash(position: &dyn Position) -> u64 {
    let mut hash = Key(0);

    // Order optimized for cache efficiency.
    let board = position.board();
    for role in Role::ALL {
        for color in [Color::Black, Color::White] {
            let piece = role.of(color);
            let index = stockfish_piece(&piece);
            for square in board.by_piece(piece) {
                hash ^= zobrist::psq(index, square as u32);
            }
        }
    }

    if let Some(ep_square) = position.ep_square(EnPassantMode::Legal) {
        hash ^= zobrist::enpassant(ep_square.file() as u32)
    }

    if position.turn() == Color::Black {
        hash ^= zobrist::side();
    }

    let castles = position.castles();
    if !castles.is_empty() {
        hash ^= zobrist::castling(stockfish_castling_right(castles));
    }

    hash.0
}

#[inline(always)]
fn stockfish_piece(piece: &Piece) -> u32 {
    match piece {
        Piece {
            color: Color::White,
            role: Role::Pawn,
        } => 1,
        Piece {
            color: Color::White,
            role: Role::Knight,
        } => 2,
        Piece {
            color: Color::White,
            role: Role::Bishop,
        } => 3,
        Piece {
            color: Color::White,
            role: Role::Rook,
        } => 4,
        Piece {
            color: Color::White,
            role: Role::Queen,
        } => 5,
        Piece {
            color: Color::White,
            role: Role::King,
        } => 6,
        Piece {
            color: Color::Black,
            role: Role::Pawn,
        } => 9,
        Piece {
            color: Color::Black,
            role: Role::Knight,
        } => 10,
        Piece {
            color: Color::Black,
            role: Role::Bishop,
        } => 11,
        Piece {
            color: Color::Black,
            role: Role::Rook,
        } => 12,
        Piece {
            color: Color::Black,
            role: Role::Queen,
        } => 13,
        Piece {
            color: Color::Black,
            role: Role::King,
        } => 14,
    }
}

#[inline(always)]
fn stockfish_castling_right(castles: &Castles) -> u32 {
    let mut castling_right: CastlingRight = CastlingRight(0);

    if castles.has(Color::White, CastlingSide::KingSide) {
        castling_right |= CastlingRight(1);
    }
    if castles.has(Color::White, CastlingSide::QueenSide) {
        castling_right |= CastlingRight(2);
    }
    if castles.has(Color::Black, CastlingSide::KingSide) {
        castling_right |= CastlingRight(4);
    }
    if castles.has(Color::Black, CastlingSide::QueenSide) {
        castling_right |= CastlingRight(8);
    }

    castling_right.0
}