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

# Basically just an enum for `Axis` types
abstract type Axis end
struct XAxis <: Axis end
struct YAxis <: Axis end

struct Fold
    axis::Axis
    index::Int
end

mutable struct Paper
    dots::BitMatrix
    folds::Vector{Fold}
end


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

function ingest(path)
    return open(path) do f

        # Start by reading all the lines from the top section of the input file
        # and parsing each one into a `CartesianIndex`.
        coordinateslist = []
        for coordinatestring in split(readuntil(f, "\n\n"))
            (col, row) = [parse(Int, s) for s in split(coordinatestring, ",")]
            push!(coordinateslist, CartesianIndex(row+1, col+1))
        end

        # Next read in the lines from the bottom section and parse each
        # one into a `Fold` struct.
        folds = []
        foldre = r"(x|y)=(\d+)"
        for foldstring in split(readchomp(f), "\n")
            m = match(foldre, foldstring)
            axis = m[1] == "x" ? XAxis() : YAxis()
            index = parse(Int, m[2]) + 1
            fold = Fold(axis, index)
            push!(folds, fold)
        end

        # Figure out how many rows/columns the `dots` Matrix should be. We
        # assume that it must be tall and wide enough to accommodate the
        # largest horizontal and vertical fold.
        rows = cols = 0
        for fold in folds
            if isa(fold.axis, XAxis)
                cols = max(cols, (fold.index * 2) - 1) 
            else
                rows = max(rows, (fold.index * 2) - 1) 
            end
        end

        # Make a large enough BitMatrix to accommodate all the 'dots' and
        # and set each coordinate found above to `true`.
        dots = falses(rows, cols)
        for coordinate in coordinateslist
            dots[coordinate] = true
        end

        Paper(dots, folds)
    end
end

# Given a BitMatrix representing the current arrangement of dots on a page
# and a fold indicating where/how to fold that paper, fold the paper and
# return the result.
function foldit(dots::BitMatrix, fold::Fold)::BitMatrix
    (rows, cols) = size(dots)

    # Need to define two views into the `dots` BitMatrix, one for the 
    # half of the paper that will stay in place, and one for the half
    # of the paper to be 'folded over'. The 'folded' view will have
    # its rows/columns reversed as appropriate.
    toprows  = bottomrows = 1:rows
    leftcols = rightcols  = 1:cols
    if isa(fold.axis, XAxis)
        leftcols   = 1:(fold.index-1)
        rightcols = cols:-1:(fold.index+1)
    else
        toprows    = 1:(fold.index-1)
        bottomrows = rows:-1:(fold.index+1)
    end

    # Take the two views, overlay them, then return the result
    still  = view(dots, toprows,    leftcols) # Stationary!
    folded = view(dots, bottomrows, rightcols)
    return (still .| folded)
end

# Take the input, fold it once, then return the number of 'dots' from
# that single fold.
function part1(input)
    dots = input.dots
    fold = input.folds[1]
    folded = foldit(dots, fold)
    return count(folded)
end

# Iteration for a `Paper` Struct ----------------------------------------------
struct PaperFoldState
    dots::BitMatrix
    lastfold::Int
end

# Initial iterator, returns the first fold
function Base.iterate(paper::Paper)
    fold = paper.folds[1]
    folded = foldit(paper.dots, fold)
    state = PaperFoldState(folded, 1)
    return (folded, state)
end

# Subsequent iterator, given the last state, return tne next state
function Base.iterate(paper::Paper, state::PaperFoldState)
    (dots, lastfold) = (state.dots, state.lastfold)
    lastfold >= length(paper.folds) && return nothing
    fold = paper.folds[lastfold + 1]
    folded = foldit(dots, fold)
    state = PaperFoldState(folded, lastfold + 1)
    return (folded, state)
end

# Needed to be able to get a specific fold state of the `Paper`
function Base.getindex(paper::Paper, i::Int)
    for (iteration, dots) in enumerate(paper)
        iteration > i && return dots
    end
    throw(BoundsError(paper, i))
end

# Some more necessary implementations so I can use the `paper[end]` syntax
# in our solution function.
Base.length(paper::Paper)    = length(paper.folds) - 1
Base.lastindex(paper::Paper) = lastindex(paper.folds) - 1
Base.eltype(::Type{Paper})   = BitMatrix

# Pretty prints a BitMatrix to make the solution to Part Two more
# readable, because reading the block characters from the default
# 1/0 display of the BitMatrix is difficult.
function prettyprint(matrix::BitMatrix)
    for line in eachrow(matrix)
        for bit in line
            if bit; print('█'); else; print(' '); end
        end
        println()
    end
end


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

# Fold the paper until all the folds are used up. The commented out print
# statement is there for solving the puzzle. The rest of it is there for 
# comparing in my test suite.
function part2(input)
    final_paper = input[end]
    # prettyprint(final_paper)
    rowsums = mapslices(sum, final_paper, dims = 2)
    colsums = mapslices(sum, final_paper, dims = 1)
    return (vec(rowsums), vec(colsums))
end

# 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