let next_secret secret =
  let secret = secret * 64 lxor secret mod 16777216 in
  let secret = secret / 32 lxor secret mod 16777216 in
  let secret = secret * 2048 lxor secret mod 16777216 in
  secret

let nth_secret n secret =
  let r = Aoc.apply_n n next_secret secret in
  r

let part1 n nums = List.map (nth_secret n) nums |> List.fold_left ( + ) 0

module Int4Tuple = struct
  type t = int * int * int * int

  let compare = Stdlib.compare
end

module Int4Map = Map.Make (Int4Tuple)

let secret_list n secret =
  let rec impl s () = Seq.Cons (s, impl (next_secret s)) in
  Seq.drop 1 (impl secret) |> Seq.take n |> List.of_seq

let rec find_sequence_costs map =
  let update_value amt = function None -> Some amt | x -> x in
  function
  | a :: b :: c :: d :: e :: t ->
      find_sequence_costs
        (Int4Map.update (b - a, c - b, d - c, e - d) (update_value e) map)
        (b :: c :: d :: e :: t)
  | _ -> map

let part2 n secrets =
  let merge_values _ x y =
    match (x, y) with
    | Some x, Some y -> Some (x + y)
    | Some x, None -> Some x
    | None, Some y -> Some y
    | None, None -> None
  in
  let costs =
    List.map (secret_list n) secrets
    |> List.map (List.map (fun x -> x mod 10))
    |> List.map (find_sequence_costs Int4Map.empty)
    |> List.fold_left
         (fun acc map -> Int4Map.merge merge_values acc map)
         Int4Map.empty
  in
  Int4Map.fold (fun _ v acc -> max acc v) costs 0

let read_file fname = Aoc.strings_of_file fname |> List.map int_of_string

let _ =
  Aoc.main read_file
    [ (string_of_int, part1 2000); (string_of_int, part2 2000) ]

(* 424 too low *)