const std = @import("std");
const aoc = @import("./aoc.zig");
const Allocator = std.mem.Allocator;
const Point = @import("./point.zig").Point(u32);
const ArrayList = std.ArrayList;

const Tile = enum {
    Empty,
    Galaxy,

    fn to_char(self: Tile) u8 {
        return switch (self) {
            .Empty => '.',
            .Galaxy => '#',
        };
    }

    fn from_char(char: u8) ?Tile {
        return switch (char) {
            '.' => .Empty,
            '#' => .Galaxy,
            else => null
        };
    }
};

const GalaxyMap = struct {
    allocator: Allocator,
    width: u32,
    height: u32,
    tiles: []Tile,

    fn deinit(self: GalaxyMap) void {
        self.allocator.free(self.tiles);
    }

    fn print(self: GalaxyMap) void {
        for (0..self.height) |y| {
            for (0..self.width) |x| {
                const idx = y * self.width + x;
                std.debug.print("{c}", .{self.tiles[idx].to_char()});
            }
            std.debug.print("\n", .{});
        }
    }
};

fn parse_input(allocator: Allocator, lines: []const []const u8) !GalaxyMap {
    var height: u32 = @intCast(lines.len);
    var width: u32 = @intCast(lines[0].len);
    var tile_map = try allocator.alloc(Tile, width * height);
    errdefer allocator.free(tile_map);

    for (0.., lines) |y, line| {
        for (0.., line) |x, char| {
            const idx = width * y + x;
            tile_map[idx] = Tile.from_char(char) orelse .Empty;
        }
    }

    return GalaxyMap{
        .allocator = allocator,
        .width = width,
        .height = height,
        .tiles = tile_map
    };
}

fn list_galaxies(allocator: Allocator, map: GalaxyMap) !ArrayList(Point) {
    var points = ArrayList(Point).init(allocator);
    errdefer points.deinit();

    for (0..map.height) |y| {
        for (0..map.width) |x| {
            const idx = y * map.width + x;
            if (map.tiles[idx] != .Galaxy) continue;

            try points.append(Point{
                .x = @intCast(x),
                .y = @intCast(y)
            });
        }
    }

    return points;
}

fn is_row_empty(map: GalaxyMap, y: u32) bool {
    for (0..map.width) |x| {
        const idx = map.width * y + x;
        if (map.tiles[idx] != .Empty) return false;
    }
    return true;
}

fn is_column_empty(map: GalaxyMap, x: u32) bool {
    for (0..map.height) |y| {
        const idx = map.width * y + x;
        if (map.tiles[idx] != .Empty) return false;
    }
    return true;
}

fn list_empty_rows(allocator: Allocator, map: GalaxyMap) !ArrayList(u32) {
    var rows = ArrayList(u32).init(allocator);
    errdefer rows.deinit();

    for (0..map.height) |y| {
        if (is_row_empty(map, @intCast(y))) {
            try rows.append(@intCast(y));
        }
    }

    return rows;
}

fn list_empty_columns(allocator: Allocator, map: GalaxyMap) !ArrayList(u32) {
    var columns = ArrayList(u32).init(allocator);
    errdefer columns.deinit();

    for (0..map.width) |x| {
        if (is_column_empty(map, @intCast(x))) {
            try columns.append(@intCast(x));
        }
    }

    return columns;
}

fn expand_points(allocator: Allocator, points: []Point, map: GalaxyMap, amount: u32) !void {
    const empty_rows = try list_empty_rows(allocator, map);
    defer empty_rows.deinit();

    const empty_columns = try list_empty_columns(allocator, map);
    defer empty_columns.deinit();

    for (points) |*p| {
        var expand_by: u32 = 0;
        for (empty_rows.items) |y| {
            if (y < p.y) {
                expand_by += 1;
            }
        }
        p.y += expand_by*amount;
    }

    for (points) |*p| {
        var expand_by: u32 = 0;
        for (empty_columns.items) |x| {
            if (x < p.x) {
                expand_by += 1;
            }
        }
        p.x += amount*expand_by;
    }
}

fn sum_min_distances(points: []Point) u64 {
    var result: u64 = 0;
    for (0..(points.len-1)) |i| {
        const p1 = points[i];
        for ((i+1)..points.len) |j| {
            const p2 = points[j];
            const dx = @as(i64, @intCast(p1.x)) - @as(i64, @intCast(p2.x));
            const dy = @as(i64, @intCast(p1.y)) - @as(i64, @intCast(p2.y));
            result += std.math.absCast(dx) + std.math.absCast(dy);
        }
    }
    return result;
}

pub fn part1(input: *aoc.Input) !aoc.Result {
    const allocator = input.allocator;
    const parsed = try parse_input(allocator, input.lines);
    defer parsed.deinit();

    var galaxies = try list_galaxies(allocator, parsed);
    defer galaxies.deinit();

    try expand_points(allocator, galaxies.items, parsed, 1);

    return .{ .uint = sum_min_distances(galaxies.items) };
}

pub fn part2(input: *aoc.Input) !aoc.Result {
    const allocator = input.allocator;
    const parsed = try parse_input(allocator, input.lines);
    defer parsed.deinit();

    var galaxies = try list_galaxies(allocator, parsed);
    defer galaxies.deinit();

    try expand_points(allocator, galaxies.items, parsed, 1000000 - 1);

    return .{ .uint = sum_min_distances(galaxies.items) };
}

const example_input = [_][]const u8{
    "...#......",
    ".......#..",
    "#.........",
    "..........",
    "......#...",
    ".#........",
    ".........#",
    "..........",
    ".......#..",
    "#...#....."
};

test "part 1 example" {
    try aoc.expectAnswerUInt(part1, 374, &example_input);
}

test "part 2 example" {
    try aoc.expectAnswerUInt(part2, 82000210, &example_input);
}