const std = @import("std");

pub fn main() anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    var gpa = &arena.allocator;

    var f = try std.fs.cwd().openFile("input", .{});
    var contents = try f.readToEndAlloc(gpa, std.math.maxInt(u32));

    var seatmap = try SeatMap.init(gpa);
    var it = std.mem.tokenize(contents, "\n");
    while (it.next()) |line| {
        try seatmap.add_line(line);
    }

    var changed : usize = 0;
    changed = try seatmap.process_round();
    while (changed != 0) {
        //changed = try seatmap.process_round();
        changed = try seatmap.process_round_part2();
    }
    std.log.info("{} occupied seats", .{seatmap.count_occupied()});
}

const NeighbourState = struct {
    count: usize = 0,
    occupied: usize = 0,
    free: usize = 0,
};

const SeatMap = struct {
    allocator: *std.mem.Allocator,
    data: std.ArrayList(u8),
    width: usize = 0,
    height: usize = 0,

    pub fn init(a: *std.mem.Allocator) !*SeatMap {
        var self = try a.create(SeatMap);
        errdefer a.destroy(self);

        // if we use self.* = .{}; then we don't forget
        // obligatory fields
        // also, when using self.* = .{} the fields with
        // defaults are properly initialized!!! it is
        // an error to not use it.
        self.* = .{
            .allocator = a,
            .data = std.ArrayList(u8).init(a),
        };
        return self;
    }

    pub fn deinit(self: *SeatMap) void {
        self.data.deinit();
        self.allocator.destroy(self);
    }

    pub fn count_occupied(self: *SeatMap) u64 {
        var c: u64 = 0;
        for (self.data.items) |i| {
            if (i == '#') {
                c += 1;
            }
        }
        return c;
    }

    pub fn print(self: *SeatMap) !void {
        var stdout = std.io.getStdOut();
        var i : usize = 0;
        while (i < self.height) : (i += 1) {
            _ = try stdout.write(self.get_row(i));
            _ = try stdout.write("\n");
        }
    }

    pub fn get_row(self: *SeatMap, row: usize) []u8 {
        return self.data.items[self.width*row..(self.width*row)+self.width];
    }

    pub fn add_line(self: *SeatMap, line: []const u8) !void {
        if (self.width != 0) {
            std.debug.assert(line.len == self.width);
        }
        else {
            self.width = line.len;
        }
        try self.data.appendSlice(line[0..]);
        self.height += 1;
    }

    // returns the number of seats changed
    pub fn process_round(self: *SeatMap) !u64 {
        // First, since operations happen concurrently,
        // we mark all of our changes to a temporary buffer
        var buffer = try self.allocator.alloc(u8, self.data.items.len);
        defer self.allocator.free(buffer);

        // Iterate above (row, col) pairs
        var row : usize = 0;
        var n_changed : u64 = 0;
        while (row < self.height) : (row += 1) {
            var col : usize = 0;
            while (col < self.width) : (col += 1) {
                if (self.index_of(row, col)) |i| {
                    buffer[i] = self.get_new_seat_state(row, col, i);
                    if (buffer[i] != self.data.items[i]) {
                        n_changed += 1;
                    }
                }
            }
        }

        // Copy buffer to self.data.items
        std.mem.copy(u8, self.data.items, buffer);
        return n_changed;
    }

    pub fn get_new_seat_state(self: *SeatMap, row: usize, col: usize, index: usize) u8 {
        // Floor remains unchanged
        var c = self.data.items[index];
        var new_c : u8 = undefined;
        if (c == '.') {
            new_c = '.';
        }

        // this gives us a struct w/ .count, .occupied, .free
        var neighbour_state = self.get_neighbour_state(row, col);
        if (c == 'L') {
            if (neighbour_state.occupied == 0) {
                new_c = '#';
            }
            else {
                new_c = c;
            }
        }
        if (c == '#') {
            if (neighbour_state.occupied >= 4) {
                new_c = 'L';
            }
            else {
                new_c = c;
            }
        }
        //std.log.debug("({},{}) was '{c}' and will be '{c}'. {} neighbors, {} occupied",
        //             .{row, col, c, new_c, neighbour_state.count, neighbour_state.occupied});
        return new_c;
    }

    pub fn get_neighbour_state(self: *SeatMap, row: usize, col: usize) NeighbourState {
        var r: usize = switch(row) {
            0 => row,
            else => row - 1,
        };
        var count: usize = 0;
        var free: usize = 0;
        var occupied: usize = 0;
        while (r <= row+1) : (r += 1) {
            var c: usize = switch(col) {
                0 => col,
                else => col - 1,
            };
            while (c <= col+1) : (c += 1) {
                if (row == r and col == c) {
                    continue;
                }
                if (self.index_of(r, c)) |i| {
                    if (self.data.items[i] == '.') {
                        // Floor spots don't contribute to neighbour,
                        // free, or occupied counts.
                        continue;
                    }
                    count += 1;
                    if (self.data.items[i] == '#') {
                        occupied += 1;
                    }
                    else {
                        free += 1;
                    }
                }
            }
        }
        return NeighbourState {
            .count = count,
            .free = free,
            .occupied = occupied,
        };
    }

    // returns the number of seats changed
    pub fn process_round_part2(self: *SeatMap) !u64 {
        // First, since operations happen concurrently,
        // we mark all of our changes to a temporary buffer
        var buffer = try self.allocator.alloc(u8, self.data.items.len);
        defer self.allocator.free(buffer);

        // Iterate above (row, col) pairs
        var row : usize = 0;
        var n_changed : u64 = 0;
        while (row < self.height) : (row += 1) {
            var col : usize = 0;
            while (col < self.width) : (col += 1) {
                if (self.index_of(row, col)) |i| {
                    buffer[i] = self.get_new_seat_state_round2(row, col, i);
                    if (buffer[i] != self.data.items[i]) {
                        n_changed += 1;
                    }
                }
            }
        }

        // Copy buffer to self.data.items
        std.mem.copy(u8, self.data.items, buffer);
        return n_changed;
    }

    pub fn get_new_seat_state_round2(self: *SeatMap, row: usize, col: usize, index: usize) u8 {
        // Floor remains unchanged
        var c = self.data.items[index];
        var new_c : u8 = undefined;
        if (c == '.') {
            new_c = '.';
        }

        // this gives us a struct w/ .count, .occupied, .free
        var neighbour_state = self.get_neighbour_state_round2(row, col);
        if (c == 'L') {
            if (neighbour_state.occupied == 0) {
                new_c = '#';
            }
            else {
                new_c = c;
            }
        }
        if (c == '#') {
            if (neighbour_state.occupied >= 5) {
                new_c = 'L';
            }
            else {
                new_c = c;
            }
        }
        //std.log.debug("({},{}) was '{c}' and will be '{c}'. {} neighbors, {} occupied",
        //             .{row, col, c, new_c, neighbour_state.count, neighbour_state.occupied});
        return new_c;
    }

    pub fn get_neighbour_state_round2(self: *SeatMap, row: usize, col: usize) NeighbourState {
        var count: usize = 0;
        var free: usize = 0;
        var occupied: usize = 0;
        var direction_vectors = [_][2]i64 {
            [_]i64{-1, -1}, // up and left,
            [_]i64{-1, 0}, // left
            [_]i64{-1, 1}, // down and left,
            [_]i64{0, 1}, // down
            [_]i64{1, 1}, // right and down
            [_]i64{1, 0}, // right
            [_]i64{1, -1}, // right and up,
            [_]i64{0, -1}, // up
        };
        var row_signed = @intCast(i64, row);
        var col_signed = @intCast(i64, col);
        for (direction_vectors) |dv| {
            var space_count : i64 = 0;
            var has_neighbour: bool = false;
            var last_neighbour_was_occupied = false;
            while (true) {
                space_count += 1;
                var r : i64 = row_signed + (dv[0]*space_count);
                var c : i64 = col_signed + (dv[1]*space_count);
                if (r < 0 or c < 0) {
                    // invalid spot with negative indices
                    break;
                }
                if(self.index_of(@intCast(usize, r), @intCast(usize, c))) |index| {
                    if (self.data.items[index] == '.') {
                        // We see through floor spaces
                        continue;
                    }
                    else {
                        has_neighbour = true;
                        if (self.data.items[index] == '#') {
                            last_neighbour_was_occupied = true;
                        }
                        break;
                    }
                }
                else {
                    // The index was not valid, so we're out of bounds in some way
                    break;
                }
            }
            if (has_neighbour) {
                count += 1;
                if (last_neighbour_was_occupied) {
                    occupied += 1;
                }
                else {
                    free += 1;
                }
            }
        }
        return NeighbourState {
            .count = count,
            .free = free,
            .occupied = occupied,
        };
    }

    pub fn index_of(self: *SeatMap, row: usize, col: usize) ?usize {
        if (col >= self.width) {
            return null;
        }
        if (row >= self.height) {
            return null;
        }
        var i = (row*self.width) + col;
        if (i < 0 or i >= self.data.items.len) {
            return null;
        }
        return i;
    }
};

