const aoc = @import("./aoc.zig");
const std = @import("std");
const Allocator = std.mem.Allocator;
const assert = std.debug.assert;

const History = std.ArrayList(i32);
const HistoryChanges = std.ArrayList([]i32);
const Input = struct {
    histories: std.ArrayList(History),

    fn deinit(self: *Input) void {
        for (self.histories.items) |history| {
            history.deinit();
        }
        self.histories.deinit();
    }
};

fn parse_history(allocator: Allocator, line: []const u8) !History {
    var history = History.init(allocator);
    var tokens = std.mem.tokenizeScalar(u8, line, ' ');
    while (tokens.next()) |token| {
        const number = try std.fmt.parseInt(i32, token, 10);
        try history.append(number);
    }
    return history;
}

fn parse_input(allocator: Allocator, lines: []const []const u8) !Input {
    var histories = std.ArrayList(History).init(allocator);
    errdefer histories.deinit();

    for (lines) |line| {
        const history = try parse_history(allocator, line);
        try histories.append(history);
    }

    return Input{ .histories = histories };
}

fn is_all_zeros(numbers: []i32) bool {
    for (numbers) |num| {
        if (num != 0) {
            return false;
        }
    }
    return true;
}

fn calc_changes(result: []i32, numbers: []i32) void {
    assert(result.len >= numbers.len-1);
    for (0..(numbers.len-1)) |i| {
        result[i] = numbers[i+1] - numbers[i];
    }
}

fn all_changes_of_history(allocator: Allocator, history: []i32) !HistoryChanges {
    const max_changes_len: usize = history.len-1;
    assert(max_changes_len >= 1);

    var changes = HistoryChanges.init(allocator);
    errdefer free_history_changes(changes);

    try changes.append(try allocator.alloc(i32, history.len-1));
    calc_changes(changes.items[0], history);

    for (1..max_changes_len) |i| {
        if (is_all_zeros(changes.items[i-1])) break;

        try changes.append(try allocator.alloc(i32, history.len-i-1));
        calc_changes(changes.items[i], changes.items[i-1]);
    }

    return changes;
}

fn free_history_changes(changes: HistoryChanges) void {
    for (changes.items) |row| {
        changes.allocator.free(row);
    }
    changes.deinit();
}

pub fn part1(input: *aoc.Input) !aoc.Result {
    const allocator = input.allocator;
    var parsed = try parse_input(allocator, input.lines);
    defer parsed.deinit();

    var result: i32 = 0;
    for (parsed.histories.items) |history| {
        const numbers = history.items;

        const changes = try all_changes_of_history(allocator, numbers);
        defer free_history_changes(changes);

        var prediction: i32 = numbers[numbers.len-1];
        for (changes.items) |change_row| {
            prediction += change_row[change_row.len-1];
        }

        result += prediction;
    }

    return .{ .int = result };
}

pub fn part2(input: *aoc.Input) !aoc.Result {
    const allocator = input.allocator;
    var parsed = try parse_input(allocator, input.lines);
    defer parsed.deinit();

    var result: i32 = 0;
    for (parsed.histories.items) |history| {
        const numbers = history.items;

        const changes = try all_changes_of_history(allocator, numbers);
        defer free_history_changes(changes);

        var prediction: i32 = 0;
        for (0..changes.items.len) |i| {
            const change_row = changes.items[changes.items.len-i-1];
            prediction = change_row[0] - prediction;
        }

        prediction = numbers[0] - prediction;

        result += prediction;
    }

    return .{ .int = result };
}

const example_input = [_][]const u8{
    "0 3 6 9 12 15",
    "1 3 6 10 15 21",
    "10 13 16 21 30 45",
};

test "part 1 example" {
    try aoc.expectAnswerInt(part1, 114, &example_input);
}

test "part 2 example" {
    try aoc.expectAnswerInt(part2, 2, &example_input);
}