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

type card = int [@@deriving sexp]
type hand = card list [@@deriving sexp]
type bid = { hand : hand; win : int } [@@deriving sexp]

let max_card = 14
let compare_card = Int.compare
let _string_of_bids bids = Sexp.to_string_hum @@ List.sexp_of_t sexp_of_bid bids

module Parse = struct
  let maybe p = option () (p >>= fun _ -> return ())
  let intP = take_while1 Char.is_digit >>| Int.of_string

  let cardP : card Angstrom.t =
    choice
      [
        satisfy Char.is_digit >>| String.of_char >>| Int.of_string;
        char 'T' *> return 10;
        char 'J' *> return 11;
        char 'Q' *> return 12;
        char 'K' *> return 13;
        char 'A' *> return 14;
      ]

  let handP : hand Angstrom.t = count 5 cardP

  let bidP =
    let open Angstrom.Let_syntax in
    let%bind hand = handP in
    let%bind _ = char ' ' in
    let%bind win = intP in
    let%bind _ = maybe (char '\n') in
    return { hand; win }

  let bidsP = many1 bidP

  exception Parse_exn of string

  let parse_input str =
    match parse_string ~consume:All bidsP str with
    | Ok ok -> ok
    | Error e -> raise (Parse_exn e)
end

type typ =
  | FiveOf
  | FourOf
  | FullHouse
  | ThreeOf
  | TwoPair
  | OnePair
  | High
  | NoType
[@@deriving sexp]

let int_of_typ = function
  | FiveOf -> 8
  | FourOf -> 7
  | FullHouse -> 6
  | ThreeOf -> 5
  | TwoPair -> 4
  | OnePair -> 3
  | High -> 2
  | NoType -> 0

let compare_typ a b = Int.compare (int_of_typ a) (int_of_typ b)

let _string_of_typ_bids bids =
  Sexp.to_string_hum
  @@ Sexp.List
       (List.map
          ~f:(fun (typ, bid) -> Sexp.List [ sexp_of_typ typ; sexp_of_bid bid ])
          bids)

module Part1 = struct
  let count_cards cards =
    let counts = Array.create ~len:(max_card + 1) 0 in
    let count c = counts.(c) <- counts.(c) + 1 in
    List.iter ~f:count cards;
    counts

  let typ_of counts =
    Array.sort ~compare:(fun a b -> Int.compare b a) counts;
    match Array.to_list counts with
    | 5 :: _ -> FiveOf
    | 4 :: _ -> FourOf
    | 3 :: 2 :: _ -> FullHouse
    | 3 :: 1 :: 1 :: _ -> ThreeOf
    | 2 :: 2 :: 1 :: _ -> TwoPair
    | 2 :: 1 :: 1 :: 1 :: _ -> OnePair
    | 1 :: 1 :: 1 :: 1 :: 1 :: _ -> High
    | _ -> NoType

  let typ_of_hand h =
    let typ = typ_of (count_cards h) in
    typ

  let compare_hands a b =
    let ta, tb = (typ_of_hand a, typ_of_hand b) in
    match compare_typ ta tb with
    | -1 -> -1
    | 1 -> 1
    | 0 -> List.compare compare_card a b
    | _ -> assert false

  let compare_bids { hand = a; _ } { hand = b; _ } = compare_hands a b
  let sort_bids = List.sort ~compare:(fun a b -> compare_bids a b)

  let _winnings bids =
    let sorted = sort_bids bids in
    let f ix ac bid = ac + ((ix + 1) * bid.win) in
    List.foldi ~f ~init:0 sorted
end

