use super::{ Puzzle, MAXN };

impl Puzzle {
    pub fn new(N: usize) -> Puzzle {
        Puzzle {
            in_col:  [0; MAXN],
            per_row: [0; MAXN],
            per_d1:  [0; 2 * MAXN - 1],
            per_d2:  [0; 2 * MAXN - 1],

            N: N,
            rng: rand::thread_rng(),
        }
    }

    pub fn clear(&mut self) {
        for i in 0..2 * self.N - 1 {
            if i < self.N {
                self.in_col[i] = 0;
                self.per_row[i] = 0;
            }
            else if i < MAXN {
                assert_eq!(self.in_col[i], 0, "Someone wrote to in_col at index {}!", i);
                assert_eq!(self.per_row[i], 0, "Someone wrote to per_row at index {}!", i);
            }
            self.per_d1[i] = 0;
            self.per_d2[i] = 0;
        }
    }

    /// O(1)
    pub fn attacks(&self, col: usize, row: usize) -> i32 {
        assert!(col < self.N && row < self.N, "Bad attacks index: {} {}!", col, row);

        1 +                 // Always a queen on a column
        self.per_row[row] + // Queens on the row
        self.per_d1[col + self.N - 1 - row] +
        self.per_d2[col + row]
        - 3 * ((self.in_col[col] == row) as i32)
        // If we're in the cell with a queen, remove repeated counts of it
    }

    // | |*| | |    |2|3|2|4|
    // | | |*| |    |3|2|4|2|
    // | | | |*|    |2|4|2|3|
    // |*| | | |    |2|2|3|2|
    pub fn example1() -> Puzzle {
        let mut r = Puzzle::new(4);
        r.in_col[0] = 3;
        r.in_col[1] = 0;
        r.in_col[2] = 1;
        r.in_col[3] = 2;
        r.per_row[0] = 1; r.per_row[1] = 1; r.per_row[2] = 1; r.per_row[3] = 1;

        r.per_d1[0] = 1;
        r.per_d1[4] = 3;

        r.per_d2[1] = 1;
        r.per_d2[3] = 2;
        r.per_d2[5] = 1;

        r
    }

    // |*|*| | |    |3|3|3|3|
    // | | | | |    |2|3|3|2|
    // | | |*|*|    |3|3|3|3|
    // | | | | |    |1|2|2|3|
    pub fn example2() -> Puzzle {
        let mut r = Puzzle::new(4);
        r.in_col[0] = 0;
        r.in_col[1] = 0;
        r.in_col[2] = 2;
        r.in_col[3] = 2;
        r.per_row[0] = 2; r.per_row[1] = 0; r.per_row[2] = 2; r.per_row[3] = 0;

        r.per_d1[3] = 2;
        r.per_d1[4] = 2;

        r.per_d2[0] = 1;
        r.per_d2[1] = 1;
        r.per_d2[4] = 1;
        r.per_d2[5] = 1;

        r
    }
}

#[test]
fn test_attacks() {
    let e1 = Puzzle::example1();
    assert_eq!(e1.attacks(0, 0), 2, "Bad attacks!"); // top-left
    assert_eq!(e1.attacks(0, 1), 3, "Bad attacks!"); // below top-left
    assert_eq!(e1.attacks(0, 2), 2, "Bad attacks!");
    assert_eq!(e1.attacks(0, 3), 2, "Bad attacks!");

    assert_eq!(e1.attacks(1, 0), 3, "Bad attacks!"); // right of top-left
    assert_eq!(e1.attacks(1, 1), 2, "Bad attacks!");
    assert_eq!(e1.attacks(1, 2), 4, "Bad attacks!");
    assert_eq!(e1.attacks(1, 3), 2, "Bad attacks!");

    assert_eq!(e1.attacks(2, 0), 2, "Bad attacks!");
    assert_eq!(e1.attacks(2, 1), 4, "Bad attacks!");
    assert_eq!(e1.attacks(2, 2), 2, "Bad attacks!");
    assert_eq!(e1.attacks(2, 3), 3, "Bad attacks!");

    assert_eq!(e1.attacks(3, 0), 4, "Bad attacks!");
    assert_eq!(e1.attacks(3, 1), 2, "Bad attacks!");
    assert_eq!(e1.attacks(3, 2), 3, "Bad attacks!");
    assert_eq!(e1.attacks(3, 3), 2, "Bad attacks!");

    let e2 = Puzzle::example2();
    assert_eq!(e2.attacks(0, 0), 3, "Bad attacks!"); // top-left
    assert_eq!(e2.attacks(0, 1), 2, "Bad attacks!"); // below top-left
    assert_eq!(e2.attacks(0, 2), 3, "Bad attacks!");
    assert_eq!(e2.attacks(0, 3), 1, "Bad attacks!");

    assert_eq!(e2.attacks(1, 0), 3, "Bad attacks!"); // right of top-left
    assert_eq!(e2.attacks(1, 1), 3, "Bad attacks!");
    assert_eq!(e2.attacks(1, 2), 3, "Bad attacks!");
    assert_eq!(e2.attacks(1, 3), 2, "Bad attacks!");

    assert_eq!(e2.attacks(2, 0), 3, "Bad attacks!");
    assert_eq!(e2.attacks(2, 1), 3, "Bad attacks!");
    assert_eq!(e2.attacks(2, 2), 3, "Bad attacks!");
    assert_eq!(e2.attacks(2, 3), 2, "Bad attacks!");

    assert_eq!(e2.attacks(3, 0), 3, "Bad attacks!");
    assert_eq!(e2.attacks(3, 1), 2, "Bad attacks!");
    assert_eq!(e2.attacks(3, 2), 3, "Bad attacks!");
    assert_eq!(e2.attacks(3, 3), 3, "Bad attacks!");
}