package main

import "core:math"
import "core:slice"
import "core:strconv"
import "core:strings"
import "core:bytes"
import "core:fmt"
import "core:os"

import "../util"

part_1 :: proc(lines: []string) {
    result: i64 = 0
    
    numbers := slice.clone_to_dynamic(
        slice.mapper(lines, proc(line: string) -> i64 {
            num, ok := strconv.parse_i64(line); assert(ok)
            return num
    }))
    slice.sort(numbers[:])

    inject_at(&numbers, 0, 0)
    append(&numbers, slice.last(numbers[:]) + 3)

    ones: i64 = 0
    threes: i64 = 0
    for _, idx in numbers[:len(numbers)-1] {
        switch numbers[idx+1] - numbers[idx] {
        case 1: ones += 1
        case 3: threes += 1
        }
    }
    result = ones * threes

    fmt.printfln("Part 1: %d", ones * threes)
}

part_2 :: proc(lines: []string) {
    result: i64 = 1
    
    numbers := slice.clone_to_dynamic(
        slice.mapper(lines, proc(line: string) -> i64 {
            num, ok := strconv.parse_i64(line); assert(ok)
            return num
    }))
    slice.sort(numbers[:])

    inject_at(&numbers, 0, 0)
    append(&numbers, slice.last(numbers[:]) + 3)
    combinations := make([]i64, len(numbers))
    combinations[len(combinations)-1] = 1

    for idx1 := len(numbers) - 2 ; idx1 >= 0; idx1 -= 1 {
        ways: i64 = 0
        for idx2 := idx1 + 1; idx2 <= idx1 + 3; idx2 += 1 {
            if idx2 >= len(numbers) do continue
            if numbers[idx2] - numbers[idx1] <= 3 do ways += combinations[idx2]
        }
        combinations[idx1] = ways
    }
    result = combinations[0]

    fmt.printfln("Part 2: %d", result)
}

main :: proc() {
    context.allocator = context.temp_allocator
    defer free_all(context.temp_allocator)

    INPUT :: #load("input.txt", string)
    lines := strings.split(INPUT, "\n")
    lines = lines[:len(lines)-1]

    part_1(lines)
    part_2(lines)
}