use crate::{
    logic::r#move::Move,
    repr::{Board, CoordPair, Piece, Winner},
};
use indicatif::{ProgressIterator, ProgressStyle};
use std::{collections::HashMap, hash::BuildHasherDefault, ops::ControlFlow};

pub struct FutureMoves {
    /// Arena containing all [`Move`]
    arena: Vec<Move>,

    /// Index of the [`Move`] tree's root node
    current_root: Option<usize>,

    /// Current generated depth of the Arena
    current_depth: usize,

    /// Color w.r.t
    agent_color: Piece,

    config: FutureMoveConfig,
}

#[derive(Copy, Clone)]
pub struct FutureMoveConfig {
    /// Max depth of that we should try and traverse
    pub max_depth: usize,

    /// the min depth an arena should fill for pruning to happen
    pub min_arena_depth: usize,

    /// when pruning, keep the top_k # of children
    pub top_k_children: usize,

    // the lower the value, the more conservative the pruning is, what level to stop pruning at?
    // a lower value allows more possible paths
    pub up_to_minus: usize,

    /// Max size of the arena, will not generate more if
    /// the arena is of that size or bigger
    pub max_arena_size: usize,

    pub do_prune: bool,

    pub print: bool,

    pub children_eval_method: ChildrenEvalMethod,
}

impl std::fmt::Display for FutureMoveConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "D{} ", self.max_depth)?;
        if self.do_prune {
            write!(f, "MD{} ", self.min_arena_depth)?;
            write!(f, "K{} ", self.top_k_children)?;
            write!(f, "UM{} ", self.up_to_minus)?;
        } else {
            write!(f, "MD_ ")?;
            write!(f, "K_ ")?;
            write!(f, "UM_ ")?;
        }
        if self.max_arena_size == usize::MAX {
            write!(f, "SMAX ")?;
        } else {
            write!(f, "S{} ", self.max_arena_size)?;
        }

        write!(f, "P{} ", self.do_prune)?;
        write!(f, "C{:?}", self.children_eval_method)?;
        Ok(())
    }
}

#[derive(Debug, Clone, Copy)]
pub enum ChildrenEvalMethod {
    /// Best (by far) strat compared to Max or Min
    Average,
    Max,
    Min,
}

impl FutureMoves {
    pub const fn new(agent_color: Piece, config: FutureMoveConfig) -> Self {
        Self {
            arena: Vec::new(),
            current_root: None,
            current_depth: 0,
            agent_color,
            config,
        }
    }

    /// Return the length of the Arena
    pub fn arena_len(&self) -> usize {
        self.arena.len()
    }

    /// Returns indexes of leaf [`Move`]s, sorted by their depth (increasing order)
    fn leaf_moves(&self) -> Vec<usize> {
        let mut indexes = (0..self.arena.len())
            // we want to select all nodes that don't have children, or are lazy (need to maybe be regenerated)
            .filter(|&idx| {
                let got = &self.arena[idx];
                !got.is_trimmed && !got.tried_children && got.winner == Winner::None
            })
            .filter(|&idx| self.is_connected_to_root(idx))
            .collect::<Vec<usize>>();

        // we want to try and make the tree even in depth
        // by first generating children for younger moves
        indexes.sort_by_key(|&x| self.depth_of(x));
        indexes
    }

    /// Find the current depth of the arena by
    /// looking at leaf moves and finding the smallest value
    fn determine_current_depth(&self) -> Option<usize> {
        // leaf_moves is sorted from min to max depth
        self.leaf_moves().first().map(|&i| self.depth_of(i))
    }

    /// Generate children for all children of `nodes`
    /// only `pub` for the sake of benchmarking
    pub fn extend_layers(&mut self) {
        // recover from partial tree extention
        if let Some(current_depth) = self.determine_current_depth() {
            self.current_depth = current_depth;
        }

        for _ in self.current_depth..self.config.max_depth {
            let pstyle_inner = if cfg!(test) || !self.config.print {
                ""
            } else {
                &format!(
                    "Generating children (depth: {}/{}): ({{pos}}/{{len}}) {{per_sec}}",
                    self.current_depth + 1,
                    self.config.max_depth
                )
            };

            let cf = self
                .leaf_moves()
                .into_iter()
                .filter(|&i| self.depth_of(i) == self.current_depth)
                .collect::<Vec<usize>>()
                .into_iter()
                .progress_with_style(ProgressStyle::with_template(pstyle_inner).unwrap())
                .try_for_each(|node_idx| {
                    self.generate_children(node_idx);
                    self.arena[node_idx].tried_children = true;

                    if self.arena_len() >= self.config.max_arena_size {
                        ControlFlow::Break(())
                    } else {
                        ControlFlow::Continue(())
                    }
                });

            self.prune_bad_children();

            if cf.is_break() {
                // pruning unfinished level
                // let by_depth = self.by_depth(0..self.arena.len());
                // let mut bdh = HashMap::new();
                // for (a, b) in by_depth {
                //     bdh.insert(a, b);
                // }

                // for &i in bdh.get(&(self.current_depth + 1)).unwrap_or(&Vec::new()) {
                //     self.remove(i);
                // }

                // for &i in bdh.get(&self.current_depth).unwrap_or(&Vec::new()) {
                //     self.arena[i].tried_children = false;
                // }

                // self.refocus_tree();

                return;
            }
            self.current_depth += 1;
        }
    }

