open Angstrom
open Base
open Core
module Hashtbl = Base.Hashtbl

type card = { id : int; winning : int list; have : int list } [@@deriving sexp]

module Parse = struct
  (* Parse cards such as

     Card 1: 41 48 83 86 17 | 83 86  6 31 17  9 48 53
     Card 2: 13 32 20 16 61 | 61 30 68 82 17 32 24 19
     Card 3:  1 21 53 59 44 | 69 82 63 72 16 21 14  1
     Card 4: 41 92 73 84 69 | 59 84 76 51 58  5 54 83
     Card 5: 87 83 26 28 32 | 88 30 70 12 93 22 82 36
     Card 6: 31 18 13 56 72 | 74 77 10 23 35 67 36 11
  *)
  let intP =
    take_while1 (function '0' .. '9' -> true | _ -> false) >>| Int.of_string

  let int_listP = skip_many (char ' ') *> sep_by1 (skip_many1 (char ' ')) intP
  let card_headP = string "Card" *> skip_many1 (char ' ') *> intP <* string ": "
  let sepP = string " | "

  let card =
    let open Angstrom.Let_syntax in
    let%bind id = card_headP in
    let%bind winning = int_listP in
    let%bind _ = sepP in
    let%bind have = int_listP in
    return { id; winning; have }

  exception ParseExn of string

  let parse_card line =
    match parse_string ~consume:All card line with
    | Ok card -> card
    | Error e -> raise (ParseExn e)
end

module Part1 = struct
  let to_set ilist = Set.of_list (module Int) ilist

  let list_intersect a b =
    let sa = to_set a in
    let sb = to_set b in
    Set.inter sa sb

  let winning_numbers card = Set.length (list_intersect card.winning card.have)
  let card_value card = Int.pow 2 (winning_numbers card) / 2

  let _test_parse ch =
    let process l =
      let card = Parse.parse_card l in
      let sexp = Sexp.to_string_hum (sexp_of_card (Parse.parse_card l)) in
      let value = card_value card in
      Printf.sprintf "%s - value: %d" sexp value
    in
    let f l = Out_channel.print_endline (process l) in
    In_channel.iter_lines ch ~f

  let _solve ch =
    let f acc line =
      let card = Parse.parse_card line in
      acc + card_value card
    in
    let result = In_channel.fold_lines ch ~init:0 ~f in
    Out_channel.printf "Total value is %d\n" result
end

module Part2 = struct
  let make_int_hashmap () : (int, int) Hashtbl.t = Hashtbl.create (module Int)

  (* format int/int hashmap *)
  let _hm_to_string hm =
    Sexp.to_string_hum (Hashtbl.sexp_of_t Int.sexp_of_t Int.sexp_of_t hm)

  (* create a counting sequence *)
  let count_seq (from : int) (upto : int) : int Sequence.t =
    let f c =
      if Int.equal c upto then None
      else Some (c, if upto > from then c + 1 else c - 1)
    in
    Sequence.unfold ~init:from ~f

  (* for a card, account for all won cards. *)
  let process_card hm card =
    (* create default entry for card *)
    Hashtbl.update hm card.id ~f:(function None -> 1 | Some x -> x);
    (* how many points have we won? *)
    let winning = Part1.winning_numbers card in
    (* how many copies does the current card have? *)
    let copies_won = Option.value ~default:1 (Hashtbl.find hm card.id) in
    (* updater function adding `copies_won` to each successive card *)
    let updater card_id =
      Hashtbl.update hm card_id ~f:(function
        | None -> 1 + copies_won
        | Some x -> x + copies_won)
    in
    (* sequence of cards to update *)
    let copied_cards = count_seq (card.id + 1) (card.id + winning + 1) in
    Sequence.iter copied_cards ~f:updater

  (* sequentially process available cards. *)
  let process_cards ch =
    let hm = make_int_hashmap () in
    let f count line =
      process_card hm (Parse.parse_card line);
      count + 1
    in
    let count = In_channel.fold_lines ch ~init:0 ~f in
    (count, hm)

  (* count all scratchcards, up to the number of available cards. *)
  let total_scratchcards hm cardcount =
    let f ~key ~data sum = if key <= cardcount then sum + data else sum in
    Hashtbl.fold hm ~init:0 ~f

  (* print solution *)
  let _solve_process_cards ch =
    let count, hm = process_cards ch in
    Out_channel.printf "%d\n" (total_scratchcards hm count)
end

let () = Part2._solve_process_cards In_channel.stdin