from aocd.models import Puzzle
from aocd import submit


def calibrate(data: list) -> int:
    """
    Solve part a.
    """
    # filter digits
    digits = [c for c in data if c.isdigit()]
    # return the sum of the first and the last digit
    return int(digits[0] + digits[-1])


def translate(s: str) -> str:
    """
    Solve part b.
    """
    nums = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
    conv = ["o1e", "t2o", "t3e", "f4r", "f5e", "s6x", "s7n", "e8t", "n9e"]

    for n, c in zip(nums, conv):
        while n in s:
            s = s.replace(n, c)
    return s


if __name__ == "__main__":
    # get puzzle and parse data
    puzzle = Puzzle(year=2023, day=1)

    # part a:
    answer_a = sum(calibrate(l) for l in puzzle.input_data.split("\n"))
    print(f"{answer_a}")
    submit(answer_a, part="a", day=1, year=2023)

    # part b:
    answer_b = sum(calibrate(translate(l)) for l in puzzle.input_data.split("\n"))
    print(f"{answer_b}")
    submit(answer_b, part="b", day=1, year=2023)