#!/usr/bin/env python
# -*- encoding: utf-8; py-indent-offset: 4 -*-

# Author: Sebastian Mark
# CC-BY-SA (https://creativecommons.org/licenses/by-sa/4.0/deed.de)

# pylint: disable=missing-module-docstring,missing-function-docstring,consider-using-f-string,no-else-return

import re
import matplotlib.pyplot as plt


def readinput():
    with open("input", "r", encoding="utf-8") as file:
        robots = []
        for line in file:
            regex = r"p=(.?\d+),(.?\d+) v=(.?\d+),(.?\d+)"
            m = re.findall(regex, line)[0]
            robots.append(
                {
                    "pos": (int(m[0]), int(m[1])),  # x,y
                    "move": (int(m[2]), int(m[3])),  # x,y
                }
            )
    return robots


def get_robot_pos_after_time(area: dict, robot: tuple, seconds: int) -> tuple:
    finx = (robot["pos"][0] + robot["move"][0] * seconds) % area["width"]
    finy = (robot["pos"][1] + robot["move"][1] * seconds) % area["height"]

    return (finx, finy)


def get_quadrant_for_pos(area: dict, position: tuple) -> int:
    middle_x = area["width"] // 2
    middle_y = area["height"] // 2

    if position[0] == middle_x or position[1] == middle_y:
        return 0

    if position[1] < middle_y:
        if position[0] < middle_x:
            return 1
        else:
            return 2
    else:
        if position[0] < middle_x:
            return 3
        else:
            return 4


def count_positions_in_quadrants(area: dict, positions: dict) -> list:
    quadrant_count = {}
    for pos, count in positions.items():
        quadrant = get_quadrant_for_pos(area, pos)
        if quadrant != 0:
            quadrant_count[quadrant] = quadrant_count.get(quadrant, 0) + count

    return quadrant_count


def draw_area(area: dict, positions: dict, filename: str):
    grid = []
    for _ in range(area["height"]):
        row = [0] * area["width"]
        grid.append(row)

    for pos in positions.keys():
        x, y = pos
        grid[y][x] = 1

    plt.figure()
    plt.imshow(grid, cmap="gray")
    plt.axis("off")
    plt.savefig(f"{filename}.jpg", bbox_inches="tight")
    plt.close()


def main():
    robots = readinput()

    area = {"width": 101, "height": 103}

    # part 1
    seconds = 100
    positions = {}
    for robot in robots:
        pos = get_robot_pos_after_time(area, robot, seconds)
        positions[pos] = positions.get(pos, 0) + 1

    safety_factor = 1
    quadrant_count = count_positions_in_quadrants(area, positions)
    for c in quadrant_count.values():
        safety_factor *= c

    print("Safety factor: %d" % safety_factor)

    # part 2
    seconds = 10_000
    for i in range(seconds):
        positions = {}
        for robot in robots:
            pos = get_robot_pos_after_time(area, robot, i)
            positions[pos] = positions.get(pos, 0) + 1
        print(f"\rSaving image for positions after {i} seconds", end="", flush=True)
        draw_area(area, positions, str("%05d" % i))


if __name__ == "__main__":
    main()