    fn remove(&mut self, index: usize) {
        if let Some(parent) = self.arena[index].parent {
            self.arena[parent].children.retain(|&j| j != index);
        }

        self.arena[index].parent = None;
    }

    /// Determines if a [`Move`] at index `idx` is connected to `self.current_root`
    /// Returns `false` if `self.current_root` is None
    fn is_connected_to_root(&self, idx: usize) -> bool {
        if let Some(root) = self.current_root {
            let mut current = Some(idx);
            while let Some(parent_idx) = current {
                if parent_idx == root {
                    return true;
                }
                current = self.arena[parent_idx].parent;
            }
        }

        false
    }

    /// Creates children for a parent (`parent_idx`)
    /// Completely unchecked, the caller should be the one who tests to make sure child generation
    /// hasn't already been tried on a parent
    fn generate_children(&mut self, parent_idx: usize) {
        let parent = &self.arena[parent_idx];

        let new_color = !parent.color;

        // use [`Board::all_positions`] here instead of [`Board::possible_moves`]
        // because we use [`Board::what_if`] later and we want to reduce calls to [`Board::propegate_from_dry`]
        let mut new: Vec<Move> = Board::all_positions()
            .flat_map(|coord| {
                parent
                    .board
                    .what_if(coord, new_color)
                    .map(move |x| (coord, x))
            })
            .map(|(coord, new_board)| {
                Move::new(Some(coord), new_board, new_color, self.agent_color)
            })
            .collect();

        if new.is_empty() {
            new.push(Move::new(None, parent.board, new_color, self.agent_color));
        }

        let start_idx = self.arena.len();
        self.arena.extend(new);

        let new_indices = start_idx..self.arena.len();

        for child_idx in new_indices {
            self.set_parent_child(parent_idx, child_idx);
        }
    }

    /// Given an index from `self.arena`, what depth is it at? 0-indexed
    fn depth_of(&self, node_idx: usize) -> usize {
        let mut depth = 0;
        let mut current = Some(node_idx);
        while let Some(parent_idx) = current {
            depth += 1;
            current = self.arena[parent_idx].parent;
        }

        depth - 1
    }

    //// PERF! pre-organize all indexes based on what depth they're at
    /// previously, I did a lookup map based on if a node was visited, still resulted in a full
    /// O(n) iteration each depth
    fn by_depth(&self, indexes: impl Iterator<Item = usize>) -> Vec<(usize, Vec<usize>)> {
        let mut by_depth: HashMap<
            usize,
            Vec<usize>,
            BuildHasherDefault<nohash_hasher::NoHashHasher<usize>>,
        > = HashMap::with_hasher(BuildHasherDefault::default());

        for idx in indexes {
            let depth = self.depth_of(idx);
            if let Some(got) = by_depth.get_mut(&depth) {
                got.push(idx);
            } else {
                by_depth.insert(depth, vec![idx]);
            }
        }
        let mut by_depth_vec: Vec<(usize, Vec<usize>)> = by_depth.into_iter().collect();
        by_depth_vec.sort_by_key(|x| x.0);
        by_depth_vec
    }

