AdventOfCode/2022/09/a.py

import os, argparse
import numpy as np

debug_print = print
debug_print = lambda x: x

def solve_task(lines):
    s = Snake()
    s.print()
    for line in lines:
        s.do(line)
        #s.print()
    print(len(s.visited))

class Snake:
    def __init__(self):
        self.H = np.array([0, 0])
        self.last_H = self.H
        self.T = np.array([0, 0])
        self.visited = set()
        self.visited.add((0,0))

    def print(self):
        print(f"H: {self.H}, T: {self.T}")
        print(self.visited)
    
    def do(self, instruction):
        command = instruction.split(' ')[0]
        distance = int(instruction.split(' ')[1])
        for _ in range(distance):
            self.go(command)
            self.adjust_tail()

    def go(self, direction):
        debug_print(f"Moving {direction}")
        self.last_H = self.H.copy()
        if direction == 'R':
            self.H[0] += 1
        elif direction == 'L':
            self.H[0] -= 1
        elif direction == 'U':
            self.H[1] += 1
        elif direction == 'D':
            self.H[1] -= 1
    
    def adjust_tail(self):
        debug_print(f"Euclidian distance from T:{self.T} to H:{self.H} is {np.linalg.norm(self.H - self.T)}")
        if np.linalg.norm(self.H - self.T) >= 2:
            debug_print(f"Tail moving from {self.T} to {self.last_H}")
            self.T = self.last_H
            self.visited.add(tuple(self.T))
        else:
            debug_print(f"Tail not moving from {self.T} to {self.last_H}")


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()