#include "aoc.h"

typedef union {
    i8 values[4];
    u32 value;
} Sequence;

typedef struct SeqEntry {
    i16 bananas;
} SeqEntry;

#define SEQ_MAP_SIZE (19 * 19 * 19 * 19)
typedef struct SeqMap {
    union {
        SeqEntry entries[19][19][19][19];
        SeqEntry flat[SEQ_MAP_SIZE];
    };
    u64 present[(SEQ_MAP_SIZE + 63) / 64];
} SeqMap;

static inline u64
prune(u64 secret) {
    secret %= 16777216;
    return secret;
}

static inline u64
derive(u64 secret) {
    secret ^= (secret * 64);
    secret = prune(secret);

    secret ^= (secret / 32);
    secret = prune(secret);

    secret ^= (secret * 2048);
    secret = prune(secret);
    return secret;
}

static inline u64
seq_map_index(Sequence seq) {
    return (seq.values[0] + 9) * 19 * 19 * 19 +
           (seq.values[1] + 9) * 19 * 19 +
           (seq.values[2] + 9) * 19 +
           (seq.values[3] + 9);
}

static inline void
seq_map_insert_if_not_present(SeqMap *map, Sequence seq, i32 bananas) {
    u64 index = seq_map_index(seq);
    SeqEntry *entry = map->flat + index;
    u64 bit = 1ULL << (index % 64);
    u8 present_index = index / 64;
    bool present = (map->present[present_index] & bit) != 0;
    if (!present) {
        map->present[present_index] |= bit;
        entry->bananas = bananas;
    }
}

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

    i8 *differences = ARENA_ALLOC_ARRAY(arena, i8, 2000);

    u64 part_1 = 0;
    SeqMap *buyer_map = ARENA_ALLOC(arena, SeqMap);
    SeqMap *map = ARENA_ALLOC(arena, SeqMap);
    for (isize i = 0; i < input.len; i++) {
        u64 secret = (u64) parse_i64(str_trim(input.tokens[i]), *arena);

        /* SeqMap buyer_map = {0}; */
        memset(buyer_map, 0, sizeof(SeqMap));
        i32 current_price = secret % 10;

        Sequence diff = {0};
        for (i32 j = 0; j < 2000; j++) {
            secret = derive(secret);
            i32 new_price = secret % 10;
            i32 diff = new_price - current_price;
            differences[j] = diff;

            if (j >= 3) {
                Sequence seq = { .values = { differences[j - 3], differences[j - 2], differences[j - 1], differences[j] } };
                i32 bananas = new_price;
                seq_map_insert_if_not_present(buyer_map, seq, bananas);
            }
            current_price = new_price;
        }
        part_1 += secret;

        // Update global map with bananas from this buyer
        SeqEntry *restrict from = buyer_map->flat;
        SeqEntry *restrict to = map->flat;
        for (isize i = 0; i < SEQ_MAP_SIZE; i++) {
            to[i].bananas += from[i].bananas;
        }
    }
    printf("%lu\n", part_1);

    // Part 2: Find the maximum number of bananas
    i32 part_2 = 0;
    SeqEntry *restrict flat = map->flat;
    for (isize i = 0; i < SEQ_MAP_SIZE; i++) {
        part_2 = MAX(flat[i].bananas, part_2);
    }
    printf("%d\n", part_2);
}