(* Common util. Should probably put this in a module or something.. *)
let get_one_line file =
  try Some (input_line file) with
    End_of_file -> None

let get_lines file =
  let rec input lines =
    match get_one_line file with
      Some line -> input (line :: lines)
    | None -> List.rev lines
  in input []

let read_file file =
  let channel = open_in(file) in
  get_lines channel

let explode s = (* convert a string to a list of chars *)
  let rec exp i l =
   if i < 0 then l else exp (i - 1) (s.[i] :: l) in
  exp (String.length s - 1) [];;

let implode l = String.of_seq (List.to_seq l)

let (<<) f g x = f(g(x))
let (>>) f g x = g(f(x))

let abs x = if x < 0 then -x else x

(* --- *)

type instruction = Mask of string | Mem of int * int (* Mem (addr, value) *)

let parse_instruction : string -> instruction =
  String.split_on_char '='
  >> List.map String.trim
  >> fun [left; right] ->
    if left = "mask"
    then Mask right
    else
      let len = String.length left in
      let addr = String.sub left 4 (len - 5) in
      Mem (int_of_string addr, int_of_string right)

let rec pow a = function
  | 0 -> 1
  | 1 -> a
  | n ->
      let b = pow a (n / 2) in
      b * b * (if n mod 2 = 0 then 1 else a)

let bin2dec =
  List.rev
  >> List.mapi (fun i d -> if d = 1 then pow 2 i else 0)
  >> List.fold_left (+) 0

let int_of_char = String.make 1 >> int_of_string

let run_mask mask num = (* 'and' the 0s and 'or' the 1s *)
  let l = explode mask in
  let ones = List.map (fun c -> if c = '1' then '1' else '0') l (* copy the mask but just the ones *)
    |> List.map int_of_char |> bin2dec in (* convert from char list to int list to decimal int *)
  let zeros = List.map (fun c -> if c = '0' then '1' else '0') l (* same for zeros but make all the zeros 1s *)
    |> List.map int_of_char |> bin2dec in
  (num lor ones) - (num land zeros) (* subtract! *)

let change_bit number position new_value =
  let mask = 1 lsl position in
  (number land (lnot mask)) lor ((new_value lsl position) land mask)

let rec find_addrs mask addr = (* return a list of all addresses after applying mask with floaty bit *)
  let mask' = List.rev @@ explode mask in
  Printf.printf "\nnew mask! %s\n" mask;

  let rec go addr' idx =
    if idx = List.length mask'
    then [addr']
    else
      let cur = List.nth mask' idx in

      if cur = '0'
      then go addr' (idx + 1)
      else
        if cur = '1'
        then go (change_bit addr' idx 1) (idx + 1)
        else (* X *)
          let addr0 = (change_bit addr' idx 0) in
          let addr1 = (change_bit addr' idx 1) in
          List.concat [go addr0 (idx + 1); go addr1 (idx + 1)]
  in go addr 0

let hash = Hashtbl.create 1000
let hash' = Hashtbl.create 1000

let rec walk mask = function
  | [] -> ()
  | x :: xs ->
      match x with
      | Mask m -> walk m xs (* update the mask and recurse *)
      | Mem (addr, value) ->
          (* problem 1 *)
          run_mask mask value
          |> Hashtbl.replace hash addr;

          (* problem 2 *)
          find_addrs mask addr
          |> List.iter (fun addr' -> Hashtbl.replace hash' addr' value;);

          walk mask xs (* recurse *)

let () =
  let lines = read_file "day14.txt" |> List.map parse_instruction in

  walk "" lines;

  let sum = Hashtbl.fold (fun k v acc -> acc + v) hash 0 in
  Printf.printf "sum: %i\n" sum;

  let sum' = Hashtbl.fold (fun k v acc -> acc + v) hash' 0 in
  Printf.printf "sum': %i\n" sum';