import sys
sys.path.append("../common")
import aoc
import networkx

adjacent = [[-1, 0], [0, -1], [1, 0], [0, 1]]

# create symbol definitions
symbols = [".", "=", "|"]
rocky, wet, narrow = 0, 1, 2
torch, gear, neither = 0, 1, 2

tooloptions = dict()
tooloptions[rocky] = (torch, gear)
tooloptions[wet] = (gear, neither)
tooloptions[narrow] = (torch, neither)

# parse input file
def parse_input(si):
    s = si.split("\n")
    depth = int(s[0].split(": ")[1])
    target = [int(v) for v in s[1].split(": ")[1].split(",")]
    return depth, target

# get erosion level from geological index
def get_erosion_level(n):
    return (n + depth) % 20183

# generate map of erosion types
def gen_map():
    grid = [[0 for x in range(mrange[0])] for y in range(mrange[1])]

    # generate values on x side
    grid[0] = [get_erosion_level(x * 16807) for x in range(mrange[0])] 

    # generate values on y side
    for y in range(mrange[1]):
        grid[y][0] = get_erosion_level(y * 48271)

    # fill in other positions
    for y in range(1, mrange[1]):
        for x in range(1, mrange[0]):
            grid[y][x] = get_erosion_level(grid[y-1][x] * grid[y][x-1])

    # mod 3 for env type
    grid = [[grid[y][x] % 3 for x in range(mrange[0])] for y in range(mrange[1])]

    # start position is rocky
    grid[target[1]][target[0]] = 0

    return grid

# define constants from input file
depth, target = parse_input(aoc.data)
mrange = (target[0] + 100, target[1] + 100) # 100 block padding for potential hook-paths

vmap = gen_map()

def risk_level(width, height):
    risk = 0
    for y in range(height + 1):
        for x in range(width + 1):
            risk += vmap[y][x]
    return risk

def adj(p):
    return [[p[0] + di[0], p[1] + di[1]] for di in adjacent]

def dijkstra():
    graph = networkx.Graph()
    for y in range(mrange[1]):
        for x in range(mrange[0]):
            tools = tooloptions[vmap[y][x]]
            # always takes 7 minutes to switch, 2 z-options for each tool
            graph.add_edge((x, y, tools[0]), (x, y, tools[1]), weight = 7)
            for (nx, ny) in adj((x, y)):
                if 0 <= nx < mrange[0] and 0 <= ny < mrange[1]:
                    ntools = tooloptions[vmap[ny][nx]]
                    # if tool is usable in both environments
                    for tool in [t for t in tools if t in ntools]:
                         # then it only takes 1 minute
                        graph.add_edge((x, y, tool), (nx, ny, tool), weight = 1)

    return networkx.dijkstra_path_length(graph,
        (0, 0, torch), (target[0], target[1], torch))

def solve1(args):
    return risk_level(target[0], target[1])

def solve2(args):
    return dijkstra()

aoc.run(solve1, solve2, sols=[11359, 976])