use std::f32::INFINITY;

use crate::{moves::{Move, MoveKind}, board::io::IO};

use super::{Grossmeister, ttable::{NodeType, TTABLE_SIZE, TranspositionTableItem}};

const SCORE_MATE: f32 = 20_000.0;

#[derive(Debug, Default, PartialEq)]
pub struct PerftResult {
    leaf_nodes: u64,
    captures: u64,
    en_passants: u64,
    castles: u64,
    checks: u64,
}

impl Grossmeister {
    // Mate distance pruning
    // If a non-zero value (mating score) is returned, branch could be safely pruned
    pub fn MDP(alpha: &mut f32, beta: &mut f32, root_distance: u8) -> f32 {
        { // Update bound in case we are winning
            let mating_score = SCORE_MATE - root_distance as f32;
            if mating_score < *beta {
                *beta = mating_score;
                if *alpha >= mating_score {
                    return mating_score
                }
            }
        }
        { // Update bound in case we are losing
            let mating_score = -SCORE_MATE + root_distance as f32;
            if mating_score > *alpha {
                *alpha = mating_score;
                if *beta <= mating_score {
                    return mating_score
                }
            }
        }
        0.0
    }

    pub fn negamax_search(&mut self, mut alpha: f32, mut beta: f32, depth_left: u8, root_distance: u8) -> (f32, Vec<Move>) {
        let mut principal_variation = Vec::new();
        let color = self.board.color();

        if self.board.threefold_repetition() {
            return (0.0, principal_variation);
        }

        if let Some(transposition) = self.transposition() {
            if transposition.depth == depth_left {
                match transposition.node_type {
                    NodeType::PV => { // PV-nodes have exact score
                        principal_variation.push(transposition.mov);
                        return (transposition.score, principal_variation);
                    }
                    NodeType::Cut => {
                        if transposition.score >= beta {
                            principal_variation.push(transposition.mov);
                            return (beta, principal_variation);
                        }
                    }
                    NodeType::All => {
                        if transposition.score <= alpha {
                            principal_variation.push(transposition.mov);
                            return (alpha, principal_variation);
                        }
                    }
                }
            }
        }

        if depth_left == 0 {
            return (self.quiscence(alpha, beta, root_distance), principal_variation);
        }

        // Mate distance pruning
        let mating_score = Grossmeister::MDP(&mut alpha, &mut beta, root_distance);
        if mating_score != 0.0 {
            return (mating_score, principal_variation)
        }

        let mut should_pv_search = true;
        let mut legal_move_found = false;

        self.cleanup_selector();
        while let Some(mov) = self.next_move() {
            let ep_target_before = self.board.ep_target;
            let castling_rights_before = self.board.castling_rights;
            let hash_before = self.board.hash;
            let captured_piece = self.board.make_move(mov);

            if !self.board.is_king_in_check(color) {
                legal_move_found = true;

                let (mut score, mut subtree_pv) = if should_pv_search {
                    // Assume PV-node is high in the list (if move ordering is good)
                    self.negamax_search(-beta, -alpha, depth_left - 1, root_distance + 1)
                } else {
                    // After we have PV-node (that raised alpha) all other nodes will be searched
                    // with zero-window first to confirm this assumption
                    // TODO: changing 0.001 -> 0.0001 leads to a weird bug
                    let result = self.negamax_search(-(alpha + 0.001), -alpha, depth_left - 1, root_distance + 1);
                    // In case some of the other nodes raises alpha, then it's true PV node now,
                    // let's research with full window to find its exact value
                    if -result.0 > alpha {
                        self.negamax_search(-beta, -alpha, depth_left - 1, root_distance + 1)
                    } else {
                        result
                    }
                };
                score *= -1.;
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);

                if score >= beta {
                    self.transposition_table[(self.board.hash % TTABLE_SIZE) as usize] = Some(TranspositionTableItem {
                        hash: self.board.hash,
                        mov,
                        depth: depth_left, // TODO: should be actual depth searched
                        node_type: NodeType::Cut,
                        score,
                    });

                    if mov.kind == MoveKind::Quiet {
                        self.register_killer(mov);
                    }

                    return (beta, principal_variation);
                }
                if score > alpha {
                    alpha = score;
                    should_pv_search = false; // Once we have PV-node we can start zero-window searching
                    principal_variation = Vec::with_capacity(depth_left as usize);
                    principal_variation.push(mov);
                    principal_variation.append(&mut subtree_pv);

                    self.transposition_table[(self.board.hash % TTABLE_SIZE) as usize] = Some(TranspositionTableItem {
                        hash: self.board.hash,
                        mov,
                        depth: depth_left, // TODO: should be actual depth searched
                        node_type: NodeType::PV,
                        score,
                    });
                } else if self.transposition().is_none() {
                    self.transposition_table[(self.board.hash % TTABLE_SIZE) as usize] = Some(TranspositionTableItem {
                        hash: self.board.hash,
                        mov,
                        depth: depth_left, // TODO: should be actual depth searched
                        node_type: NodeType::All,
                        score,
                    });
                }
            } else {
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
            }

            // Could not finish in time, return what we have so far
            if self.should_halt.load(std::sync::atomic::Ordering::SeqCst) {
                break;
            }
        }

