// use crate::networks::PolicyNetwork;
use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold};
use crate::net::encoding::{encode_board_state_perspective, encode_move};
use crate::net::model::ChessModel;
use burn::prelude::Device;
use chess::BoardStatus::{Checkmate, Stalemate};
use chess::Color::{Black, White};
use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
use rand::SeedableRng;
use std::collections::HashMap;

//TODO retain tree after move selection?

#[derive(Clone)]
pub struct Node {
    pub prior: f32,
    pub children: Vec<usize>,
    pub visit_count: u32,
    pub value_sum: f32,
    pub board_state: BoardState,
    pub last_move: Option<ChessMove>, // move that produced this node from its parent
}

impl Node {
    pub fn new(prior: f32, board_state: BoardState, last_move: Option<ChessMove>) -> Node {
        Node {
            prior,
            children: vec![],
            visit_count: 0,
            value_sum: 0.0,
            board_state,
            last_move,
        }
    }

    pub fn value(&self) -> f32 {
        self.value_sum / self.visit_count as f32
    }

    pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize {
        self.children
            .iter()
            .copied()
            .max_by(|&a, &b| {
                ucb_score(self, &arena[a], c_puct)
                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
                    .unwrap()
            })
            .expect("select_child on leaf")
    }
}

fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
    let prior_score = c_puct * child.prior * ((parent.visit_count + 1) as f32).sqrt()
        / (1.0 + child.visit_count as f32);

    let value_score: f32;
    if child.visit_count > 0 {
        value_score = -child.value(); // value from opposing side
    } else {
        value_score = 0.0;
    }

    value_score + prior_score
}

#[derive(Clone, Debug)]
pub struct MctsResults {
    pub board_state: BoardState,
    // Compact encoded move distribution: (encoded_move_index, probability)
    pub move_dist: Vec<(usize, f32)>,
    pub value: f32,
}

impl MctsResults {
    pub fn new(board_state: BoardState, move_dist: Vec<(usize, f32)>, value: f32) -> MctsResults {
        MctsResults {
            board_state,
            move_dist,
            value,
        }
    }
}

#[derive(Debug)]
pub struct MctsConfig {
    pub num_simulations: usize,
    pub c_puct: f32,
    pub dirichlet_alpha: f32,
    pub dirichlet_epsilon: f32,
    pub batch_max: usize,
}

impl MctsConfig {
    pub fn new(
        num_simulations: usize,
        c_puct: f32,
        dirichlet_alpha: f32,
        dirichlet_epsilon: f32,
        batch_max: usize,
    ) -> MctsConfig {
        MctsConfig {
            num_simulations,
            c_puct,
            dirichlet_alpha,
            dirichlet_epsilon,
            batch_max,
        }
    }
}

impl Default for MctsConfig {
    fn default() -> MctsConfig {
        MctsConfig::new(400, 1.0, 0.05, 0.25, 64)
    }
}

pub struct Mcts {
    pub config: MctsConfig,
}

fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
    match status {
        Ongoing => None,
        BoardStateStatus::Stalemate | Threefold | FiftyMove | InsufficientMaterial => Some(0.0),
        BoardStateStatus::WhiteWinner => Some(if stm == White { 1.0 } else { -1.0 }),
        BoardStateStatus::BlackWinner => Some(if stm == Black { 1.0 } else { -1.0 }),
    }
}

fn softmax(logits: &[f32]) -> Vec<f32> {
    if logits.is_empty() {
        return Vec::new();
    }
    let mut max_logit = f32::NEG_INFINITY;
    for &v in logits {
        if v > max_logit {
            max_logit = v;
        }
    }
    if !max_logit.is_finite() {
        max_logit = 0.0;
    }
    let mut exps_sum = 0.0f32;
    let mut probs = vec![0.0f32; logits.len()];
    for (i, &v) in logits.iter().enumerate() {
        let x = (v - max_logit).exp();
        if x.is_finite() {
            probs[i] = x;
            exps_sum += x;
        }
    }
    if exps_sum > 0.0 && exps_sum.is_finite() {
        for p in &mut probs {
            *p /= exps_sum;
        }
    } else {
        let uniform = 1.0 / logits.len() as f32;
        for p in &mut probs {
            *p = uniform;
        }
    }
    probs
}

