<reponame>jcabre04/advent-of-code
import collections

# Adjacency list using code tweaked from: 
# https://www.educative.io/edpresso/how-to-implement-a-graph-in-python

# Add a vertex to the dictionary
def add_vertex(graph, v):
    if v in graph:
        pass
        #print("Vertex ", v, " already exists.")
    else:
        graph[v] = []

# Add an edge between vertex v1 and v2 with edge weight e
def add_edge(graph, v1, v2, e):
    # Check if vertex v1 is a valid vertex
    if v1 not in graph:
        pass
        print("Vertex ", v1, " does not exist.")
    # Check if vertex v2 is a valid vertex
    elif v2 not in graph:
        pass
        print("Vertex ", v2, " does not exist.")
    else:
        # Since this code is not restricted to a directed or 
        # an undirected graph, an edge between v1 v2 does not
        # imply that an edge exists between v2 and v1
        temp = [v2, e]
        graph[v1].append(temp)

# Print the graph
def print_graph(graph):
    for vertex in graph:
        for edges in graph[vertex]:
            print(vertex, " -> ", edges[0], " edge weight: ", edges[1])

def build_graph(lines):
    graph = {}

    #build graph:
    for line in lines:
        start, end = line.split("-")
        add_vertex(graph, start)
        add_vertex(graph, end)
        add_edge(graph, start, end, 1)
        add_edge(graph, end, start, 1)
    return graph

# edge is a list of [vertex aka graph key, edge weight]
# path is a list of vertices representing the path taken
def travel(graph, edge, path):
    if edge[0] == 'end':
        return 1

    total_paths = 0
    new_path = path + [edge[0]]
    current_cave = graph[edge[0]]
    
    for new_edge in current_cave:
        if new_edge[0].islower() and new_edge[0] not in new_path:
            total_paths += travel(graph, new_edge, new_path)
        elif new_edge[0].isupper():
            total_paths += travel(graph, new_edge, new_path)
    
    return total_paths

def part1(lines):
    graph = build_graph(lines)
    paths = 0
    
    start = graph["start"]
    for edge in start:
        path = ["start"]
        paths += travel(graph, edge, path)

    return paths

def lower_cave_checks(cave_name, new_path):
    if cave_name == "end":
        return True

    if cave_name == "start":
        return False
    
    path_lower_names = [name for name in new_path if name.islower()]
    lower_name_counter = collections.Counter(path_lower_names)

    if cave_name in lower_name_counter and 2 in set(lower_name_counter.values()):
        return False

    return True

def travel_plus(graph, edge, path):
    if edge[0] == 'end':
        return 1

    total_paths = 0
    new_path = path + [edge[0]]
    current = graph[edge[0]]
    
    for new_edge in current:
        if new_edge[0].islower() and lower_cave_checks(new_edge[0], new_path):
            total_paths += travel_plus(graph, new_edge, new_path)
        elif new_edge[0].isupper():
            total_paths += travel_plus(graph, new_edge, new_path)
    return total_paths

def part2(lines):
    graph = build_graph(lines)
    paths = 0
    
    start = graph["start"]
    for edge in start:
        path = ["start"]
        paths += travel_plus(graph, edge, path)

    return paths

if __name__ == "__main__":
    with open("../input.txt") as file:
        lines = [line.strip() for line in file]

    print(part1(lines))
    print(part2(lines))