    /// Compute `Move.value`, propegating upwards from the furthest out Moves
    /// in the Arena.
    fn compute_values(&mut self, indexes: impl Iterator<Item = usize>) {
        let by_depth_vec = self.by_depth(indexes);

        // reversed so we build up the value of the closest (in time) moves from the future
        for (_, nodes) in by_depth_vec.into_iter().rev() {
            for idx in nodes {
                let children_values = self.arena[idx]
                    .children
                    .iter()
                    .flat_map(|&child| self.arena[child].value)
                    .collect::<Vec<_>>();

                let children_value = match self.config.children_eval_method {
                    ChildrenEvalMethod::Average => children_values
                        .into_iter()
                        .sum::<i128>()
                        .checked_div(self.arena[idx].children.len() as i128),
                    ChildrenEvalMethod::Max => children_values.into_iter().max(),
                    ChildrenEvalMethod::Min => children_values.into_iter().min(),
                }
                .unwrap_or(0);

                // we use `depth` and divided `self_value` by it, idk if this is worth it
                // we should really setup some sort of ELO rating for each commit, playing them against
                // each other or something, could be cool to benchmark these more subjective things, not
                // just performance (cycles/time wise)
                self.arena[idx].value = Some(self.arena[idx].self_value as i128 + children_value);
            }
        }
    }

    /// Return the best move which is a child of `self.current_root`
    pub fn best_move(&self) -> Option<Option<CoordPair>> {
        self.current_root
            .and_then(|x| {
                self.arena[x]
                    .children
                    .iter()
                    // this would be considered `minimax`
                    .max_by_key(|&&idx| self.arena[idx].value)
            })
            .inspect(|&&x| {
                assert_eq!(
                    self.arena[x].color, self.agent_color,
                    "selected move color should be the same as the color of the agent"
                );
            })
            .map(|&x| self.arena[x].coord)
    }

    /// Updates `FutureMoves` based on the current state of the board
    /// The board is supposed to be after the opposing move
    /// Returns whether or not the arena was regenerated (bool)
    pub fn update_from_board(&mut self, board: &Board) -> bool {
        let curr_board = (0..self.arena.len())
            // needs to be every other generation in order to
            // match the agent_color usually root or great-grand child
            .filter(|&idx| self.depth_of(idx) % 2 == 0)
            .find(|&idx| {
                &self.arena[idx].board == board && self.arena[idx].color == !self.agent_color
            });

        if let Some(curr_board_idx) = curr_board {
            self.set_root_idx_raw(curr_board_idx);
            false
        } else {
            if self.config.print && !self.arena.is_empty() {
                println!("regenerating arena from board");
            }
            self.set_root_from_board(*board);
            true
        }
    }

    /// Clear the arena and create and set a root which contains a Board
    pub fn set_root_from_board(&mut self, board: Board) {
        self.arena.clear();
        self.arena
            .push(Move::new(None, board, !self.agent_color, self.agent_color));

        // because we have to regenerate root from a [`Board`]
        // we need to reset the current_depth (fixes `skip_move_recovery`)
        self.current_depth = 0;

        self.set_root_idx_raw(0);
    }

    /// Update current root without modifying or pruning the Arena
    fn update_root_idx_raw(&mut self, idx: usize) {
        self.current_root = Some(idx);
        self.current_depth -= self.depth_of(idx);
    }

    /// Update current root index while pruning and extending the tree (also recalculate values)
    fn set_root_idx_raw(&mut self, idx: usize) {
        self.update_root_idx_raw(idx);

        self.refocus_tree();
        self.extend_layers();
        if self.config.print {
            println!("# of moves stored: {}", self.arena_len());
        }
        self.compute_values(0..self.arena.len());

        // check arena's consistancy
        debug_assert_eq!(self.check_arena().join("\n"), "");
    }

    pub fn set_parent_child(&mut self, parent: usize, child: usize) {
        self.arena[parent].children.push(child);
        self.arena[child].parent = Some(parent);
    }

    /// Checks the consistancy of the Arena (parents and children)
    /// returns a vector of errors ([`String`])
    pub fn check_arena(&self) -> Vec<String> {
        let mut errors = vec![];
        for idx in 0..self.arena.len() {
            let m = &self.arena[idx];
            if let Some(parent) = m.parent {
                if !(0..self.arena.len()).contains(&parent) {
                    errors.push(format!("{}: parent is out of range ({})", idx, parent));
                }

                if !self.arena[parent].children.contains(&idx) {
                    errors.push(format!(
                        "{}: parent ({}) doesn't list {} as child",
                        idx, parent, idx
                    ));
                }
            }

            for &child_idx in &m.children {
                if !(0..self.arena.len()).contains(&child_idx) {
                    errors.push(format!("{}: parent is out of range ({})", idx, child_idx));
                }

                if self.arena[child_idx].parent != Some(idx) {
                    errors.push(format!(
                        "{}: child ({}) does not list self as parent",
                        idx, child_idx
                    ));
                }
            }

            if !self.is_connected_to_root(idx) {
                errors.push(format!("{}: not connected to root in any way", idx));
            }
        }
        errors
    }