impl Mcts {
    pub fn search(
        &mut self,
        board_state: &BoardState,
        model: &ChessModel,
        device: &Device,
        training: bool,
    ) -> MctsResults {
        let mut nodes = Vec::<Node>::new();

        let root = 0;
        nodes.push(Node::new(0.0, board_state.clone(), None));

        self.expand(root, &mut nodes, model, device);

        if training {
            // APPLY DIRICHLET NOISE HERE
            self.add_dirichlet_noise(root, &mut nodes);
        }

        let num_sims = self.config.num_simulations;

        for _ in 0..num_sims {
            let mut path = vec![root];
            let mut current = root;

            while !nodes[current].children.is_empty() {
                current = nodes[current].select_child(&nodes, &self.config.c_puct);
                path.push(current);
            }

            let value: f32 = self.expand(current, &mut nodes, model, device);

            let color = nodes[current].board_state.board.side_to_move();
            self.backpropagate(&mut nodes, &path, value, color);
        }

        // Build compact move distribution
        let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
        let stm = board_state.board.side_to_move();
        let total_root_visits: u32 = nodes[root]
            .children
            .iter()
            .map(|&c| nodes[c].visit_count)
            .sum();
        let denom = if total_root_visits > 0 {
            total_root_visits as f32
        } else {
            1.0
        };
        for child_idx in nodes[root].children.iter() {
            let mv = nodes[*child_idx].last_move.expect("move didnt exist");
            if let Some(enc) = encode_move(mv, stm) {
                let prob = nodes[*child_idx].visit_count as f32 / denom;
                move_dist.push((enc, prob));
            }
        }

        MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
    }

    fn expand(
        &mut self,
        node_idx: usize,
        arena: &mut Vec<Node>,
        model: &ChessModel,
        device: &Device,
    ) -> f32 {
        let status = arena[node_idx].board_state.status;
        let stm = arena[node_idx].board_state.board.side_to_move();

        if let Some(term_val) = terminal_value(status, stm) {
            term_val
        } else {
            let state = encode_board_state_perspective(&arena[node_idx].board_state, device);
            let (policy_head, value_head) = model.forward(state.unsqueeze::<4>());

            let legal_moves: Vec<ChessMove> =
                MoveGen::new_legal(&arena[node_idx].board_state.board).collect();

            let logits = policy_head.into_data().to_vec::<f32>().unwrap();
            let probs = softmax(&logits);

            let mut move_priors = Vec::with_capacity(legal_moves.len());
            let mut prior_sum = 0.0f32;

            for &mv in &legal_moves {
                if let Some(enc_idx) = encode_move(mv, stm) {
                    let p = probs[enc_idx];
                    move_priors.push((mv, p));
                    prior_sum += p;
                }
            }

            for (mv, p) in move_priors {
                let normalized_prior = if prior_sum > 0.0 {
                    p / prior_sum
                } else {
                    1.0 / legal_moves.len() as f32
                };

                let mut new_board = arena[node_idx].board_state.clone();
                new_board.apply_move(mv);

                let child_idx = arena.len();
                arena.push(Node::new(normalized_prior, new_board, Some(mv)));
                arena[node_idx].children.push(child_idx);
            }

            value_head.into_data().to_vec::<f32>().unwrap()[0]
        }
    }

    fn backpropagate(&mut self, nodes: &mut [Node], path: &[usize], value: f32, color: Color) {
        for &idx in path {
            let node = &mut nodes[idx];

            if node.board_state.board.side_to_move() == color {
                node.value_sum += value;
            } else {
                node.value_sum -= value;
            }

            node.visit_count += 1;
        }
    }

    fn add_dirichlet_noise(&mut self, node_id: usize, nodes: &mut Vec<Node>) {
        let node_children = &nodes[node_id].children.clone();

        let n = node_children.len();
        if n == 0 {
            return;
        }

        let alpha = self.config.dirichlet_alpha;
        let epsilon = self.config.dirichlet_epsilon;

        let noise = dirichlet_sample(n, alpha);

        for (node_idx, noise_val) in node_children.iter().zip(noise) {
            let prev_prior = nodes[*node_idx].prior;
            nodes[*node_idx].prior = (1.0 - epsilon) * prev_prior + epsilon * noise_val;
        }
    }
}

fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
    use rand_distr::{Distribution, Gamma};

    let gamma = Gamma::new(alpha as f64, 1.0).unwrap();

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap();
    let seed = now.as_nanos() as u64;
    let mut rng = rand::rngs::SmallRng::seed_from_u64(seed);

    let mut samples: Vec<f32> = (0..size).map(|_| gamma.sample(&mut rng) as f32).collect();

    let sum: f32 = samples.iter().sum();

    if sum > 0.0 && sum.is_finite() {
        for x in &mut samples {
            *x /= sum;
        }
    } else if size > 0 {
        let uniform = 1.0 / size as f32;
        for x in &mut samples {
            *x = uniform;
        }
    }

    samples
}

pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> {
    // Choose best child of root by visit count (standard MCTS behavior)
    let root_node = &nodes[root];
    if root_node.children.is_empty() {
        return None;
    }

    let mut best_visits = 0u32;
    let mut best_child_idx: Option<usize> = None;
    for &child_idx in &root_node.children {
        let child = &nodes[child_idx];
        if child.visit_count > best_visits {
            best_visits = child.visit_count;
            best_child_idx = Some(child_idx);
        }
    }

    let child_idx = match best_child_idx {
        Some(i) => i,
        None => return None,
    };

    let bn = &nodes[child_idx];

    println!("best_child visits: {}", bn.visit_count);
    let uci_move = bn.last_move?.to_string();
    Some((bn.board_state.clone(), uci_move))
}

pub fn print_tree_arena(nodes: &[Node], root: usize, depth: usize) {
    let node = &nodes[root];

    println!(
        "{}- Value: {:.2}, Visits: {}, Prior: {:.2}, Side: {:?}",
        "  ".repeat(depth),
        node.value_sum,
        node.visit_count,
        node.prior,
        node.board_state.board.side_to_move()
    );

    for &child_idx in &node.children {
        print_tree_arena(nodes, child_idx, depth + 1);
    }
}

fn material_of(piece: Piece) -> f32 {
    match piece {
        Queen => 9.0,
        Rook => 5.0,
        Bishop => 3.0,
        Knight => 3.0,
        Pawn => 1.0,
        _ => 0.0,
    }
}

const MATERIAL_VALUE_DIVISOR: f32 = 40.0;

pub fn heuristic_eval(board: &Board, perspective: Color) -> f32 {
    let mut value = 0.0;

    // material
    for color in ALL_COLORS {
        for piece in ALL_PIECES {
            let bitboard = board.color_combined(color).0 & board.pieces(piece).0;
            let total_val =
                (bitboard.count_ones() as f32 * material_of(piece)) / MATERIAL_VALUE_DIVISOR;

            if color == perspective {
                value += total_val;
            } else {
                value -= total_val;
            }
        }
    }

    value
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BoardStateStatus {
    Ongoing,
    Stalemate,
    WhiteWinner,
    BlackWinner,
    Threefold,
    FiftyMove,
    InsufficientMaterial,
}

#[derive(Clone, Debug)]
pub struct BoardState {
    pub board: Board,
    pub halfmove_clock: u8,
    pub repetition_table: HashMap<u64, u8>,
    pub status: BoardStateStatus,
}

impl BoardState {
    pub fn new(board: Board, halfmove_clock: u8, repetition_table: HashMap<u64, u8>) -> BoardState {
        BoardState {
            board,
            halfmove_clock,
            repetition_table,
            status: Ongoing,
        }
    }

    pub fn apply_move(&mut self, mv: ChessMove) {
        let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn);
        let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some();

        let is_en_passant = is_pawn_move
            && self.board.en_passant() == Some(mv.get_dest())
            && self.board.piece_on(mv.get_dest()).is_none();

        self.halfmove_clock = self.halfmove_clock.saturating_add(1);

        if is_pawn_move || is_normal_capture || is_en_passant {
            self.halfmove_clock = 0;
        }

        self.board = self.board.make_move_new(mv);

        let board_hash = self.board.get_hash();
        let new_rep = self.repetition_table.get(&board_hash).copied().unwrap_or(0) + 1;
        self.repetition_table.insert(board_hash, new_rep);

        match self.board.status() {
            Checkmate => {
                self.status = if self.board.side_to_move() == White {
                    BoardStateStatus::BlackWinner
                } else {
                    BoardStateStatus::WhiteWinner
                };
                return;
            }
            Stalemate => {
                self.status = BoardStateStatus::Stalemate;
                return;
            }
            _ => {}
        }

        if is_insufficient_material(&self.board) {
            self.status = InsufficientMaterial;
            return;
        }

        if self.halfmove_clock >= 100 {
            self.status = FiftyMove;
            return;
        }

        if new_rep >= 3 {
            self.status = Threefold;
            return;
        }

        self.status = BoardStateStatus::Ongoing;
    }
}

impl Default for BoardState {
    fn default() -> BoardState {
        let board = Board::default();
        let repetition_table = HashMap::from([(board.get_hash(), 1)]);
        BoardState::new(board, 0, repetition_table)
    }
}

// Could be made more comprehensive
pub fn is_insufficient_material(board: &Board) -> bool {
    let pawns = board.pieces(Pawn).0;
    let rooks = board.pieces(Rook).0;
    let queens = board.pieces(Queen).0;

    // Any pawn/rook/queen means this is not trivially dead.
    if pawns != 0 || rooks != 0 || queens != 0 {
        return false;
    }

    let bishops = board.pieces(Bishop).0;
    let knights = board.pieces(Knight).0;

    let bishop_count = bishops.count_ones();
    let knight_count = knights.count_ones();

    match (bishop_count, knight_count) {
        // K vs K
        (0, 0) => true,

        // K+B vs K
        // K+N vs K
        (1, 0) | (0, 1) => true,

        // Other configurations require more analysis.
        _ => false,
    }
}