module HittingSet

export HittingSetProblem, solve

HittingSetProblem{T} = Pair{Set{T}, Vector{Pair{T, Set{Set{T}}}}}

# `targets` should be a collection of Set objects
function HittingSetProblem(targets, chosen = Set())
  wholeset = union(targets...)
  T = eltype(wholeset)
  unsorted_moves = [
    elt => Set(filter(s -> elt ∉ s, targets))
    for elt in wholeset
  ]
  moves = sort(unsorted_moves, by = pair -> length(pair.second))
  Set{T}(chosen) => moves
end

function Base.display(problem::HittingSetProblem{T}) where T
  println("HittingSetProblem{$T}")
  
  chosen = problem.first
  println(" {", join(string.(chosen), ", "), "}")
  
  moves = problem.second
  for (choice, missed) in moves
    println(" | ", choice)
    for s in missed
      println(" | | {", join(string.(s), ", "), "}")
    end
  end
  println()
end

function solve(pblm::HittingSetProblem{T}, maxdepth = Inf) where T
  problems = Dict(pblm)
  while length(first(problems).first) < maxdepth
    subproblems = typeof(problems)()
    for (chosen, moves) in problems
      if isempty(moves)
        return chosen
      else
        for (choice, missed) in moves
          to_be_chosen = union(chosen, Set([choice]))
          if isempty(missed)
            return to_be_chosen
          elseif !haskey(subproblems, to_be_chosen)
            push!(subproblems, HittingSetProblem(missed, to_be_chosen))
          end
        end
      end
    end
    problems = subproblems
  end
  problems
end

function test(n = 1)
  T = [Int64, Int64, Symbol, Symbol][n]
  targets = Set{T}.([
    [
      [1, 3, 5],
      [2, 3, 4],
      [1, 4],
      [2, 3, 4, 5],
      [4, 5]
    ],
    # example from Amit Chakrabarti's graduate-level algorithms class (CS 105)
    # notes by Valika K. Wan and Khanh Do Ba, Winter 2005
    # https://www.cs.dartmouth.edu/~ac/Teach/CS105-Winter05/
    [
      [1, 3], [1,    4], [1,    5],
      [1, 3], [1, 2, 4], [1, 2, 5],
      [4, 3], [   2, 4], [   2, 5],
      [6, 3], [6,    4], [      5]
    ],
    [
      [:w, :x, :y],
      [:x, :y, :z],
      [:w, :z],
      [:x, :y]
    ],
    # Wikipedia showcases this as an example of a problem where the greedy
    # algorithm performs especially poorly
    [
      [:a, :x, :t1],
      [:a, :y, :t2],
      [:a, :y, :t3],
      [:a, :z, :t4],
      [:a, :z, :t5],
      [:a, :z, :t6],
      [:a, :z, :t7],
      [:b, :x, :t8],
      [:b, :y, :t9],
      [:b, :y, :t10],
      [:b, :z, :t11],
      [:b, :z, :t12],
      [:b, :z, :t13],
      [:b, :z, :t14]
    ]
  ][n])
  problem = HittingSetProblem(targets)
  if isa(problem, HittingSetProblem{T})
    println("Correct type")
  else
    println("Wrong type: ", typeof(problem))
  end
  problem
end

end