AdventOfCode/2022/09/b.py

import os, argparse
import numpy as np

debug_print = print
debug_print = lambda x: x

def solve_task(lines):
    s = Snake(9)
    s.print()
    for line in lines:
        s.do(line)
        #s.print()
    print(f"Tail visited {s.tail_visited()} spaces")

class Snake:
    def __init__(self, length = 1):
        self.length = length
        self.pos = np.array([0, 0])
        if length > 0:
            self.body = Snake(length-1)
        self.visited = set()
        self.visited.add(tuple(self.pos))
    
    def tail_visited(self):
        if self.length == 0:
            return len(self.visited)
        else:
            return self.body.tail_visited()

    def print(self):
        print(f"{self.length}: {self.pos}")
        #print(self.visited)
        if self.length > 0:
            self.body.print()
    
    def do(self, instruction):
        command = instruction.split(' ')[0]
        distance = int(instruction.split(' ')[1])
        for _ in range(distance):
            self.move(command)
            self.body.follow(self.pos)

    def move(self, direction):
        debug_print(f"Moving {direction}")
        if direction == 'R':
            self.pos[0] += 1
        elif direction == 'L':
            self.pos[0] -= 1
        elif direction == 'U':
            self.pos[1] += 1
        elif direction == 'D':
            self.pos[1] -= 1
    
    def follow(self, new_pos):
        debug_print(f"Euclidian distance from T:{self.pos} to H:{new_pos} to is {np.linalg.norm(self.pos - new_pos)}")
        if np.linalg.norm(self.pos - new_pos) >= 2:
            diff = new_pos - self.pos
            if diff[0] > 0:
                self.pos[0] += 1
            elif diff[0] < 0:
                self.pos[0] -= 1
            
            if diff[1] > 0:
                self.pos[1] += 1
            elif diff[1] < 0:
                self.pos[1] -= 1

            debug_print(f"Tail moving from {self.pos}")
            self.visited.add(tuple(self.pos))
            if self.length > 0:
                self.body.follow(self.pos)
        else:
            debug_print(f"Tail not moving from {self.pos}")


def read_lines(filename):
    lines = []
    with open(filename) as infile:
        for raw_line in infile:
            line = raw_line.rstrip()
            lines.append(line)
    return lines

def parse_arguments():
    parser = argparse.ArgumentParser(description="Script that solves the case",epilog="Have a nice day!")
    parser.add_argument('filename', nargs='?', default="example.txt", help='Input file')
    args = parser.parse_args()
    return args

def main():
    args = parse_arguments()
    lines = read_lines(args.filename)
    solve_task(lines)

if __name__ == "__main__":
    main()