zerosleeps

from unittest import TestCase
from utils import get_raw_input


class Rope:
    def __init__(self, number_of_knots=1):
        self.head = [0, 0]
        self.knots = [[0, 0] for _ in range(number_of_knots)]
        self.tail_visits = set()

    def move_head(self, direction):
        if direction == "U":
            self.head[1] += 1
        elif direction == "R":
            self.head[0] += 1
        elif direction == "D":
            self.head[1] -= 1
        elif direction == "L":
            self.head[0] -= 1
        self.move_knots()

    def move_knots(self):
        for i, knot in enumerate(self.knots):
            if i == 0:
                knot_to_follow = self.head
            else:
                knot_to_follow = self.knots[i - 1]

            x_difference = knot_to_follow[0] - knot[0]
            y_difference = knot_to_follow[1] - knot[1]

            if x_difference > 1 or (x_difference == 1 and abs(y_difference) == 2):
                knot[0] += 1
            if x_difference < -1 or (x_difference == -1 and abs(y_difference) == 2):
                knot[0] -= 1
            if y_difference > 1 or (y_difference == 1 and abs(x_difference) == 2):
                knot[1] += 1
            if y_difference < -1 or (y_difference == -1 and abs(x_difference) == 2):
                knot[1] -= 1

            if i + 1 == len(self.knots):
                self.tail_visits.add(tuple(knot))

    @property
    def number_of_tail_visits(self):
        return len(self.tail_visits)


def run(rope, moves):
    for move in moves:
        direction, distance = move.split()
        for i in range(int(distance)):
            rope.move_head(direction)
    return rope


def part_one(input):
    rope = Rope()
    run(rope, input)
    return rope.number_of_tail_visits


def part_two(input):
    rope = Rope(9)
    run(rope, input)
    return rope.number_of_tail_visits


if __name__ == "__main__":
    input = get_raw_input("day_09_input.txt")
    print(f"Part one: {part_one(input)}")
    print(f"Part two: {part_two(input)}")


class TextExamples(TestCase):
    def setUp(self):
        self.first_example = [
            "R 4",
            "U 4",
            "L 3",
            "D 1",
            "R 4",
            "D 1",
            "L 5",
            "R 2",
        ]
        self.second_example = [
            "R 5",
            "U 8",
            "L 8",
            "D 3",
            "R 17",
            "D 10",
            "L 25",
            "U 20",
        ]

    def test_part_one_example(self):
        self.assertEqual(part_one(self.first_example), 13)

    def test_part_one_example_one(self):
        self.assertEqual(part_two(self.first_example), 1)

    def test_part_one_example_two(self):
        self.assertEqual(
            part_two(self.second_example),
            36,
        )