const std = @import("std");

pub fn main() anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    var gpa = &arena.allocator;

    var f = try std.fs.cwd().openFile("input", .{});
    var contents = try f.readToEndAlloc(gpa, std.math.maxInt(u32));
    var it = std.mem.tokenize(contents, "\n");
    var joltages = std.ArrayList(u32).init(gpa);
    defer joltages.deinit();
    while (it.next()) |line| {
        try joltages.append(atoi(line));
    }
    var chain = try JoltageChain.construct_chain(gpa, joltages.items[0..]);
    defer chain.deinit();
    std.log.info("Product of one_diffs and three_diffs of joltage chain: {}",
                 .{chain.part_one()});

    var permutations = try chain.calculate_permutations();
    std.log.info("{} permutations", .{permutations});
}

const Tree = struct {
    children : std.ArrayList(*Tree),
    allocator: *std.mem.Allocator,
    value : u32 = 0,
    dirty : bool = true,
    perms : u64 = 0,

    fn init(allocator: *std.mem.Allocator, value: u32) !*Tree {
        var self = try allocator.create(Tree);
        errdefer allocator.destroy(self);
        self.allocator = allocator;
        self.value = value;
        self.children = std.ArrayList(*Tree).init(allocator);
        return self;
    }

    fn deinit(self: *Tree) void {
        self.children.deinit();
        self.allocator.destroy(self);
    }

    fn add_child(self: *Tree, child: *Tree) !void {
        try self.children.append(child);
        self.dirty = true;
    }

    fn find_child_by_value(self: *Tree, value: u32) ?*Tree {
        if (self.value == value) {
            return self;
        }
        for (self.children.items) |c| {
            var t = c.find_child_by_value(value);
            if (t) |_t| {
                return _t;
            }
        }
        return null;
    }

    // caching the calculation of the perm count makes it possible to
    // finish reasonbly quickly when revisiting nodes that have already
    // been seen
    fn perm_count(self: *Tree) u64 {
        if (!self.dirty) {
           return self.perms;
        }
        var count : u64 = 0;
        if (self.children.items.len == 0) {
            count = 0; // No permutation, since have can only go one way
        }
        else if (self.children.items.len == 1) {
            // We only have ourself
            count = std.math.max(1, self.children.items[0].perm_count());
        }
        else {
            for(self.children.items) |c| {
                count += std.math.max(1, c.perm_count());
            }
        }
        self.perms = count;
        self.dirty = false;
        return count;
    }

    fn print_tree(self: *Tree) void {
        std.log.warn("{} children:\n", .{self.value});
        for(self.children.items) |c| {
            c.print_tree();
        }
    }
};