module Part2 = struct
  let convert_part1_hand { hand; win } =
    let f = function 11 -> 1 | x -> x in
    { hand = List.map ~f hand; win }

  type cardcount = { card : int; count : int } [@@deriving sexp]

  let compare_cardcount { count = count1; _ } { count = count2; _ } =
    Int.compare count1 count2

  let count_cards cards =
    let cardcounts =
      Sequence.unfold ~init:0 ~f:(fun s ->
          if s <= max_card + 1 then Some ({ card = s; count = 0 }, s + 1)
          else None)
    in
    let counts = Sequence.to_array cardcounts in
    let count c =
      counts.(c) <- { (counts.(c)) with count = counts.(c).count + 1 }
    in
    List.iter ~f:count cards;
    counts

  let wildcard = 1

  let wildcards =
    List.fold ~init:0 ~f:(fun a { card; count } ->
        if Int.equal card wildcard then a + count else a)

  let typ_of counts =
    Array.sort ~compare:(fun a b -> compare_cardcount b a) counts;
    let countsl = Array.to_list counts in
    let nwild = wildcards countsl in
    let countsl =
      List.filter ~f:(fun { card; _ } -> not (Int.equal card wildcard)) countsl
    in
    (*Out_channel.printf "counts %s\n" (Sexp.to_string_hum (List.sexp_of_t sexp_of_cardcount countsl));
    *)
    match countsl with
    | { count = x; _ } :: _ when Int.(nwild + x = 5) -> FiveOf
    | { count = x; _ } :: { count = y; _ } :: _
      when Int.(nwild + x = 4 || (x = 4 && nwild + y = 1)) ->
        FourOf
    | { count = x; _ } :: { count = y; _ } :: _
      when Int.((nwild + x = 3 && y = 2) || (nwild + y = 2 && x = 3)) ->
        FullHouse
    | { count = x; _ } :: { count = 1; _ } :: { count = z; _ } :: _
      when Int.((nwild + x = 3 && z = 1) || (x = 3 && nwild + z = 1)) ->
        ThreeOf
    | { count = 2; _ } :: { count = y; _ } :: { count = z; _ } :: _
      when Int.((nwild + y = 2 && z = 1) || (y = 2 && nwild + z = 1)) ->
        TwoPair
    | { count = x; _ }
      :: { count = y; _ }
      :: { count = z; _ }
      :: { count = a; _ }
      :: _
      when Int.(
             (x + nwild = 2 && y = 1 && z = 1 && a = 1)
             || (x = 2 && y = 1 && z = 1 && a + nwild = 1)) ->
        OnePair
    | { count = 1; _ }
      :: { count = 1; _ }
      :: { count = 1; _ }
      :: { count = 1; _ }
      :: { count = x; _ }
      :: _
      when Int.(nwild + x = 1) ->
        High
    | _ -> assert false

  let typ_of_hand h =
    (*Out_channel.printf "hand %s -> " (Sexp.to_string_hum (List.sexp_of_t Int.sexp_of_t h));
    *)
    let typ = typ_of (count_cards h) in
    (*Out_channel.printf "%s -> %s\n"
      (Sexp.to_string_hum (sexp_of_typ typ))
      (Sexp.to_string_hum @@ List.sexp_of_t Int.sexp_of_t h);
    *)
    typ

  let compare_hands a b =
    let ta, tb = (typ_of_hand a, typ_of_hand b) in
    match compare_typ ta tb with
    | -1 -> -1
    | 1 -> 1
    | 0 -> List.compare compare_card a b
    | _ -> assert false

  let sort_bids bids =
    let f bid = (typ_of_hand bid.hand, bid) in
    let typ_hands = List.map ~f bids in
    List.sort
      ~compare:(fun (_, a) (_, b) -> compare_hands a.hand b.hand)
      typ_hands

  let winnings bids =
    let sorted = sort_bids @@ List.map ~f:convert_part1_hand bids in
    Out_channel.print_endline (_string_of_typ_bids sorted);
    let f ix ac (_, bid) = ac + ((ix + 1) * bid.win) in
    List.foldi ~f ~init:0 sorted
end

let () =
  let inp = In_channel.input_all In_channel.stdin in
  let bids = Parse.parse_input inp in
  Out_channel.printf "Total winnings: %d\n" (Part2.winnings bids)