test "part_one" {
    var d = [_][]const u8 {
        "L.LL.LL.LL",
        "LLLLLLL.LL",
        "L.L.L..L..",
        "LLLL.LL.LL",
        "L.LL.LL.LL",
        "L.LLLLL.LL",
        "..L.L.....",
        "LLLLLLLLLL",
        "L.LLLLLL.L",
        "L.LLLLL.LL",
    };

    var print = true;
    var stdout = std.io.getStdOut();

    var sm = try SeatMap.init(std.testing.allocator);
    defer sm.deinit();
    for (d) |l| {
        try sm.add_line(l[0..]);
    }

    var states = [_][]const u8 {
        // 0
        "L.LL.LL.LLLLLLLLL.LLL.L.L..L..LLLL.LL.LLL.LL.LL.LLL.LLLLL.LL..L.L.....LLLLLLLLLLL.LLLLLL.LL.LLLLL.LL",
        // 1
        "#.##.##.#########.###.#.#..#..####.##.###.##.##.###.#####.##..#.#.....###########.######.##.#####.##",
        // 2
        "#.LL.L#.###LLLLLL.L#L.L.L..L..#LLL.LL.L##.LL.LL.LL#.LLLL#.##..L.L.....#LLLLLLLL##.LLLLLL.L#.#LLLL.##",
        // 3
        "#.##.L#.###L###LL.L#L.#.#..#..#L##.##.L##.##.LL.LL#.###L#.##..#.#.....#L######L##.LL###L.L#.#L###.##",
        // 4
        "#.#L.L#.###LLL#LL.L#L.L.L..#..#LLL.##.L##.LL.LL.LL#.LL#L#.##..L.L.....#L#LLLL#L##.LLLLLL.L#.#L#L#.##",
        // 5
        "#.#L.L#.###LLL#LL.L#L.#.L..#..#L##.##.L##.#L.LL.LL#.#L#L#.##..L.L.....#L#L##L#L##.LLLLLL.L#.#L#L#.##"
    };
    // Confirm that initial map is correct
    std.testing.expect(std.mem.eql(u8, sm.data.items, states[0]));

    if (print) {
        _ = try stdout.write("\nInitial map:\n");
        try sm.print();
        _ = try stdout.write("\n");
    }

    var round: u64 = 0;
    var changed = try sm.process_round();
    while (changed != 0) {
        round += 1;
        var matches = std.mem.eql(u8, sm.data.items, states[round]);
        //std.log.warn("Round {} matches expected state: {}", .{round, matches});
        //std.log.warn("len {}: {}", .{sm.data.items.len, sm.data.items});
        //std.log.warn("len {}: {}", .{states[round].len, states[round]});
        //std.log.warn("index of first diff {}",
        //             .{std.mem.indexOfDiff(u8, sm.data.items, states[round])});
        std.testing.expect(matches);
        if (print) {
            var title = try std.fmt.allocPrint(std.testing.allocator, "\nRound {}:\n",
                                           .{round});
            _ = try stdout.write(title);
            std.testing.allocator.free(title);
            try sm.print();
            _ = try stdout.write("\n");
        }
        changed = try sm.process_round();
    }
    if (print) {
        _ = try stdout.write("Final state: \n");
        try sm.print();
        _ = try stdout.write("\n");
    }
    std.testing.expect(std.mem.eql(u8, sm.data.items, states[round]));

    var n_occupied = sm.count_occupied();
    std.testing.expectEqual(n_occupied, 37);
}

