import sys
import os
sys.path.insert(1, os.path.abspath('../../'))
from python_tools.aoc_utils import *

TH = []
LIMIT = []
def find_trail_heads(input):
    for ci, i in enumerate(input):
        for cj, j in enumerate(i):
            if j == 0:
                TH.append([ci, cj])
VISITED = set()
def trail_scorep2(node, map):
    score = 0
    for k in [[-1,0],[1,0],[0,-1],[0,1]]:
        i = k[0]
        j = k[1]
        if ((0 <= node[0]+i < len(map)) and (0 <= node[1]+j < len(map[0]))):
            if(tuple([node[0]+i, node[1]+j]) not in VISITED):
                if map[node[0]][node[1]] + 1 == map[node[0]+i][node[1]+j]:
                    if map[node[0]+i][node[1]+j] == 9:
                        score += 1
                    else:
                        score += trail_scorep2([node[0]+i, node[1]+j], map)
    return score

def trail_score(node, map):
    VISITED.add(tuple(node))
    score = 0
    for k in [[-1,0],[1,0],[0,-1],[0,1]]:
        i = k[0]
        j = k[1]
        if ((0 <= node[0]+i < len(map)) and (0 <= node[1]+j < len(map[0]))):
            if(tuple([node[0]+i, node[1]+j]) not in VISITED):
                if map[node[0]][node[1]] + 1 == map[node[0]+i][node[1]+j]:
                    if map[node[0]+i][node[1]+j] == 9:
                        VISITED.add(tuple([node[0]+i, node[1]+j]))
                        score += 1
                    else:
                        score += trail_score([node[0]+i, node[1]+j], map)
    return score

def part_one(input):
    p1 = 0
    find_trail_heads(input)
    for i in TH:
        VISITED.clear()
        out = trail_score(i, input)
        p1 += out
        # print(out)
    print("Part One: ", p1)


def part_two(input):
    p2 = 0
    for i in TH:
        VISITED.clear()
        out = trail_scorep2(i, input)
        p2 += out
        # print(out)
    print("Part Two: ", p2)

def main():
    f = open('input.txt', 'r')
    contents = [list(map(int, list(x.strip()))) for x in f.readlines()]
    # for i in contents:
    #     print(''.join(str(i)))
    part_one(contents)
    part_two(contents)

if __name__ == "__main__":
    main()