module StringMap = Map.Make (String)
module StringSet = Set.Make (String)

let add_connection map (a, b) =
  let update s = function
    | None -> Some (StringSet.add s StringSet.empty)
    | Some set -> Some (StringSet.add s set)
  in
  let map = StringMap.update a (update b) map in
  let map = StringMap.update b (update a) map in
  map

let make_pairs = function
  | [ a; b ] -> (a, b)
  | _ -> raise (invalid_arg "make_pairs")

let load_file fname =
  Aoc.strings_of_file fname
  |> List.map (String.split_on_char '-')
  |> List.map make_pairs
  |> List.fold_left add_connection StringMap.empty

let rec find_second_member acc connections visited a candidates =
  let rec impl acc set = function
    | [] -> acc
    | c :: t -> impl (StringSet.add c set :: acc) set t
  in
  match StringSet.choose_opt candidates with
  | None -> acc
  | Some h ->
      let candidates = StringSet.remove h candidates in
      if StringSet.mem h visited then
        find_second_member acc connections visited a candidates
      else
        let visited = StringSet.add h visited in
        let anh = StringSet.inter (StringMap.find h connections) candidates in
        let acc =
          impl acc (StringSet.of_list [ a; h ]) (StringSet.to_list anh)
        in
        find_second_member acc connections visited a candidates

let rec find_rings acc visited connections = function
  | [] -> acc
  | h :: t ->
      if StringSet.mem h visited then find_rings acc visited connections t
      else
        let visited = StringSet.add h visited in
        let acc =
          find_second_member acc connections visited h
            (StringMap.find h connections)
        in
        find_rings acc visited connections t

(** [starts_with_t set] returns true if any member of [set] starts with the
    letter ['t']. *)
let starts_with_t set = StringSet.exists (fun x -> x.[0] = 't') set

let part1 connections =
  StringMap.to_list connections
  |> List.map fst
  |> find_rings [] StringSet.empty connections
  |> List.filter starts_with_t |> List.length |> string_of_int

(** [search_candidate max_lst connections current candidates] searches for the
    longest star network that contains all the computers in [current] and maybe
    some of the elements of [candidates].

    [max_lst] is the biggest network found so far. [connections] is the map of
    computers to direct connections.

    The initial call should look something like:

    [search_candidate [] connections k (StringSet.to_list (StringMap.find k
     connections))] *)
let rec search_candidate max_lst connections current candidates =
  match candidates with
  | [] -> if List.length current > List.length max_lst then current else max_lst
  | h :: t ->
      let map = StringMap.find h connections in
      let current' = h :: current in
      let candidates' = List.filter (Fun.flip StringSet.mem map) candidates in
      let max_lst = search_candidate max_lst connections current' candidates' in
      search_candidate max_lst connections current t

(** [find_max_set connections] returns a list containing the largest number of
    computers in a star network (that is for every pair of elements in the list
    there is a connection between them).

    [connections] is the map of connections keyed by computer with the value
    being a set of all direct connections. [connections] must be bi-directional
    that is if: [StringSet.mem a (StringMap.find b connections)] then
    [StringSet.mem b (StringMap.find a connections)]. Note that
    [StringSet.mem a (StringMap.find a connections)] must return [false]. *)
let find_max_set connections =
  let rec impl max_lst = function
    | [] -> max_lst
    | (k, v) :: t ->
        impl
          (search_candidate max_lst connections [ k ] (StringSet.to_list v))
          t
  in
  impl [] (StringMap.to_list connections)

let part2 connections =
  find_max_set connections |> List.sort compare |> String.concat ","

let _ = Aoc.main load_file [ (Fun.id, part1); (Fun.id, part2) ]