use super::Puzzle;
use rand::Rng;

impl Puzzle {
    /// O(N)
    pub fn max_queen_col(&mut self) -> usize {
        let (mut max_attack, mut max_queens) = (i32::MIN, Vec::new());
        for col in 0..self.N {
            let attacks = self.attacks(col, self.in_col[col]);

            if attacks > max_attack {
                max_attack = attacks;
                max_queens.clear();
            }
            if attacks == max_attack {
                max_queens.push(col);
            }
        }

        if max_queens.len() > 1 {
            return max_queens[self.rng.gen_range(0..max_queens.len())];
        }
        return max_queens[0];
    }

    /// O(N)
    pub fn min_col_row(&mut self, colI: usize) -> usize {
        let (mut min_row, mut min_ind) = (i32::MAX, Vec::new());
        for row in 0..self.N {
            let attacks = self.attacks(colI, row);

            if attacks < min_row {
                min_row = attacks;
                min_ind.clear();
            }
            if attacks == min_row {
                min_ind.push(row);
            }
        }

        if min_ind.len() > 1 {
            return min_ind[self.rng.gen_range(0..min_ind.len())];
        }
        return min_ind[0];
    }

    /// O(1)
    pub fn update(&mut self, colI: usize, val: i32) {
        assert!(colI < self.N && (val == 1 || val == -1), "Bad arguments to update: {} {}!", colI, val);

        self.per_row[self.in_col[colI]]                    += val;
        self.per_d1[colI + self.N - 1 - self.in_col[colI]] += val;
        self.per_d2[colI + self.in_col[colI]]              += val;

        if val < 0 {
            // Sanity check when removing a queen
            self.in_col[colI] = self.N;
        }
    }

    /// O(1)
    pub fn move_queen(&mut self, colI: usize, rowI: usize) {
        if rowI == self.in_col[colI] { return; }

        self.update(colI, -1);
        self.in_col[colI] = rowI;
        self.update(colI, 1);
    }
}

#[test]
fn test_max_queen_col() {
    let mut e1 = Puzzle::example1();
    assert_eq!(e1.max_queen_col(), 2, "Bad max_queen_col on single max!");

    let mut e2 = Puzzle::example2();
    let mut got = [false; 4];
    for _ in 0..10000 {
        got[e2.max_queen_col()] = true;
    }
    assert!(got.iter().all(|&x| x), "Bad max_queen_col on multiple max!");
}

#[test]
fn test_min_col_row() {
    let mut e1 = Puzzle::example1();
    let mut got = [false; 4];
    for _ in 0..10000 {
        got[e1.min_col_row(0)] = true;
    }
    assert!(got[0] && !got[1] && got[2] && got[3], "Bad min_col_row on multiple (3) min!");
    got = [false; 4];
    for _ in 0..10000 {
        got[e1.min_col_row(2)] = true;
    }
    assert!(got[0] && !got[1] && got[2] && !got[3], "Bad min_col_row on multiple (2) min!");

    let mut e2 = Puzzle::example2();
    assert_eq!(e2.min_col_row(0), 3, "Bad min_col_row on single min!");
    assert_eq!(e2.min_col_row(3), 1, "Bad min_col_row on single min!");
}

#[test]
fn test_update() {
    let mut e1 = Puzzle::example1();
    e1.update(0, -1);
    e1.update(2, -1);
    e1.in_col[0] = 3;
    e1.update(0, 1);
    e1.in_col[2] = 1;
    e1.update(2, 1);
    let e1s = Puzzle::example1();
    assert_eq!(e1, e1s, "Bad update!");
}

#[test]
fn test_move_queen() {
    let e1s = Puzzle::example1();
    let mut e1 = Puzzle::example1();
    e1.move_queen(0, 0);
    e1.move_queen(0, 3);
    assert_eq!(e1, e1s, "Bad move_queen identity!");

    let mut e1 = Puzzle::example1();
    e1.move_queen(2, 0);
    e1.move_queen(2, 1);
    assert_eq!(e1, e1s, "Bad move_queen identity on moved max!");

    let mut e1 = Puzzle::example1();
    e1.move_queen(0, 0);
    e1.move_queen(2, 2);
    let e2 = Puzzle::example2();
    assert_eq!(e1, e2, "Bad move_queen on changed maxes!");
}