open Printf;;
open Buffer;;
open Str;;

let read_whole_file f =
  let s = really_input_string f (in_channel_length f) in
  close_in f;
  s

let rec find_all_numbers pos str  =
  try
    let re = Str.regexp "-?[0-9]+" in
    let new_pos = Str.search_forward re str pos in
    let n_str = Str.matched_group 0 str in
    let n = int_of_string n_str in
    n :: find_all_numbers (new_pos + String.length n_str) str 
  with Not_found -> []

let take_3_values lst =
  match lst with
  | n1 :: n2 :: n3 :: tail ->
    ((n1, n2, n3), tail)
  | _ -> raise (Invalid_argument "Not enough elements")

let get_combo_operand (a, b, c) operand =
  match operand with
  | 4 -> a
  | 5 -> b
  | 6 -> c
  | n ->
    if n >= 0 && n < 4 then 
      n
    else
      raise (Invalid_argument "Invalid combo operand")

let read_instruction program program_counter =
  try 
    let opcode = List.nth program program_counter in
    let operand = List.nth program (program_counter + 1) in
    Some (opcode, operand)
  with Failure _ -> None

let rec cpu_tick program (a, b, c, program_counter) =
  match read_instruction program program_counter with
  | Some (opcode, operand) -> begin
    let next_program_counter = program_counter + 2 in
    match opcode with
    | 0 ->
      let result = Int.shift_right a (get_combo_operand (a, b, c) operand) in
      ((result, b, c, next_program_counter), None, false)
    | 1 ->
      let result = Int.logxor b operand in
      ((a, result, c, next_program_counter), None, false)
    | 2 ->
      let result = (get_combo_operand (a, b, c) operand) mod 8 in
      ((a, result, c, next_program_counter), None, false)
    | 3 ->
      if a = 0 then
        ((a, b, c, next_program_counter), None, false)
      else
        ((a, b, c, operand), None, false)
    | 4 ->
      let result = Int.logxor b c in
      ((a, result, c, next_program_counter), None, false)
    | 5 ->
      let out = (get_combo_operand (a, b, c) operand) mod 8 in
      ((a, b, c, next_program_counter), Some out, false)
    | 6 ->
      let result = Int.shift_right a (get_combo_operand (a, b, c) operand) in
      ((a, result, c, next_program_counter), None, false)
    | 7 ->
      let result = Int.shift_right a (get_combo_operand (a, b, c) operand) in
      ((a, b, result, next_program_counter), None, false)
    | _ -> raise (Invalid_argument "Invalid opcode")
    end
  | None -> ((a, b, c, program_counter), None, true)

let rec run_program_inner program state =
  match cpu_tick program state with
  | (_, _, true) -> []
  | (new_state, out, false) ->
    match out with
    | Some n ->
      n :: run_program_inner program new_state
    | None ->
      run_program_inner program new_state

let run_program (init_a, init_b, init_c) program =
  let result = run_program_inner program (init_a, init_b, init_c, 0) in
  result

let rec join_output lst =
  lst
  |> List.map string_of_int
  |> String.concat ","

let rec construct_a_consider_all_options reverse_mapping program possibilities prev_mask =
  match possibilities with
  | next :: possibilities_tail -> begin
    match construct_a_try_option reverse_mapping program (Some (Int.shift_right next 3)) prev_mask with
    | Some result -> begin
        Some (Int.logor (Int.logand next 7) (Int.shift_left result 3))
      end
    | None -> construct_a_consider_all_options reverse_mapping program possibilities_tail prev_mask
  end
  | _ -> None

and construct_a_try_option reverse_mapping program prev_opt prev_mask =
  match program with
  | n :: tail -> begin
    let this_possibilities = reverse_mapping.(n) in
    let restrict_possibilities_with_prev this_possibilities =
      match prev_opt with
      | Some prev ->
        List.filter (fun x -> (Int.logand x prev_mask) = prev) this_possibilities
      | None -> this_possibilities in
    let this_filtered = restrict_possibilities_with_prev this_possibilities in
    construct_a_consider_all_options reverse_mapping tail this_filtered prev_mask
  end
  | _ ->
    match prev_opt with
    | Some prev -> if prev = 0 then Some 0 else None
    | None -> Some 0

let get_initial_possibilities lookahead_bits =
  Int.shift_left 1 lookahead_bits

let rec get_prev_mask lookahead_bits =
  if lookahead_bits > 3 then
    Int.logor 1 (Int.shift_left (get_prev_mask (lookahead_bits - 1)) 1 )
  else
    0

(* The amount of lookahead bits needs to be at least equal to amount of bits of a that influences a single output digit *)
let construct_a_based_on_program program lookahead_bits =
  let initial_possibilities = get_initial_possibilities lookahead_bits in
  let prev_mask = get_prev_mask lookahead_bits in
  let reverse_mapping = Array.make 8 [] in
  Seq.ints 0
    |> Seq.take initial_possibilities
    |> Seq.map (fun a -> run_program (a, 0, 0) program |> List.hd)
    |> Seq.iteri (fun i a -> reverse_mapping.(a) <- List.append reverse_mapping.(a) [i]);
  construct_a_try_option reverse_mapping program None prev_mask

let () =
  let f = open_in_bin "input.txt" in
  let (registers, program) = f
    |> read_whole_file
    |> find_all_numbers 0
    |> take_3_values in
  let result = join_output @@ run_program registers program in
  printf "%s\n" result;
  let result2 = construct_a_based_on_program program 10 |> Option.get in
  printf "%d\n" result2;