#include "aoc.h"
#include <stdlib.h>

typedef struct {
    i16 x, y;
    u8 letter;
} Antenna;

static int
cmp_antennas(const void *a, const void *b) {
    Antenna *aa = (Antenna *)a;
    Antenna *bb = (Antenna *)b;
    return aa->letter - bb->letter;
}

static i32
sum(const u8 *restrict values, isize len) {
    i32 sum = 0;
    for (isize i = 0; i < len; i++) {
        sum += values[i];
    }
    return sum;
}

int main(int argc, char **argv) {
    Arena *arena = make_arena(Megabytes(1));
    str input = read_file(arena, argv[1]);
    Grid grid = parse_grid(input, arena);

    struct {
        Antenna *data;
        isize len, cap;
    } antennas = {0};
    for (isize y = 0; y < grid.height; y++) {
        for (isize x = 0; x < grid.width; x++) {
            u8 letter;
            if ((letter = grid.grid[y * grid.width + x]) != '.') {
                *push(&antennas, arena) = (Antenna){x, y, letter};
            }
        }
    }

    // sort/group antennas by frequency (i.e. letter)
    qsort(antennas.data, antennas.len, sizeof(Antenna), cmp_antennas);

    u8 *antinodes = ARENA_ALLOC_ARRAY(arena, u8, grid.width * grid.height);

    isize group_start = 0;
    for (isize i = 0; i < antennas.len; i++) {
        if (antennas.data[i].letter != antennas.data[group_start].letter) {
            group_start = i;
        }
        for (isize j = group_start; j < i; j++) {
            ASSERT(antennas.data[i].letter == antennas.data[j].letter);
            i32 dx = antennas.data[i].x - antennas.data[j].x;
            i32 dy = antennas.data[i].y - antennas.data[j].y;

            // first antinode
            i32 antinode_x = antennas.data[i].x + dx;
            i32 antinode_y = antennas.data[i].y + dy;
            if (antinode_x >= 0 && antinode_y >= 0 && antinode_x < grid.width && antinode_y < grid.height) {
                antinodes[antinode_y * grid.width + antinode_x] = 1;
            }

            // second antinode
            antinode_x = antennas.data[j].x - dx;
            antinode_y = antennas.data[j].y - dy;
            if (antinode_x >= 0 && antinode_y >= 0 && antinode_x < grid.width && antinode_y < grid.height) {
                antinodes[antinode_y * grid.width + antinode_x] = 1;
            }
        }
    }

    i32 part_1 = sum(antinodes, grid.width * grid.height);
    printf("%d\n", part_1);

    // Part 2
    group_start = 0;
    for (isize i = 0; i < antennas.len; i++) {
        if (antennas.data[i].letter != antennas.data[group_start].letter) {
            group_start = i;
        }
        for (isize j = group_start; j < i; j++) {
            ASSERT(antennas.data[i].letter == antennas.data[j].letter);
            i32 dx = antennas.data[i].x - antennas.data[j].x;
            i32 dy = antennas.data[i].y - antennas.data[j].y;

            i32 x = antennas.data[i].x;
            i32 y = antennas.data[i].y;

            while (x >= 0 && y >= 0 && x < grid.width && y < grid.height) {
                antinodes[y * grid.width + x] = 1;
                x += dx;
                y += dy;
            }

            x = antennas.data[i].x;
            y = antennas.data[i].y;
            while (x >= 0 && y >= 0 && x < grid.width && y < grid.height) {
                antinodes[y * grid.width + x] = 1;
                x -= dx;
                y -= dy;
            }
        }
    }

    i32 part_2 = sum(antinodes, grid.width * grid.height);
    printf("%d\n", part_2);
}