//! Find locations.

use alloc::vec::Vec;
use hashbrown::hash_set::HashSet;
use mino::input::Kicks;
use mino::piece::{Shape, Spawn};
use mino::{input, Loc, Mat, Movement, Piece};

// Generic arguments legend
// ========================
// - T: piece type (e.g. srs::PieceType)
// - 'm: matrix lifetime
// - 'c: T::cells() lifetime
// - 'k: T::kicks() lifetime
// - 'a: general purpose lifetime

static ALL_INPUTS: &[Movement] = &[
    Movement::LEFT,
    Movement::RIGHT,
    Movement::CW,
    Movement::CCW,
    Movement::FLIP,
];

/// Yields all of the locations reachable by the given peice on the given matrix.
pub fn find_locations<'a, 'c, 'k, T>(matrix: &'a Mat, piece_ty: T) -> FindLocations<T>
where
    T: Shape<'c> + Kicks<'k> + Spawn + Clone + 'a,
{
    FindLocations::new(matrix, piece_ty)
}

/// Algorithm used to search for reachable locations on a given board state. The current
/// algorithm is just depth-first search but may be subject to change after benchmarking
/// experiments.
///
/// [`FindLocations`] is an [`Iterator`], so you can use that interface to obtain the next
/// available location or loop over all of them.
pub struct FindLocations<'m, T> {
    matrix: &'m Mat,
    piece_ty: T,
    // TODO: allow these two collections to be extracted and reused
    stack: Vec<Loc>,
    visited: HashSet<Loc, core::hash::BuildHasherDefault<ahash::AHasher>>,
}

impl<'m, 'c, 'k, T> FindLocations<'m, T>
where
    T: Shape<'c> + Kicks<'k> + Spawn + Clone,
{
    /// Constructs a new instance of the location finding algorithm.
    fn new(matrix: &'m Mat, piece_ty: T) -> Self {
        // TODO: use with_capacity
        let mut stack = Vec::default();
        let mut visited = HashSet::default();

        let init_pc = Piece::spawn(piece_ty.clone());
        let game_over = init_pc.cells().intersects(matrix);
        if !game_over {
            stack.push(init_pc.loc);
            visited.insert(init_pc.loc);
        }

        Self {
            matrix,
            piece_ty,
            stack,
            visited,
        }
    }

    fn push(&mut self, loc: Loc) {
        // 'visited' set prevents cycles
        if self.visited.insert(loc) {
            self.stack.push(loc);
            // self.stats.max_height = max(self.stats.max_height, self.stack.len());
        }
    }

    fn pop(&mut self) -> Option<Piece<T>> {
        let ty = self.piece_ty.clone();
        self.stack.pop().map(|loc| Piece { ty, loc })
    }
}

impl<'m, 'c, 'k, T> Iterator for FindLocations<'m, T>
where
    T: Shape<'c> + Kicks<'k> + Spawn + Clone,
{
    type Item = Loc;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(pc0) = self.pop() {
            // TODO: configure capabilities, e.g. is 180 enabled, are kicks enabled

            // push all locations reachable by performing an input
            for &inp in ALL_INPUTS.iter() {
                let mut pc = pc0.clone();
                if inp.perform(&mut pc, self.matrix) {
                    self.push(pc.loc);
                }
            }

            let mut pc = pc0;
            if input::drop(&mut pc, self.matrix) {
                // piece was floating, so drop it and analyze that later
                // TODO: configure capability to do soft drop
                self.push(pc.loc);
            } else {
                // piece was not floating so it's a valid final piece location.
                return Some(pc.loc);
            }
        }
        None
    }
}

impl<'m, 'c, 'k, T: Shape<'c> + Kicks<'k> + Spawn + Clone> core::iter::FusedIterator
    for FindLocations<'m, T>
{
    // CORRECTNESS: once `self.stack` is empty, `next()` will always return `None`.
}

#[cfg(test)]
mod test {
    use super::*;
    use mino::srs::PieceType;
    use mino::{mat, Rot};

    use alloc::vec::Vec;

    fn find_missing(ty: PieceType, lhs: &[Loc], rhs: &[Loc]) -> Option<Loc> {
        lhs.iter().find_map(|&loc_l| {
            let in_rhs = rhs.iter().any(|&loc_r| {
                let pc_l = Piece { ty, loc: loc_l };
                let pc_r = Piece { ty, loc: loc_r };
                pc_l.cells() == pc_r.cells()
            });
            if in_rhs {
                None
            } else {
                Some(loc_l)
            }
        })
    }