test "part two" {
    var d = [_][]const u8 {
        "L.LL.LL.LL",
        "LLLLLLL.LL",
        "L.L.L..L..",
        "LLLL.LL.LL",
        "L.LL.LL.LL",
        "L.LLLLL.LL",
        "..L.L.....",
        "LLLLLLLLLL",
        "L.LLLLLL.L",
        "L.LLLLL.LL",
    };

    var print = true;
    var stdout = std.io.getStdOut();

    var sm = try SeatMap.init(std.testing.allocator);
    defer sm.deinit();
    for (d) |l| {
        try sm.add_line(l[0..]);
    }

    var states = [_][]const u8 {
        // 0
        "L.LL.LL.LLLLLLLLL.LLL.L.L..L..LLLL.LL.LLL.LL.LL.LLL.LLLLL.LL..L.L.....LLLLLLLLLLL.LLLLLL.LL.LLLLL.LL",
        // 1
        "#.##.##.#########.###.#.#..#..####.##.###.##.##.###.#####.##..#.#.....###########.######.##.#####.##",
        // 2
        "#.LL.LL.L##LLLLLL.LLL.L.L..L..LLLL.LL.LLL.LL.LL.LLL.LLLLL.LL..L.L.....LLLLLLLLL##.LLLLLL.L#.LLLLL.L#",
        // 3
        "#.L#.##.L##L#####.LLL.#.#..#..##L#.##.###.##.#L.###.#####.#L..#.#.....LLL####LL##.L#####.L#.L####.L#",
        // 4
        "#.L#.L#.L##LLLLLL.LLL.L.L..#..##LL.LL.L#L.LL.LL.L##.LLLLL.LL..L.L.....LLLLLLLLL##.LLLLL#.L#.L#LL#.L#",
        // 5
        "#.L#.L#.L##LLLLLL.LLL.L.L..#..##L#.#L.L#L.L#.#L.L##.L####.LL..#.#.....LLL###LLL##.LLLLL#.L#.L#LL#.L#",
        // 6
        "#.L#.L#.L##LLLLLL.LLL.L.L..#..##L#.#L.L#L.L#.LL.L##.LLLL#.LL..#.L.....LLL###LLL##.LLLLL#.L#.L#LL#.L#",
    };
    // Confirm that initial map is correct
    std.testing.expect(std.mem.eql(u8, sm.data.items, states[0]));

    if (print) {
        _ = try stdout.write("\nInitial map:\n");
        try sm.print();
        _ = try stdout.write("\n");
    }

    var round: u64 = 0;
    var changed = try sm.process_round();
    while (changed != 0) {
        round += 1;
        var matches = std.mem.eql(u8, sm.data.items, states[round]);
        //std.log.warn("Round {} matches expected state: {}", .{round, matches});
        //std.log.warn("len {}: {}", .{sm.data.items.len, sm.data.items});
        //std.log.warn("len {}: {}", .{states[round].len, states[round]});
        //std.log.warn("index of first diff {}",
        //             .{std.mem.indexOfDiff(u8, sm.data.items, states[round])});
        std.testing.expect(matches);
        if (print) {
            var title = try std.fmt.allocPrint(std.testing.allocator, "\nRound {}:\n",
                                           .{round});
            _ = try stdout.write(title);
            std.testing.allocator.free(title);
            try sm.print();
            _ = try stdout.write("\n");
        }
        changed = try sm.process_round_part2();
    }
    if (print) {
        _ = try stdout.write("Final state: \n");
        try sm.print();
        _ = try stdout.write("\n");
    }
    std.testing.expect(std.mem.eql(u8, sm.data.items, states[round]));

    var n_occupied = sm.count_occupied();
    std.testing.expectEqual(n_occupied, 26);
}