        if !legal_move_found {
            if self.board.is_king_in_check(color) {
                return (-SCORE_MATE + root_distance as f32, principal_variation);
            } else {
                return (0.0, principal_variation);
            }
        }

        (alpha, principal_variation)
    }

    pub fn quiscence(&mut self, mut alpha: f32, mut beta: f32, root_distance: u8) -> f32 {
        let color = self.board.color();

        if self.board.threefold_repetition() {
            return 0.0
        }

        // Mate distance pruning
        let mating_score = Grossmeister::MDP(&mut alpha, &mut beta, root_distance);
        if mating_score != 0.0 {
            return mating_score
        }


        let mut tactical_only = false;
        if !self.board.is_king_in_check(color) {
            // If we are not in check, we can evaluate stand pat
            let stand_pat = self.evaluate();

            if stand_pat >= beta {
                return beta;
            }
            if alpha < stand_pat {
                alpha = stand_pat;
            }

            // If we are not in check, we can only search tactical moves
            tactical_only = true;
        }

        let mut legal_move_found = false;
        self.cleanup_selector();
        while let Some(mov) = if tactical_only { self.next_tactical() } else { self.next_move() } {
            let ep_target_before = self.board.ep_target;
            let castling_rights_before = self.board.castling_rights;
            let hash_before = self.board.hash;
            let captured_piece = self.board.make_move(mov);

            if !self.board.is_king_in_check(color) {
                legal_move_found = true;
                let evaluation = -self.quiscence(-beta, -alpha, root_distance + 1);
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);

                if evaluation >= beta {
                    return beta; // Fail-hard beta-cutoff
                }
                if evaluation > alpha {
                    alpha = evaluation;
                }
            } else {
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
            }
        }

        if !legal_move_found && self.board.is_king_in_check(color) {
            return -SCORE_MATE + root_distance as f32
        }

        alpha
    }

    pub fn iterative_deepening(&mut self, max_depth: u8) -> (f32, Vec<Move>) {
        let mut result = None;
        let mut depth = 1;
        let mut alpha = -INFINITY;
        let mut beta = INFINITY;
        let window_size = 0.25;
        let mut gradual_widening_counter = 0;

        while depth <= max_depth {
            if self.debug {
                println!("info string window {:?}", (alpha, beta));
            }

            let search_result = self.negamax_search(alpha, beta, depth, 0);

            if self.should_halt.load(std::sync::atomic::Ordering::SeqCst) {
                println!("info string halting search");
                break;
            }

            if search_result.0 <= alpha { // Alpha-cutoff
                gradual_widening_counter += 1;
                beta = alpha + window_size * 0.1;
                alpha = search_result.0 - window_size * 2.0f32.powi(gradual_widening_counter);
                println!("info depth {} score upperbound {:.0}", depth, beta * 100.0);
                continue;
            }
            if search_result.0 >= beta { // Beta-cutoff
                gradual_widening_counter += 1;
                alpha = beta - window_size * 0.1;
                beta = search_result.0 + window_size * 2.0f32.powi(gradual_widening_counter);
                println!("info depth {} score lowerbound {:.0}", depth, alpha * 100.0);
                continue;
            }

            if search_result.1.is_empty() {
                panic!("Why the fuck no moves?");
            }

            // Print UCI info
            print!("info depth {} score ", depth);
            if search_result.0.abs() >= SCORE_MATE - 128f32 { // TODO: VALUE_WIN - MAX_PLY
                let mate_distance = (SCORE_MATE - search_result.0.abs()) * if search_result.0 > 0.0 {
                    1.
                } else {
                    -1. // -N means we are mated in N plies
                };
                print!("mate {:.0} ", (mate_distance / 2.0).ceil());
            } else {
                print!("cp {:.0} ", search_result.0 * 100.0)
            }
            print!("pv ");
            for mov in &search_result.1 {
                print!("{} ", mov);
            }
            println!();

            // If our score is mate in N, we break at depth N
            if depth as f32 >= SCORE_MATE - search_result.0.abs() {
                result = Some(search_result);
                break;
            }

            depth += 1;
            gradual_widening_counter = 0;
            alpha = search_result.0 - window_size;
            beta = search_result.0 + window_size;
            result = Some(search_result);
        }

        if result.is_none() {
            println!("info string could not find move in time, will re-run depth-1 search to avoid panic");
            result = Some(self.negamax_search(-INFINITY, INFINITY, 1, 0))
        }

        println!("info hashfull {}", 1000 * self.transposition_table.iter().filter(|item| item.is_some()).count() / self.transposition_table.len());

        match result {
            Some(r) => {
                print!("bestmove {}", r.1[0]);
                if r.1.len() > 1 {
                    print!(" ponder {}", r.1[1])
                }
                println!();
                r
            }
            None => panic!("No move found by search")
        }
    }

    pub fn perft(&mut self, depth: u8, print: bool) -> PerftResult {
        let mut result = PerftResult::default();

        if depth == 0 {
            result.leaf_nodes = 1;
            return result;
        }
        let color = self.board.color();

        if print {
            // println!("Running perft for depth {}. Color to move is {:?}\n{} moves available", depth, color, moves.len());
            // println!("{} moves available", moves.len());
        }

        self.cleanup_selector();
        while let Some(mov) = self.next_move() {
            let ep_target_before = self.board.ep_target;
            let castling_rights_before = self.board.castling_rights;
            let hash_before = self.board.hash;
            let captured_piece = self.board.make_move(mov);
            // King can not be in check after our own move
            if !self.board.is_king_in_check(color) {
                if depth == 1 {
                    match mov.kind {
                        MoveKind::Capture => {
                            result.captures += 1;
                        }
                        MoveKind::EnPassant => {
                            result.en_passants += 1;
                            result.captures += 1;
                        }
                        MoveKind::Castle => {
                            result.castles += 1;
                        }
                        _ => {}
                    }
                    if self.board.is_king_in_check(color.flip()) {
                        result.checks += 1;
                    }
                }

                if print {
                    println!("{:?}", mov);
                    self.board.print();
                }
                let subtree_result = self.perft(depth - 1, print);

                result.leaf_nodes += subtree_result.leaf_nodes;
                result.captures += subtree_result.captures;
                result.checks += subtree_result.checks;
                result.castles += subtree_result.castles;
                result.en_passants += subtree_result.en_passants;

            }
            self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
        }

        if print {
            println!("Found {} leaf nodes in this subtree (depth {})", result.leaf_nodes, depth);
        }

        result
    }
}