    fn test_find_locs<E>(matrix: &Mat, ty: PieceType, expected: E)
    where
        E: IntoIterator,
        Loc: From<E::Item>,
    {
        let expected = expected.into_iter().map(Loc::from).collect::<Vec<_>>();
        let found = find_locations(matrix, ty).collect::<Vec<_>>();
        if let Some(exp) = find_missing(ty, &expected, &found) {
            panic!("{exp:?} expected but not found");
        } else if let Some(fnd) = find_missing(ty, &found, &expected) {
            panic!("{fnd:?} returned but not expected");
        }

        // XXX(iitalics): currently its OK for find_locations() to yield locations with
        // identical cells, since a transposition table will deduplicate them later on in
        // the engine logic. however it may turn out to be a performance win to do some
        // deduplication early to reduce pressure on the transposition table.
        assert!(expected.len() <= found.len());
        //assert_eq!(expected.len(), found.len());
    }

    #[test]
    fn test_o_empty() {
        test_find_locs(Mat::EMPTY, PieceType::O, (0..=8).map(|x| (x, 0, Rot::N)));
    }

    #[test]
    fn test_i_empty() {
        test_find_locs(
            Mat::EMPTY,
            PieceType::I,
            [
                // N/S
                (1..=7, 0, Rot::N),
                // E/W
                (0..=9, 1, Rot::W),
            ]
            .into_iter()
            .flat_map(|(xs, y, r)| xs.map(move |x| (x, y, r))),
        );
    }

    #[test]
    fn test_o_bumpy() {
        const MAT: &Mat = mat! {
            ".........x";
            ".........x";
            ".x...xx..x";
        };
        test_find_locs(
            MAT,
            PieceType::O,
            [
                (0, 1, Rot::N),
                (1, 1, Rot::N),
                (2, 0, Rot::N),
                (3, 0, Rot::N),
                (4, 1, Rot::N),
                (5, 1, Rot::N),
                (6, 1, Rot::N),
                (7, 0, Rot::N),
                (8, 3, Rot::N),
            ],
        );
    }

    #[test]
    fn test_t_spin() {
        const MAT: &Mat = mat! {
            ".......xx.";
            "xxxxx...xx";
            "xxxxxx.xxx";
        };
        test_find_locs(
            MAT,
            PieceType::T,
            [
                // top: N
                (1..=5, 2, Rot::N),
                (6..=8, 3, Rot::N),
                // top: E
                (0..=4, 3, Rot::E),
                (5..=5, 2, Rot::E),
                (6..=6, 3, Rot::E),
                (7..=8, 4, Rot::E),
                // top: S
                (1..=4, 3, Rot::S),
                (5..=5, 2, Rot::S),
                (6..=6, 3, Rot::S),
                (7..=8, 4, Rot::S),
                // top: W
                (1..=4, 3, Rot::W),
                (5..=5, 2, Rot::W),
                (6..=6, 1, Rot::W),
                (7..=8, 4, Rot::W),
                (9..=9, 3, Rot::W),
                // TSS
                (6..=6, 1, Rot::E),
                (6..=6, 1, Rot::N),
                // TSD
                (6..=6, 1, Rot::S),
            ]
            .into_iter()
            .flat_map(|(xs, y, r)| xs.map(move |x| (x, y, r))),
        );
    }

    #[test]
    fn test_s_spin() {
        const MAT: &Mat = mat! {
            "xxxxxx..xx";
            "xxxxx..xxx";
        };
        test_find_locs(
            MAT,
            PieceType::S,
            [
                // N/S
                (1..=6, 2, Rot::N),
                (7..=7, 1, Rot::N),
                (8..=8, 2, Rot::N),
                // E/W
                (0..=4, 3, Rot::E),
                (5..=6, 2, Rot::E),
                (7..=8, 3, Rot::E),
                // spin
                (6..=6, 1, Rot::S), // equiv (6..=6, 0, Rot::N),
            ]
            .into_iter()
            .flat_map(|(xs, y, r)| xs.map(move |x| (x, y, r))),
        );
    }

    #[test]
    fn test_z_spin() {
        const MAT: &Mat = mat! {
            "xxxxx..xxx";
            "xxxxxx..xx";
        };
        test_find_locs(
            MAT,
            PieceType::Z,
            [
                // N/S
                (1..=4, 2, Rot::N),
                (5..=5, 1, Rot::N),
                (6..=8, 2, Rot::N),
                // E/W
                (0..=4, 3, Rot::E),
                (5..=6, 2, Rot::E),
                (7..=8, 3, Rot::E),
                // spin
                (6..=6, 1, Rot::S), // equiv (6..=6, 0, Rot::N),
            ]
            .into_iter()
            .flat_map(|(xs, y, r)| xs.map(move |x| (x, y, r))),
        );
    }

    #[test]
    fn test_top_out() {
        // giant pillar in the center should cause top-out
        let rows = [0b0000100000u16; 22];
        let mat = Mat::new(&rows);
        test_find_locs(mat, PieceType::I, core::iter::empty::<Loc>());
    }
}