from collections import Counter
from typing import List, Tuple

from aoc import BaseAssignment

Vector = Tuple[Tuple[int, int], Tuple[int, int]]

class Assignment(BaseAssignment):
    def parse_item(self, item: str) -> Vector:
        start, end = item.split('->')
        return (
           tuple(int(_) for _ in start.strip().split(',')),
           tuple(int(_) for _ in end.strip().split(',')),
        )

    def read_input(self, example = False) -> List[Vector]:
        return list(super().read_input(example))


def is_horizontal(vector: Vector):
    return vector[0][0] == vector[1][0]

def is_vertical(vector: Vector):
    return vector[0][1] == vector[1][1]

def is_diagonal(vector: Vector):
    x_start = vector[0][0]
    x_end = vector[1][0]
    y_start = vector[0][1]
    y_end = vector[1][1]

    return abs(x_end - x_start) == abs(y_end - y_start)

def points_in_vector(vector: Vector, includes_diagonals: bool = False) -> List[Tuple[int, int]]:
    x_start = vector[0][0]
    x_end = vector[1][0]
    y_start = vector[0][1]
    y_end = vector[1][1]
    
    delta_x = x_end - x_start
    try:
        direction_x = int(delta_x / abs(delta_x))
    except ZeroDivisionError:
        direction_x = 0

    delta_y = y_end - y_start
    try:
        direction_y = int(delta_y / abs(delta_y))
    except ZeroDivisionError:
        direction_y = 0

    return [
        (x_start, y_start + y_delta)
        for y_delta
        in range(0, delta_y + direction_y, direction_y)
    ] if is_horizontal(vector) else [
        (x_start + x_delta, y_start)
        for x_delta
        in range(0, delta_x + direction_x, direction_x)
    ] if is_vertical(vector) else [
        (x_start + x_delta, y_start + y_delta)
        for x_delta, y_delta
        in zip(
           range(0, delta_x + direction_x, direction_x),
           range(0, delta_y + direction_y, direction_y)
        )
    ] if includes_diagonals and is_diagonal(vector) else []


class AssignmentOne(Assignment):
    example_result = 5

    def run(self, input: List[Vector]) -> int:
        coordinates = []

        for vector in input:
            coordinates += points_in_vector(vector)

        return len([c for c in Counter(coordinates).values() if c >= 2])


class AssignmentTwo(Assignment):
    example_result = 12

    def run(self, input: List[Vector]) -> int:
        coordinates = []

        for vector in input:
            coordinates += points_in_vector(vector, True)

        return len([c for c in Counter(coordinates).values() if c >= 2])