#[cfg(test)]
mod tests {
    use crate::{board::{Board, io::IO}, square::Square, moves::{Move, MoveKind}, grossmeister::{Grossmeister, search::PerftResult}};
    use super::SCORE_MATE;

    #[test]
    fn perft() {
        let board = Board::new();
        let mut gm = Grossmeister::new(board);

        assert_eq!(gm.perft(0, false), PerftResult { leaf_nodes: 1, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(1, false), PerftResult { leaf_nodes: 20, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(2, false), PerftResult { leaf_nodes: 400, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(3, false), PerftResult { leaf_nodes: 8902, captures: 34, en_passants: 0, castles: 0 , checks: 12 });
        assert_eq!(gm.perft(4, false), PerftResult { leaf_nodes: 197281, captures: 1576, en_passants: 0, castles: 0 , checks: 469 });
        // assert_eq!(board.perft(5, false), PerftResult { leaf_nodes: 4865609, captures: 82719, en_passants: 258, castles: 0, checks: 27351 });
    }

    #[test]
    fn position_perft() {
        let fen = String::from("r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq - ");
        let board = Board::from_FEN(fen);
        let mut gm = Grossmeister::new(board);

        assert_eq!(gm.perft(0, false), PerftResult { leaf_nodes: 1, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(1, false), PerftResult { leaf_nodes: 48, captures: 8, en_passants: 0, castles: 2 , checks: 0 });
        assert_eq!(gm.perft(2, false), PerftResult { leaf_nodes: 2039, captures: 351, en_passants: 1, castles: 91 , checks: 3 });
        assert_eq!(gm.perft(3, false), PerftResult { leaf_nodes: 97862, captures: 17102, en_passants: 45, castles: 3162, checks: 993 });
        // assert_eq!(board.perft(4, false), PerftResult { leaf_nodes: 4085603, captures: 757163, en_passants: 1929, castles: 128013, checks: 25523 });
    }

    #[test]
    fn endgame_perft() {
        let fen = String::from("8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - ");
        let board = Board::from_FEN(fen);
        let mut gm = Grossmeister::new(board);

        assert_eq!(gm.perft(1, false), PerftResult { leaf_nodes: 14, captures: 1, en_passants: 0, castles: 0 , checks: 2 });
        assert_eq!(gm.perft(2, false), PerftResult { leaf_nodes: 191, captures: 14, en_passants: 0, castles: 0 , checks: 10 });
        // assert_eq!(board.perft(3, false), PerftResult { leaf_nodes: 2812, captures: 209, en_passants: 2, castles: 0 , checks: 267 });
    }

    #[test]
    fn mate_in_3() {
        let fen = String::from("2kr1b1r/pp1npppp/2p1bn2/7q/5B2/2NB1Q1P/PPP1N1P1/2KR3R w - - 0 1");
        let board = Board::from_FEN(fen);
        let mut gm = Grossmeister::new(board);
        let (score, pv) = gm.iterative_deepening(8);

        assert_eq!(score, SCORE_MATE - 3.0);
        assert_eq!(pv, vec![
            Move { source: Square::F3, target: Square::C6, kind: MoveKind::Capture },
            Move { source: Square::B7, target: Square::C6, kind: MoveKind::Capture },
            Move { source: Square::D3, target: Square::A6, kind: MoveKind::Quiet },
        ]);
    }

    #[test]
    fn stupid_knight_sac() {
        let fen = String::from("r3k1r1/pp3ppp/1q6/2ppPn2/6P1/1PPP1P2/P1N3KP/R2QR3 b - - 0 18");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        let (_, pv) = gm.iterative_deepening(6);
        assert_eq!(
            pv[0],
            Move { source: Square::F5, target: Square::H4, kind: MoveKind::Quiet },
            "You should save this poor knight from danger!"
        );
    }

    #[test]
    fn weird_bishop_sac() {
        let fen = String::from("r1b1k1nr/p4pp1/1pp1p3/4n2p/1b1qP3/1B1P3N/PPPBQPPP/RN2K2R w KQkq - 7 10");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        let (_, pv) = gm.iterative_deepening(5);
        assert_eq!(
            pv[0],
            Move { source: Square::C2, target: Square::C3, kind: MoveKind::Quiet },
            "You should fork this bastard!"
        );
    }

    #[test]
    fn tricky_mate_them_in_2() {
        let fen = String::from("7k/2P2Q2/3K1b2/p4N1p/P6P/6P1/8/8 w - - 10 61");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert_eq!(SCORE_MATE - score, 3.0); // Mate in 3 plies
    }

    #[test]
    fn tricky_mate_us_in_2() {
        let fen = String::from("8/2P2Q1k/3K1b2/p4N1p/P6P/6P1/8/8 b - - 5 58");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert_eq!(score + SCORE_MATE, 4.0); // Mate in 4 plies
    }

    #[test]
    fn another_mate_in_2() {
        let fen = String::from("7r/2pr1pR1/2p2p2/1p2pN2/p3Pk1P/P2P2b1/1PP1K3/R7 w - - 0 33");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert_eq!(SCORE_MATE - score, 3.0); // Mate in 3 plies
    }
}