# Some Useful Data Structures -------------------------------------------------

# Yes, there are four types of caves. Large caves, small caves, start caves, and
# end caves. Yes, they have different fields. Sue me. Each cave indicates the size, 
# name, and index of the cave, if it's important for that type of cave. The 
# indices are used later to determine whether a cave has already been visited.
abstract type Cave end

struct StartCave <: Cave end
struct EndCave   <: Cave end
struct LargeCave <: Cave name::String end
struct SmallCave <: Cave
    name::String
    index::Int
end

# Given a name and index, produce the appropriate type of Cave depending
# on whether the name is uppercase or not.
issmallcave(S) = all(islowercase, S) && S ∉ ["start", "end"]
Cave(s::AbstractString, idx::Int)::SmallCave = SmallCave(s, idx)
function Cave(s::AbstractString, ::Nothing)::Cave
    s == "start" && return StartCave()
    s == "end"   && return EndCave()
    return LargeCave(s)
end


# Ingest the Data File ---------------------------------------------------------

# Read in each line of the input file and parse it as an 'edge'. The final product
# here will be a Dict, where the keys are `Cave`s and the values are the list of
# other `Cave`s the key Cave is connected to. This behaves similarly to an 
# adjacency list, except you can look up cave connections in O(1) time.
function ingest(path)
    paths = open(path) do f
        [split(line, "-") for line in readlines(f)]
    end

    cavemap = Dict{Cave, Vector{Cave}}()
    indexes = Dict{String,Int}()
    for path in paths
        # Need to add entries to the output dictionary for both paths, from
        # the left cave to the right, and from the right cave to the left.
        (left, right) = path

        # Get the appropriate index for each cave. If there's already an 
        # assigned index for that cave name, just use it, otherwise use one more
        # than the total number of caves already indexed. Only do this if the
        # cave name is all lowercase, a small cave.
        if issmallcave(left) && get(indexes, left, 0) == 0
            indexes[left] = length(indexes) + 1
        end
        if issmallcave(right) && get(indexes, right, 0) == 0
            indexes[right] = length(indexes) + 1
        end

        # Make the `Caves`!
        leftcave  = Cave(left,  get(indexes, left, nothing))
        rightcave = Cave(right, get(indexes, right, nothing))

        # Add the left cave to the cavemap
        destinations = get(cavemap, leftcave, Vector{Cave}())
        push!(destinations, rightcave)
        cavemap[leftcave] = destinations

        # Add the right cave to the cavemap
        destinations = get(cavemap, rightcave, Vector{Cave}())
        push!(destinations, leftcave)
        cavemap[rightcave] = destinations
    end
    return cavemap
end

# Some Useful Helper Functions and Types ---------------------------------------

# Just some handy type aliases for clarity
const CaveMap = Dict{Cave, Vector{Cave}}
const ExploreStack = Vector{Tuple{BitVector,Cave}}


# We only need to keep track of the small caves. If moving into a large
# cave, we only need to pass back the set of visited Caves. If we're moving
# into the end cave, then we don't need to know about the visited caves 
# anymore, we can just use an empty BitVector.
nextpath(visited::BitVector, ::LargeCave) = visited
nextpath(::BitVector, ::EndCave) = BitVector()
function nextpath(visited::BitVector, nextcave::SmallCave) 
    visited = deepcopy(visited)
    visited[nextcave.index] = true
    return visited
end

# Determines whether we should skip entering a cave. We should always enter a 
# Large cave (don't skip it) or the end cave, never re-enter the start cave,
# and only enter a small cave if we've never been inside it before (meaning
# it's corresponding index in the visited vector will be false).
shouldskip(::BitVector, ::LargeCave) = false
shouldskip(::BitVector, ::EndCave) = false
shouldskip(::BitVector, ::StartCave) = true
shouldskip(visited::BitVector, cave::SmallCave) = visited[cave.index]


# Solve Part One ---------------------------------------------------------------

# This is a pretty standard implementation of a depth-first search, with a bit
# of a twist. Most of the twist is handled by the different helper functions
# and types implemented above.
function part1(input)
    # Start by creating a "stack" and initializing it with the "start" cave 
    # and a BitVector large enough to hold the indices of all the small caves
    smallcavecount = mapreduce(C -> C isa SmallCave, +, keys(input))
    stack = ExploreStack([(falses(smallcavecount), StartCave())])
    paths = 0

    # While there are still caves to explore...
    while !isempty(stack)
        # Pop the last cave and the list of visited caves off the stack...
        (visitedsofar, cave) = pop!(stack)

        # If we've reached the end, just add one to our count of paths
        # and move on to the next item in the stack
        if cave isa EndCave
            paths += 1
            continue
        end
        
        # Otherwise, get the list of all the caves the current cave is
        # connected to, and add them to the stack to be visited next
        for nextcave in get(input, cave, [])
            # If we found the start, or if our next cave is in our list of
            # visited caves, don't go there.
            shouldskip(visitedsofar, nextcave) && continue

            # We only need to keep track of the small caves we've visited
            visited = nextpath(visitedsofar, nextcave)
            push!(stack, (visited, nextcave))
        end
    end
    return paths
end

# Some More Useful Structs and Methods -----------------------------------------

# We need to change the types of our list of visited caves and exploration 
# stack to take advantage of the new methods for identifying the next path
# and whether we should skip the next cave.
const Path = Tuple{BitVector,Bool}
const ExploreStackTwo = Vector{Tuple{Path,Cave}}