    fn prune_bad_children(&mut self) {
        if self.current_depth < self.config.min_arena_depth || !self.config.do_prune {
            return;
        }

        // values are needed in order to prune and see what's best
        self.compute_values(0..self.arena_len());

        let by_depth = self.by_depth(0..self.arena.len());

        for (depth, indexes) in by_depth {
            // TODO! maybe update by_depth every iteration or something?
            if depth > self.current_depth.saturating_sub(self.config.up_to_minus) {
                return;
            }

            // only prune moves of the agent
            if indexes.first().map(|&i| self.arena[i].color) != Some(self.agent_color) {
                continue;
            }

            for idx in indexes {
                let mut m = self.arena[idx].clone();
                if m.is_trimmed {
                    continue;
                }
                m.is_trimmed = true;
                m.sort_children(&self.arena);
                if m.children.len() > self.config.top_k_children {
                    let drained = m.children.drain(self.config.top_k_children..);
                    for idx in drained {
                        self.arena[idx].parent = None;
                    }
                }
                self.arena[idx] = m;
            }
        }

        // rebuild tree to exclude the things that were pruned
        self.refocus_tree();
    }

    /// Rebuilds the Arena based on `self.current_root`, prunes unrelated nodes
    fn refocus_tree(&mut self) {
        let Some(root) = self.current_root else {
            return;
        };

        // make sure `root` doesn't reference another node
        self.arena[root].parent = None;

        let mut retain = vec![false; self.arena.len()];

        // stack is going to be AT MAXIMUM, the size of the array,
        // so lets just pre-allocate it
        let mut stack: Vec<usize> = Vec::with_capacity(self.arena.len());
        stack.push(root);

        // traverse children of the current root
        while let Some(idx) = stack.pop() {
            retain[idx] = true;
            stack.extend(self.arena[idx].children.iter());
        }

        let mut index_map = vec![None; self.arena.len()];

        let (indexes, moves): (Vec<(usize, usize)>, Vec<Move>) = retain
            .into_iter()
            .enumerate() // old_idx
            .zip(self.arena.drain(..))
            .filter(|&((_, keep), _)| keep) // filter out un-related nodes
            .map(|((old_idx, _), node)| (old_idx, node))
            .enumerate() // new_idx
            .map(|(new_idx, (old_idx, node))| ((new_idx, old_idx), node))
            .unzip();

        for (new_idx, old_idx) in indexes {
            index_map[old_idx] = Some(new_idx);
        }

        self.arena = moves
            .into_iter()
            .map(|mut node| {
                node.parent = node
                    .parent
                    .and_then(|parent| index_map.get(parent).and_then(|&x| x));

                for c in node.children.as_mut_slice() {
                    *c = index_map
                        .get(*c)
                        .and_then(|&x| x)
                        .expect("index_map should contain the child's index");
                }

                node
            })
            .collect();

        self.current_root = index_map[root];
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const FUTURE_MOVES_CONFIG: FutureMoveConfig = FutureMoveConfig {
        max_depth: 3, // we want great-grand children for traversing moves
        min_arena_depth: 0,
        top_k_children: 1,
        up_to_minus: 0,
        max_arena_size: 100,
        do_prune: false,
        print: false,
        children_eval_method: ChildrenEvalMethod::Max,
    };

    #[test]
    fn prune_tree_test() {
        let mut futm = FutureMoves::new(Piece::Black, FUTURE_MOVES_CONFIG);

        futm.arena.push(Move {
            coord: None,
            board: Board::new(),
            winner: Winner::None,
            parent: None,
            children: Vec::new(),
            value: None,
            self_value: 0,
            color: Piece::Black,
            is_trimmed: false,
            tried_children: false,
        });

        futm.update_root_idx_raw(0);

        // child 1
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));
        futm.set_parent_child(0, 1);

        // dummy (2)
        futm.arena.push(Move::new(
            Some((9, 9).into()),
            Board::new(),
            Piece::White,
            Piece::Black,
        ));

