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):
    inputs = [
        [
            (37, "test_input07"),
            (336120, "input07")
        ],
        [
            (168, "test_input07"),
            (96864235, "input07")
        ]
    ]

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

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


if __name__ == '__main__':
    day = Day(7)
    day.run(verbose=True)