open Day7_syntax

let parse_lexbuf = Day7_parser.rules Day7_lexer.token

let parse_file file = Bracket.infile file
  ~act:(fun chan -> Lexing.from_channel chan |> parse_lexbuf)


module BagOrd = struct type t = bag [@@deriving ord] end
module BagGraph = Graph.Make(BagOrd)
module BagSet = Set.Make(BagOrd)
module BagMap = Map.Make(BagOrd)

let make_graph edges =
  let flatten (l1, ps) = List.map (fun (l2, i) -> l1, l2, i) ps in
  BagGraph.make_e (List.concat_map flatten edges)

let graph_from_file file = parse_file file |> make_graph

let search_from wanted graph =
  let seen = ref BagSet.empty in
  let add _ vout _ = if vout <> wanted then seen := BagSet.add vout !seen in
  BagGraph.iter_dfs add graph wanted;
  !seen

let count_from wanted graph =
  let totals = ref BagMap.empty in
  let find0 bag = BagMap.find_opt bag !totals |> Option.value ~default:0 in
  let put bag amount = totals := BagMap.add bag (amount + find0 bag) !totals in
  let add vin vout e = put vin (e * (find0 vout + 1)) in
  BagGraph.iter_dfs add graph wanted;
  !totals


let print_map map =
  let cmp_snd (_, x) (_, y) = Int.compare x y in
  let print1 (bag, x) = Format.printf "%s: %d\n" bag x in
  BagMap.to_seq map |> List.of_seq |> List.sort cmp_snd |> List.iter print1


let%test_module _ = (module struct
  let parse_string str = Lexing.from_string str |> parse_lexbuf

  let example1 = {|
    light red bags contain 1 bright white bag, 2 muted yellow bags.
    dark orange bags contain 3 bright white bags, 4 muted yellow bags.
    bright white bags contain 1 shiny gold bag.
    muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
    shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
    dark olive bags contain 3 faded blue bags, 4 dotted black bags.
    vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
    faded blue bags contain no other bags.
    dotted black bags contain no other bags.
  |}

  let example_graph1 = parse_string example1 |> make_graph
  let example_graph1_rev = BagGraph.rev example_graph1

  let%expect_test _ =
    search_from "shiny gold" example_graph1_rev |> BagSet.iter print_endline;
    [%expect{|
      bright white
      dark orange
      light red
      muted yellow |}]

  let%expect_test _ =
    count_from "shiny gold" example_graph1 |> print_map;
    [%expect{|
      dark olive: 7
      vibrant plum: 11
      shiny gold: 32 |}]


  let example2 = {|
    shiny gold bags contain 2 dark red bags.
    dark red bags contain 2 dark orange bags.
    dark orange bags contain 2 dark yellow bags.
    dark yellow bags contain 2 dark green bags.
    dark green bags contain 2 dark blue bags.
    dark blue bags contain 2 dark violet bags.
    dark violet bags contain no other bags.
  |}

  let example_graph2 = parse_string example2 |> make_graph

  let%expect_test _ =
    count_from "shiny gold" example_graph2 |> print_map;
    [%expect{|
      dark blue: 2
      dark green: 6
      dark yellow: 14
      dark orange: 30
      dark red: 62
      shiny gold: 126 |}]
end)


let main1 file =
  let graph = graph_from_file file |> BagGraph.rev in
  let seen = search_from "shiny gold" graph in
  Format.printf "%d seen out of %d\n"
    (BagSet.cardinal seen) (BagGraph.count_vertices graph)

let%expect_test "part 1" =
  main1 "../../data/day7";
  [%expect{| 278 seen out of 2988 |}]


let main2 file =
  graph_from_file file |> count_from "shiny gold" |> print_map

let%expect_test "part 2" =
  main2 "../../data/day7";
  [%expect{|
    muted chartreuse: 1
    dull lime: 4
    plaid turquoise: 4
    dull lavender: 6
    light lavender: 10
    shiny maroon: 15
    pale green: 40
    pale gold: 67
    faded plum: 97
    drab beige: 272
    mirrored turquoise: 1467
    dim lavender: 7482
    dark coral: 22577
    shiny gold: 45157 |}]

let main = Misc.main 7 [|main1; main2|]