#include "aoc.h"

typedef struct {
    i32 x, y;
    i32 vx, vy;
} Robot;

static void
parse_pair(str s, i32 *x, i32 *y, Arena arena) {
    s = str_sub(s, 2, s.len);
    str lhs, rhs;
    str_split(s, STR(","), &lhs, &rhs);
    *x = (i32) parse_i64(lhs, arena);
    *y = (i32) parse_i64(rhs, arena);
}

static Robot
parse_robot(str line, Arena arena) {
    str p, v;
    str_split(line, STR(" "), &p, &v);
    Robot robot;
    parse_pair(p, &robot.x, &robot.y, arena);
    parse_pair(v, &robot.vx, &robot.vy, arena);
    return robot;
}

static void
dump_grid(Arena arena, Robot *robots, isize count, int width, int height) {
    u8 *grid = ARENA_ALLOC_ARRAY(&arena, u8, width * height);

    for (isize i = 0; i < count; i++) {
        Robot *r = &robots[i];
        grid[r->y * width + r->x]++;
    }

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            u8 count = grid[y * width + x];
            if (count == 0) {
                putchar('.');
            } else {
                putchar('0' + count);
            }
        }
        putchar('\n');
    }
}

static void
step(Robot *robots, isize count, int width, int height) {
    for (isize i = 0; i < count; i++) {
        Robot *r = &robots[i];
        r->x = (r->x + r->vx + width) % width;
        r->y = (r->y + r->vy + height) % height;
    }
}

static bool
has_tree(u8 *grid, int width, int height) {
    for (int x = 0; x < width; x++) {
        i32 consecutive = 0;
        for (int y = 1; y < height; y++) {
            if (grid[y * width + x] && grid[(y - 1) * width + x]) {
                consecutive++;
            } else {
                consecutive = 0;
            }
            if (consecutive >= width / 4) {
                return true;
            }
        }
    }
    return false;
}

static f32
compute_entropy(u8 *grid, int width, int height) {
    i32 counts[256] = {0};
    for (int i = 0; i < width * height; i++) {
        counts[grid[i]]++;
    }
    f32 entropy = 0.0f;
    for (int i = 0; i < 256; i++) {
        if (counts[i]) {
            f32 prob = (f32) counts[i] / (f32) (width * height);
            entropy -= prob * log2f(prob);
        }
    }
    return entropy;
}

int main(int argc, char **argv) {
    Arena *arena = make_arena(Megabytes(1));
    Tokens lines = read_lines(arena, argv[1]);

    int count = lines.len;
    bool is_example = true;

    Robot *robots = ARENA_ALLOC_ARRAY(arena, Robot, lines.len);
    for (isize i = 0; i < lines.len; i++) {
        robots[i] = parse_robot(str_trim(lines.tokens[i]), *arena);
        if (robots[i].x >= 11) {
            is_example = false;
        }
    }
    int width = is_example ? 11 : 101;
    int height = is_example ? 7 : 103;
    ASSERT(width % 2 == 1);
    ASSERT(height % 2 == 1);

    /* dump_grid(*arena, robots, count, width, height); */
    /* puts(""); */

    i64 second = 0;
    for (int i = 0; i < 100; i++) {
        step(robots, count, width, height);
        second++;
    }

    int cx = width / 2;
    int cy = height / 2;

    i64 counts[4] = {0};
    for (isize i = 0; i < count; i++) {
        Robot *r = &robots[i];
        if (r->x != width / 2 && r->y != height / 2) {
            int qx = r->x < cx ? 0 : 1;
            int qy = r->y < cy ? 0 : 1;
            counts[qy * 2 + qx]++;
        }
    }

    i64 part_1  = counts[0] * counts[1] * counts[2] * counts[3];
    printf("%ld\n", part_1);

    // Part 2
    u8 *grid = ARENA_ALLOC_ARRAY(arena, u8, width * height);
    f32 best = 1.0f;
    i32 best_second = 0;
    i64 part_2 = 0;
    for (;;) {
        step(robots, count, width, height);
        second++;
        memset(grid, 0, width * height);
        for (isize i = 0; i < count; i++) {
            Robot *r = &robots[i];
            grid[r->y * width + r->x]++;
        }
#if 0
        f32 entropy = compute_entropy(grid, width, height);
        if (entropy < best) {
            best = entropy;
            best_second = second;
            printf("%ld %f\n", second, entropy);
            dump_grid(*arena, robots, count, width, height);
            puts("--------------------");
        }
#else
        if (has_tree(grid, width, height)) {
            printf("Tree detected on second %ld\n", second);
            part_2 = second;
            break;
        }
#endif
    }
    dump_grid(*arena, robots, count, width, height);
    printf("%ld\n", part_2);
}