(** [parse_machine line_a line_b line_p] parse the machine description and
    returns a triple [(a, b, p)] describing the actions of the 'A' and 'B'
    buttons along with the location of the prize. *)
let parse_machine line_a line_b line_p =
  let get_xy re s =
    let _ = Str.search_forward re s 0 in
    ( int_of_string (Str.matched_group 1 s),
      int_of_string (Str.matched_group 2 s) )
  in
  let re_ab = Str.regexp {|Button [AB]: X\+\([0-9]+\), Y\+\([0-9]+\)|} in
  let re_pos = Str.regexp {|Prize: X=\([0-9]+\), Y=\([0-9]+\)|} in
  let a = get_xy re_ab line_a in
  let b = get_xy re_ab line_b in
  let prize = get_xy re_pos line_p in
  (a, b, prize)

(** [parse_machines lst] returns a list of the machines parsed from the input
    list of strings. *)
let parse_machines =
  let rec impl acc = function
    | "" :: t -> impl acc t
    | a :: b :: p :: t -> impl (parse_machine a b p :: acc) t
    | [] -> acc
    | _ -> failwith "parse_machines.impl"
  in
  impl []

(** [machines_of_file fname] returns the list of machines described in the file
    [fname]. *)
let machines_of_file fname = Aoc.strings_of_file fname |> parse_machines

(** [calc_tokens (a, b p)] calculates how many tokens are needed to get to [p]
    by pressing button A (moving [a] amount) and button B (moving [b] amount).
    Returns [None] if no solution possible, or [Some t] if the prize can be got
    by spending [t] tokens.

    Note the problem says minimize but if there is a solution there is only one
    solution. *)
let calc_tokens ((ax, ay), (bx, by), (x, y)) =
  (* Solve as a sequence of linear equations: 
     We want to find A & B in:
     A * ax + B * bx = X (1)
     A * ay + B * by = Y (2) 

     dividing (1) by bx and (2) by by gives us:

     A * ax / bx + B = X / bx (3)
     A * ay / by + B = Y / by (4)

     (3) - (4) gives:

     A * (ax / bx - ay / by) = X / bx - Y / by.

     Multiplying through by (bx * by) gives:

     A * (ax * by - ay * bx) = X * by - Y * bx.

     Dividing by (ax * by - ay * by) gives:

     A = (X * by - Y * bx) / (ax * by - ay * by)

     If ax * by - ay * by is 0 we have an infinite number of answers, and we'll
     deal with that if that becomes an issue (spoiler alert: it doesn't).

    A needs to be a whole number so (X * by - Y * bx) mod (ax * by - ay * by)
    must be 0, otherwise there is no solution. *)
  let a_n = (x * by) - (y * bx) in
  let a_d = (ax * by) - (ay * bx) in
  if a_d = 0 then None
  else if a_n mod a_d <> 0 then None
  else if a_n / a_d <= 0 then None
  else
    let a = a_n / a_d in
    let b = (x - (ax * a)) / bx in
    Some ((3 * a) + b)

(** [add_offset offset machine] offsets the prize location for [machine] by
    [(offset, offset)]. *)
let add_offset offset (a, b, (x, y)) = (a, b, (x + offset, y + offset))

(** [part offset machines] calculates the number of tokens needed to win as many
    prizes as possible from [machines]. All machines prizes are offset by
    [(offset, offset)]. *)
let part offset machines =
  List.map (add_offset offset) machines
  |> List.filter_map calc_tokens
  |> List.fold_left ( + ) 0

let _ =
  Aoc.main machines_of_file
    [ (string_of_int, part 0); (string_of_int, part 10000000000000) ]