const std = @import("std");

pub fn main() !void {
    try std.fmt.format(std.io.getStdOut().writer(), "day 9 pt1: {}\n", .{try solve_pt1(std.heap.page_allocator, puzzle_input)});
    try std.fmt.format(std.io.getStdOut().writer(), "day 9 pt2: {}\n", .{try solve_pt2(std.heap.page_allocator, puzzle_input)});
}

fn solve_pt1(a: std.mem.Allocator, input: []const u8) !i64 {
    var sum: i64 = 0;
    var spl = std.mem.split(u8, input, "\n");
    while (spl.next()) |line| {
        const nums = try readnums(a, line);
        defer nums.deinit();
        const nxt = nums.getLast() + try predictNext(a, nums);
        // std.log.warn("nxt {}", .{nxt});
        sum += nxt;
    }
    return sum;
}

// note uses recursion!
fn predictNext(a: std.mem.Allocator, list: std.ArrayList(i64)) !i64 {
    if (std.mem.allEqual(i64, list.items, 0)) {
        return 0;
    }
    var sublist = std.ArrayList(i64).init(a);
    defer sublist.deinit();
    for (1..list.items.len) |ix| {
        try sublist.append(list.items[ix] - list.items[ix - 1]);
    }
    const nxt_diff = try predictNext(a, sublist);
    return sublist.getLast() + nxt_diff;
}

fn readnums(a: std.mem.Allocator, input: []const u8) !std.ArrayList(i64) {
    var toks = std.mem.tokenize(u8, input, " ");
    var res = std.ArrayList(i64).init(a);
    while (toks.next()) |tok| {
        try res.append(try std.fmt.parseInt(i64, tok, 10));
    }
    return res;
}

test "pt1" {
    try std.testing.expectEqual(@as(i64, 114), try solve_pt1(std.testing.allocator, test_input));
}

fn solve_pt2(a: std.mem.Allocator, input: []const u8) !i64 {
    var sum: i64 = 0;
    var spl = std.mem.split(u8, input, "\n");
    while (spl.next()) |line| {
        const nums = try readnums(a, line);
        defer nums.deinit();
        const prv = nums.items[0] - try predictPrev(a, nums);
        // std.log.warn("prv {}", .{prv});
        sum += prv;
    }
    return sum;
}

// note uses recursion!
fn predictPrev(a: std.mem.Allocator, list: std.ArrayList(i64)) !i64 {
    if (std.mem.allEqual(i64, list.items, 0)) {
        return 0;
    }
    var sublist = std.ArrayList(i64).init(a);
    defer sublist.deinit();
    for (1..list.items.len) |ix| {
        try sublist.append(list.items[ix] - list.items[ix - 1]);
    }
    return sublist.items[0] - try predictPrev(a, sublist);
}

test "pt2" {
    try std.testing.expectEqual(@as(i64, 2), try solve_pt2(std.testing.allocator, test_input));
}

const test_input =
    \\0 3 6 9 12 15
    \\1 3 6 10 15 21
    \\10 13 16 21 30 45
;

const puzzle_input = @embedFile("day9.in");