#!/usr/bin/env python3

import collections
import sys

input_file = sys.argv[1]

with open(input_file) as fd:
    lines = [line.rstrip() for line in fd.readlines()]

height = len(lines)
width = len(lines[0])

# Adjust parameters for part/file
part = int(sys.argv[2])
assert part in (1, 2)

if input_file.startswith("input"):
    minic = 100
else:
    if part == 1:
        minic = 1
    else:
        minic = 50
skips = 1 if part == 1 else 20
canon: collections.Counter[int] = collections.Counter()
demo = {}
if input_file == "demo":
    if skips == 1:
        demo = {2: 14, 4: 14, 6: 2, 8: 4, 10: 2} | (
            {12: 3, 20: 1, 36: 1, 38: 1, 40: 1, 64: 1}
        )
    elif skips == 20:
        demo = {50: 32, 52: 31, 54: 29, 56: 39, 58: 25, 60: 23} | (
            {62: 20, 64: 19, 66: 12, 68: 14, 70: 12, 72: 22, 74: 4, 76: 3}
        )
    for k, v in demo.items():
        canon[k] = v

vec = tuple[int, int]
directions = [
    (-1, 0),  # ^ North
    (0, 1),  # > East
    (1, 0),  # v South
    (0, -1),  # < West
]

# Find start position
for i, line in enumerate(lines):
    if "S" in line:
        j = line.index("S")
        start = i, j


# Visit normally
normal = None
visited: list[list[int | None]] = list()
for _ in range(height):
    visited.append([None] * width)
visited[start[0]][start[1]] = 0
stack: set[vec] = {start}
s = 0
while stack:
    s += 1
    nstack: set[vec] = set()
    for pos in stack:
        i, j = pos
        for d, direction in enumerate(directions):
            ii, jj = i + direction[0], j + direction[1]

            previs = visited[ii][jj]
            if previs is not None and previs < s:
                continue
            visited[ii][jj] = s

            cchar = lines[ii][jj]
            if cchar == "#":
                continue
            elif cchar == "E":
                if normal is None:
                    normal = s
            nstack.add((ii, jj))
    stack = nstack
assert normal

# Print
for i in range(height):
    line = ""
    for j in range(width):
        char = lines[i][j]
        if visited[i][j] is not None:
            if char == "#":
                char = "@"
            else:
                char = "O"
        line += char
    print(line)
print()


# Find cheats
saves: collections.Counter[int] = collections.Counter()
for i in range(1, height - 1):
    if height > 100:
        print(103, i, "/", height)
    for j in range(1, width - 1):
        char = lines[i][j]
        if char == "#":
            continue
        ovis = visited[i][j]
        if ovis is None:
            continue
        if ovis >= normal:
            continue
        # for di in range(-skips, skips):
        #     ii = i + di
        #     G
        for ii in range(1, height - 1):
            for jj in range(1, width - 1):
                manh = abs(i - ii) + abs(j - jj)
                if manh > skips:
                    continue
                cchar = lines[ii][jj]
                if cchar == "#":
                    continue
                nvis = visited[ii][jj]
                if nvis is None:
                    continue
                orem = normal - ovis
                nrem = abs(normal - nvis) + manh
                save = orem - nrem
                if save < minic:
                    continue
                saves[save] += 1


print(f"{normal=}")
print(f"{dict(sorted(saves.items()))=}")
if demo:
    print(f"{dict(sorted(canon.items()))=}")
    diff = canon.copy()
    diff.subtract(saves)
    print(f"{dict(sorted(diff.items()))=}")
    print(f"{(saves == canon)=}")
    print(f"{saves.total()=}")
    print(f"{canon.total()=}")
    difft = 0
    for v in diff.values():
        difft += abs(v)
    print(f"{difft=}")
print(saves.total())
# 1119834 too high
# 982425 correct!