const std = @import("std");

const path = "data/day14/input.txt";

const RetType = u43;
const Str = []const u8;

const Rules = std.StringHashMap(u8);
const Polymer = struct {
    template: Str,
    rules: Rules,
};

fn parseInput(a: std.mem.Allocator) anyerror!Polymer {
    const input = @embedFile(path);
    var lines = std.mem.split(u8, input, "\n");

    var ret = Polymer{
        .template = undefined,
        .rules = Rules.init(a),
    };

    ret.template = lines.next().?;
    _ = lines.next(); // skip empty line

    // parse rules
    while (lines.next()) |rule_line| {
        if (rule_line.len == 0) break;
        var rule = std.mem.tokenize(u8, rule_line, " -> ");
        try ret.rules.put(rule.next().?, rule.next().?[0]);
    }

    return ret;
}

pub fn second(allocator: ?std.mem.Allocator) anyerror!RetType {
    var p = try parseInput(allocator.?);
    defer p.rules.deinit();

    var pairs = std.StringHashMap(RetType).init(allocator.?);
    defer pairs.deinit();

    // initialize pairs
    for (p.template) |_, idx| {
        if (idx == p.template.len - 1) break;
        const ps = p.template[idx .. idx + 2];
        var item = try pairs.getOrPut(ps);
        if (item.found_existing) item.value_ptr.* += 1 else item.value_ptr.* = 1;
    }

    // arena for allocPrint()
    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    var step: usize = 0;
    while (step < 40) : (step += 1) {
        var next_pairs = std.StringHashMap(RetType).init(allocator.?);

        var it = pairs.iterator();
        while (it.next()) |item| {
            const next_left = try std.fmt.allocPrint(arena.allocator(), "{c}{c}", .{ item.key_ptr.*[0], p.rules.get(item.key_ptr.*).? });
            var next_item = try next_pairs.getOrPut(next_left);
            if (next_item.found_existing) next_item.value_ptr.* += item.value_ptr.* else next_item.value_ptr.* = item.value_ptr.*;

            const next_right = try std.fmt.allocPrint(arena.allocator(), "{c}{c}", .{ p.rules.get(item.key_ptr.*).?, item.key_ptr.*[1] });
            next_item = try next_pairs.getOrPut(next_right);
            if (next_item.found_existing) next_item.value_ptr.* += item.value_ptr.* else next_item.value_ptr.* = item.value_ptr.*;
        }

        pairs.deinit();
        pairs = next_pairs;
    }

    var chars = std.AutoHashMap(u8, RetType).init(allocator.?);
    defer chars.deinit();

    // put in the last char as the counter below skips it
    try chars.put(p.template[p.template.len - 1], 1);

    var pit = pairs.iterator();
    while (pit.next()) |pair| {
        const char = pair.key_ptr.*[0];
        const item = try chars.getOrPut(char);
        if (item.found_existing) item.value_ptr.* += pair.value_ptr.* else item.value_ptr.* = pair.value_ptr.*;
    }

    var min = ~@as(RetType, 0);
    var max: RetType = 0;
    var cit = chars.valueIterator();
    while (cit.next()) |c| {
        if (c.* < min) {
            min = c.*;
        }
        if (c.* > max) {
            max = c.*;
        }
    }
    return max - min;
}

pub fn main() anyerror!void {
    var buf: [500_000]u8 = undefined;
    var fba = std.heap.FixedBufferAllocator.init(&buf);

    var timer = try std.time.Timer.start();
    const ret = try second(fba.allocator());
    const t = timer.lap() / 1000;

    try std.testing.expectEqual(@as(RetType, 7477815755570), ret);

    std.debug.print("Day 14b result: {d} \ttime: {d}us\n", .{ ret, t });
}

test "day14b" {
    try std.testing.expectEqual(@as(RetType, 7477815755570), try second(std.testing.allocator));
}