#!/usr/bin/env julia

#= --- Part Two ---

Next, you need to find the largest basins so you know what areas are most important to avoid.

A basin is all locations that eventually flow downward to a single low point. Therefore, every low point has a basin, although some basins are very small. Locations of height 9 do not count as being in any basin, and all other locations will always be part of exactly one basin.

The size of a basin is the number of locations within the basin, including the low point. The example above has four basins.

The top-left basin, size 3:

2199943210
3987894921
9856789892
8767896789
9899965678

The top-right basin, size 9:

2199943210
3987894921
9856789892
8767896789
9899965678

The middle basin, size 14:

2199943210
3987894921
9856789892
8767896789
9899965678

The bottom-right basin, size 9:

2199943210
3987894921
9856789892
8767896789
9899965678

Find the three largest basins and multiply their sizes together. In the above example, this is 9 * 14 * 9 = 1134.

What do you get if you multiply together the sizes of the three largest basins?
=#

infile = length(ARGS) > 0 ? ARGS[1] : "input.txt"
println("infile = ", infile)

depth_map = reduce(vcat, [parse.(Int, split(line, ""))'
                          for line in  eachline(infile)])

function neighbor_idx(nrows, ncols, y, x)
    v = CartesianIndex[]
    if x > 1
        push!(v, CartesianIndex(y, x-1))
    end
    if x < ncols
        push!(v, CartesianIndex(y, x+1))
    end
    if y > 1
        push!(v, CartesianIndex(y-1, x))
    end
    if y < nrows
        push!(v, CartesianIndex(y+1, x))
    end
    return v
end


function neighbors(m, y, x)
    nrows, ncols = size(m)
    v = eltype(m)[]
    if x > 1
        push!(v, m[y, x-1])
    end
    if x < ncols
        push!(v, m[y, x+1])
    end
    if y > 1
        push!(v, m[y-1, x])
    end
    if y < nrows
        push!(v, m[y+1, x])
    end
    return v
end

function flow_neighbors(m, y, x, flow_map, flow_count)
    nrows, ncols = size(m)
    nidx = neighbor_idx(nrows, ncols, y, x)
    v = m[y, x]
    if v == 9
        flow_count[y, x] = 0
        return
    end
    smallest_v = v
    flow_idx = flow_map[y, x]
    for idx in nidx
        if m[idx] < smallest_v
            smallest_v = m[idx]
            flow_idx = flow_map[idx]
            while flow_map[flow_idx] != flow_idx
                flow_idx = flow_map[flow_idx]
            end
        end
    end
    if smallest_v != v
        flow_count[flow_idx] += flow_count[y, x]
        flow_count[y, x] = 0
        flow_map[y, x] = flow_idx
    end
    #println(y, "x", x)
    #display(flow_count)
    #println()
end

function flow_all(m)
    total = 0
    nrows, ncols = size(m)
    flow_count = ones(Int, nrows, ncols)
    flow_map = reshape([CartesianIndex(y, x)
                        for x in 1:ncols
                        for y in 1:nrows], nrows, ncols)
    for x in 1:ncols
        for y in 1:nrows
            flow_neighbors(m, y, x, flow_map, flow_count)
        end
    end
    return flow_map, flow_count
end

map, count = flow_all(depth_map)

if (length(depth_map) < 100)
    display(count)
    println()
end

basins = sort(filter(x -> x > 0, count))
display(basins)
println()

println("nbasins = ", length(basins))
println("largest_prod = ", prod(basins[end-2:end]))