from tools.aoc import AOCDay
from tools.int_seq import triangular
from typing import Any, List


def getMinFuelUse(crabs: List[int], increased: bool = False) -> int:
    min_num = min(crabs)
    max_num = max(crabs)
    avg = (min_num + max_num) // 2
    while 1:
        if increased:
            avg_fuel = sum(triangular(abs(avg - c)) for c in crabs)
            avg_fuel_p = sum(triangular(abs(avg + 1 - c)) for c in crabs)
            avg_fuel_m = sum(triangular(abs(avg - 1 - c)) for c in crabs)
        else:
            avg_fuel = sum(abs(avg - c) for c in crabs)
            avg_fuel_p = sum(abs(avg + 1 - c) for c in crabs)
            avg_fuel_m = sum(abs(avg - 1 - c) for c in crabs)

        if avg_fuel_m < avg_fuel:
            max_num = avg
            avg = (min_num + avg) // 2
        elif avg_fuel_p < avg_fuel:
            min_num = avg
            avg = (max_num + avg) // 2
        else:
            return avg_fuel


class Day(AOCDay):
    test_solutions_p1 = [37]
    test_solutions_p2 = [168]

    def part1(self) -> Any:
        crabs = self.getInputAsArraySplit(",", int)
        return getMinFuelUse(crabs)

    def part2(self) -> Any:
        crabs = self.getInputAsArraySplit(",", int)
        return getMinFuelUse(crabs, True)