include("utils.jl")

function parse_input(filename)
    map(x->parse.(Int, x), split.(non_empty_lines(filename), ""))
end

function sink_locations(width, height, heights)
    sinks = []
    for x in 1:width
        for y in 1:height
            minimum = true
            if x > 1
                minimum &= heights[x - 1, y] > heights[x, y]
            end
            if x < width
                minimum &= heights[x + 1, y] > heights[x, y]
            end
            if y > 1
                minimum &= heights[x, y - 1] > heights[x, y]
            end
            if y < height
                minimum &= heights[x, y + 1] > heights[x, y]
            end
            if minimum
                push!(sinks, (x, y))
            end
        end
    end
    sinks
end

function solution9_1()
    input = collect(parse_input("9_ex.data"))
    width = length(input[1])
    height = length(input)
    heights = reshape(reduce(vcat, input),(width,height))
    sinks = sink_locations(width, height, heights)
    risk = 0
    for sink in sinks
        risk += heights[sink[1], sink[2]] + 1
    end
    risk
end

function basin_size(x, y, width, height, heights)
    heights = copy(heights)
    to_visit = [CartesianIndex(x, y)]
    visited = []
    while length(to_visit) > 0
        current = pop!(to_visit)
        if current in visited
            continue
        end
        push!(visited, current)
        if heights[current] == 9
            continue
        end
        if current[1] > 1
            push!(to_visit, CartesianIndex(-1, 0) + current)
        end
        if current[1] < width
            push!(to_visit, CartesianIndex(+1, 0) + current)
        end
        if current[2] > 1
            push!(to_visit, CartesianIndex(0, -1) + current)
        end
        if current[2] < height
            push!(to_visit, CartesianIndex(0, +1) + current)
        end
        heights[current] = -1
    end
    length(filter(==(-1), heights))
end

function solution9_2()
    input = collect(parse_input("9.data"))
    width = length(input[1])
    height = length(input)
    heights = reshape(reduce(vcat, input),(width,height))
    sizes = map(sink_locations(width, height, heights)) do sink
        size = basin_size(sink[1], sink[2], width, height, heights)
    end
    reduce(*, Iterators.take(reverse(sort(sizes)), 3))
end

solution9_2()