        // 3
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));
        futm.set_parent_child(0, 3);

        // 4
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));
        futm.set_parent_child(0, 4);

        assert_eq!(futm.arena_len(), 5);

        futm.refocus_tree();
        assert_eq!(futm.arena_len(), 4);
        assert_eq!(futm.arena[0].children.len(), 3);

        assert_ne!(
            futm.arena[2].coord,
            Some((9, 9).into()),
            "dummy value still exists"
        );
    }

    #[test]
    fn expand_layer_test() {
        let mut futm = FutureMoves::new(Piece::Black, FUTURE_MOVES_CONFIG);
        futm.config.max_depth = 1;

        futm.arena.push(Move::new(
            None,
            Board::new().starting_pos(),
            Piece::Black,
            Piece::Black,
        ));

        futm.update_root_idx_raw(0);

        futm.extend_layers();
        assert_eq!(futm.arena_len(), 5);

        // move to a child
        futm.update_root_idx_raw(1);
        futm.refocus_tree();
        assert_eq!(futm.arena_len(), 1);

        // make sure current_root is properly updated
        assert_eq!(futm.current_root, Some(0));

        futm.extend_layers();
        assert!(
            futm.arena_len() > 1,
            "extend_layer didn't grow arena after refocus"
        );
    }

    #[test]
    fn depth_of_test() {
        let mut futm = FutureMoves::new(Piece::Black, FUTURE_MOVES_CONFIG);

        futm.arena.push(Move {
            coord: None,
            board: Board::new(),
            winner: Winner::None,
            parent: None,
            children: vec![],
            value: None,
            self_value: 0,
            color: Piece::Black,
            is_trimmed: false,
            tried_children: false,
        });

        futm.update_root_idx_raw(0);

        // child 1
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));
        futm.set_parent_child(0, 1);

        // dummy
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));

        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));

        futm.set_parent_child(1, 3);

        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));

        futm.set_parent_child(0, 4);

        assert_eq!(futm.depth_of(3), 2);
    }

    #[test]
    fn by_depth_test() {
        let mut futm = FutureMoves::new(Piece::Black, FUTURE_MOVES_CONFIG);

        futm.arena.push(Move {
            coord: None,
            board: Board::new(),
            winner: Winner::None,
            parent: None,
            children: vec![1],
            value: None,
            self_value: 0,
            color: Piece::Black,
            is_trimmed: false,
            tried_children: false,
        });

        futm.update_root_idx_raw(0);

        // child 1
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));
        futm.set_parent_child(0, 1);

        // dummy
        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));

        futm.arena
            .push(Move::new(None, Board::new(), Piece::White, Piece::Black));
        futm.set_parent_child(1, 3);

        assert_eq!(
            futm.by_depth(0..futm.arena.len()),
            vec![(0, vec![0, 2]), (1, vec![1]), (2, vec![3])]
        );
    }

    /// tests whether or not FutureMoves can recover from multiple skips and then manually regenerating the arena
    #[test]
    fn skip_move_recovery() {
        let mut futm = FutureMoves::new(Piece::Black, FUTURE_MOVES_CONFIG);
        let mut board = Board::new().starting_pos();

        // replay of a test I did
        // TODO! make this as small of a test as possible
        let moves = vec![
            (Some((5, 4)), Piece::Black),
            (Some((5, 5)), Piece::White),
            (Some((5, 6)), Piece::Black),
            (Some((6, 4)), Piece::White),
            (Some((7, 3)), Piece::Black),
            (Some((7, 4)), Piece::White),
            (Some((7, 5)), Piece::Black),
            (Some((2, 4)), Piece::White),
            (Some((1, 4)), Piece::Black),
            (Some((1, 5)), Piece::White),
            (Some((1, 6)), Piece::Black),
            (Some((0, 6)), Piece::White),
            (Some((3, 2)), Piece::Black),
            (Some((1, 7)), Piece::White),
            (None, Piece::Black), // black skips a move
            (Some((0, 4)), Piece::White),
            (None, Piece::Black), // black skips a move
            (Some((4, 2)), Piece::White),
        ];

        for (i, (coords, color)) in moves.into_iter().enumerate() {
            if color == futm.agent_color {
                // my turn

                // seperate variable because we want it to always be executed
                let update_result = futm.update_from_board(&board);

                // make sure that the arena should only be
                // regenerated on the first move
                if i > 0 && update_result {
                    panic!("board regenerated on move #{}", i);
                }

                let best_move = futm.best_move();
                assert!(
                    best_move.is_some(),
                    "best_move (#{}) resulted in ABSOLUTE None: {:?}",
                    i,
                    best_move
                );

                if coords.is_none() {
                    assert_eq!(best_move, Some(None));
                } else {
                    assert_ne!(best_move, Some(None));
                }
            }

            if let Some(coord) = coords {
                board.place(coord.into(), color).unwrap();
            }
        }
    }
}