# The logic is the same for large caves and end caves, but now we need to
# account for whether we've visited a single small cave twice when 
# updating our `Path`
nextpath(path::Path, ::LargeCave)= path
nextpath(::Path, ::EndCave) = (BitVector(), false)
function nextpath(path::Path, nextcave::SmallCave)
    (visited, twovisits) = path
    twovisits = twovisits || visited[nextcave.index]
    visited = deepcopy(visited)
    visited[nextcave.index] = true
    return (visited, twovisits)
end

# The logic for large caves, end caves, and start caves remains unchanged, but
# we needed to define them again because our first version was too specialized.
# These versions using `::Any` would work for part one and part two. The big
# change is the logic for determining whether to skip a small cave.
shouldskip(::Any, ::LargeCave) = false
shouldskip(::Any, ::EndCave) = false
shouldskip(::Any, ::StartCave) = true
function shouldskip(path::Path, cave::SmallCave) 
    (visited, twovisits) = path
    return visited[cave.index] && twovisits
end


# Solve Part Two ---------------------------------------------------------------

# This bit is *exactly* the same as part one. Well, except for the first few
# lines. Here, we've changed the types of the arguments in our stack to use
# the new logic methods from above. Otherwise, the algorithm is exactly the
# same.
function part2(input)
    # Start by creating a "stack" and initializing it with the "start" cave , 
    # a BitVector large enough to hold the indices of all the small caves, and
    # a boolean value to indicate whether we've seen a small cave twice.
    smallcavecount = mapreduce(C -> C isa SmallCave, +, keys(input))
    initialpath = (falses(smallcavecount), false)
    stack = ExploreStackTwo([(initialpath, StartCave())])
    paths = 0

    # While there are still caves to explore...
    while !isempty(stack)
        # Pop the last cave and the `Path` off the stack...
        (pathsofar, cave) = pop!(stack)

        # If we've reached the end, just add one to our count of paths
        # and move on to the next item in the stack
        if cave isa EndCave
            paths += 1
            continue
        end
        
        # Otherwise, get the list of all the caves the current cave is
        # connected to, and add them to the stack to be visited next
        for nextcave in get(input, cave, [])
            # If we found the start, or if our puzzle logic tells us to skip
            # the next cave, don't go there.
            shouldskip(pathsofar, nextcave) && continue

            # We only need to keep track of the small caves we've visited, still
            path = nextpath(pathsofar, nextcave)
            push!(stack, (path, nextcave))
        end
    end
    return paths
end

# Read in the input data file, parsing it into a matrix of numbers
function ingest(path)
    outvectors = open(path) do f
        [collect(line) for line in readlines(f)]
    end
    toint(x) = parse(Int8, x)

    # I transposed the matrix to make debugging easier, since 
    # otherwise it won't print in the same orientation as the input.
    (outmatrix
     =  outvectors
     |> (x -> reduce(hcat, x))
     |> (x -> map(toint, x))
     |> transpose
     |> collect)
    return outmatrix
end

# Helper Functions -------------------------------------------------------------

# Return a view of the matrix `M` centered on the index `idx` and encompassing
# the eight surrounding indices. Accounts for corners and edges.
function neighborhoodof(idx::CartesianIndex, M::AbstractMatrix)
    (rows, cols) = size(M)
    (row, col) = Tuple(idx)

    top    = row > 1    ? row - 1 : row
    left   = col > 1    ? col - 1 : col
    bottom = row < rows ? row + 1 : row
    right  = col < cols ? col + 1 : col

    return view(M, top:bottom, left:right)
end

# Given a matrix representing our octopuses, advance the state of the 
# school of octopuses, mutating the input and returning the number of
# octopuses who flashed this round.
function advance!(octopuses::Matrix{Int8})::Int
    # Increase the value of all octopuses
    octopuses .+= 1

    # Identify the octopuses who have flashed
    flashed = justflashed = octopuses .> 9

    # Shorthand to fetch octopuses surrounding an index. There's no need to
    # exclude the central octopus, since they've already flashed.
    getneighborhood(idx) = neighborhoodof(idx, octopuses)

    # So long as any octopus just crossed the threshold...
    while any(justflashed)
        # Get all the neighborhoods of the octopuses who just flashed
        neighborhoods = map(getneighborhood, findall(justflashed))

        # For each neighborhood, increase the energy level of all
        # octopuses by 1
        foreach(N -> N .+= 1, neighborhoods)

        # Identify the octopuses who just crossed the threshold
        justflashed = (octopuses .> 9) .& .!flashed

        # Update the listing of all octopuses who have flashed this round
        flashed .|= justflashed
    end

    # Reset all the octopuses who have flashed and return the count
    octopuses[flashed] .= 0
    return count(flashed)
end


# Advance the school of octopuses 100 times, counting the number of octopuses
# who flash each round and returning the count of all octopuses flashing over
# 100 rounds
function part1(input)
    flashes = 0
    octopuses = deepcopy(input)
    for _ in 1:100
        flashes += advance!(octopuses)
    end
    return flashes
end

# Now, we advance the octopuses until we reach a state where all octopuses
# have been reset at the same time, meaning they all flashed simultaneously.
function part2(input)
    octopuses = deepcopy(input)
    rounds = 0
    while any(octopuses .> 0)
        advance!(octopuses)
        rounds += 1
    end
    return rounds
end