const JoltageChain = struct {
    allocator: *std.mem.Allocator,
    joltages : std.ArrayList(u32),

    pub fn construct_chain(allocator: *std.mem.Allocator, j: []u32) !*JoltageChain {
        var self = try allocator.create(JoltageChain);
        errdefer allocator.destroy(self);

        self.joltages = std.ArrayList(u32).init(allocator);
        self.allocator = allocator;

        // This has a side effect of modifying j...
        std.sort.insertionSort(u32, j[0..], {}, comptime std.sort.asc(u32));
        try self.joltages.append(0); // Always start at zero
        var last_joltage : u32 = 0;
        for (j) |x| {
            if ((x-last_joltage)<4) {
                try self.joltages.append(x);
                last_joltage = x;
            }
            else {
                break;
            }
        }
        try self.joltages.append(last_joltage + 3);
        return self;
    }

    pub fn deinit(self: *JoltageChain) void {
        self.joltages.deinit();
        self.allocator.destroy(self);
    }

    pub fn part_one(self: *JoltageChain) u32 {
        var one_diffs : u32 = 0;
        var three_diffs : u32 = 0;
        for (self.joltages.items) |v, k| {
            if (k == 0) {
                continue;
            }
            var delta : u32 = v - self.joltages.items[k-1];
            if (delta == 1) {
                one_diffs += 1;
            }
            if (delta == 3) {
                three_diffs += 1;
            }
        }
        std.log.debug("Found {} one diffs, and {} three diffs",
                      .{one_diffs, three_diffs});
        return one_diffs * three_diffs;
    }

    // this is a garbage way of doing it and just crushes CPU/memory
    // it does work for small datasets though
    pub fn calculate_permutations(self: *JoltageChain) !u32 {
        var tree = try Tree.init(self.allocator, self.joltages.items[0]);
        var current_node : *Tree = undefined;
        var count : u32 = 0;

        // Finding in the tree seems v. slow, will use a hash map for lookup
        // it is also much easier to try and deinit
        var map = std.hash_map.AutoHashMap(u32, *Tree).init(self.allocator);
        try map.ensureCapacity(@intCast(u32, self.joltages.items.len));
        map.putAssumeCapacityNoClobber(self.joltages.items[0], tree);
        for (self.joltages.items) |v, k| {
            var _t = map.get(v);
            if (_t) |t| {
                current_node = t;
            }
            else {
                unreachable;
            }
            //std.log.warn("Current node: {}", .{current_node});
            var n : usize = k+1;
            while (n < self.joltages.items.len) : (n += 1) {
                //std.log.warn("{}", .{n});
                if ((self.joltages.items[n] - v) <= 3) {
                    var value = self.joltages.items[n];
                    _t = map.get(value);
                    if (_t) |t| {
                        try current_node.add_child(t);
                        //std.log.warn("Added existing child {} to {}", .{value, v});
                    }
                    else {
                        var t = try Tree.init(current_node.allocator, value);
                        try current_node.add_child(t);
                        map.putAssumeCapacityNoClobber(value, t);
                        //std.log.warn("Added new child {} to {}", .{value, v});
                    }
                }
            }
        }
        std.log.warn("tree count: {}", .{tree.perm_count()});
        // Need to do some cleanup
        var it = map.iterator();
        while (it.next()) |kv| {
            kv.value.deinit();
            // don't remove from map, since it may modifying the underlying structure? (source?)
            // and we're about to deinit it anyway
        }
        map.deinit();
        return count;
    }
};

test "joltage_chain" {
    var joltages = [_]u32 {
        16,
        10,
        15,
        5,
        1,
        11,
        7,
        19,
        6,
        12,
        4,
    };
    var chain = try JoltageChain.construct_chain(std.testing.allocator, joltages[0..]);
    var diff = chain.part_one();
    std.testing.expectEqual(diff, 7*5);
    var perm = chain.calculate_permutations();

    chain.deinit();
}

test "joltage_chain" {
    var joltages = [_]u32 {
        28,
        33,
        18,
        42,
        31,
        14,
        46,
        20,
        48,
        47,
        24,
        23,
        49,
        45,
        19,
        38,
        39,
        11,
        1,
        32,
        25,
        35,
        8,
        17,
        7,
        9,
        4,
        2,
        34,
        10,
        3,
    };
    var chain = try JoltageChain.construct_chain(std.testing.allocator, joltages[0..]);
    var diff = chain.part_one();
    std.testing.expectEqual(diff, 220);
    var perm = chain.calculate_permutations();

    chain.deinit();
}

fn atoi(a: []const u8) u32 {
    var i : u32 = 0;
    var mul : u32 = 1;
    var start : usize = 0;
    if (a[0] == '-' or a[0] == '+') {
        start = 1;
        if (a[0] == '-') {
            //mul *= -1;
            unreachable;
        }
    }
    for(a[start..]) |v, k| {
        if (! std.ascii.isDigit(v)) {
            std.log.warn("Byte {x} is not a digit", .{v});
            continue;
        }
        // 48 is '0' in ascii
        std.debug.assert(v >= 48 and v < 58);
        i += @as(u32, @as(u32, (v - 48)) * std.math.pow(u32, 10, @intCast(u32, a.len - k - 1 - start)));
    }
    //std.log.debug("0x{x} --> {}", .{a, i});
    return i * mul;
}

test "atoi_regular" {
    var i = atoi("1234");
    std.testing.expectEqual(i, 1234);
}

test "atoi_pos" {
    var i = atoi("+1234");
    std.testing.expectEqual(i, 1234);
}