open Printf;;

let rec list_of_lines in_file =
  try
    let line = input_line in_file in
    line :: list_of_lines(in_file)
  with End_of_file ->
    close_in in_file;
    []

let parse_filesystem_block index length =
  if index mod 2 == 0 then
    Array.init length (fun x -> index / 2)
  else
    Array.init length (fun x -> (-1))

let digit_code_to_int c =
  int_of_char c - int_of_char '0'

let rec fix_file_layout_inner front_index back_index file_layout =
  if front_index >= back_index then
    ()
  else if file_layout.(front_index) >= 0 then
    fix_file_layout_inner (front_index + 1) back_index file_layout
  else if file_layout.(back_index) < 0 then
    fix_file_layout_inner front_index (back_index - 1) file_layout
  else begin
    file_layout.(front_index) <- file_layout.(back_index);
    file_layout.(back_index) <- (-1);
    fix_file_layout_inner (front_index + 1) (back_index - 1) file_layout
  end

let fix_file_layout file_layout =
  let array_copy = Array.copy file_layout in
  fix_file_layout_inner 0 (Array.length array_copy - 1) array_copy;
  array_copy

let rec get_repeated_count array value index step =
  try 
    if array.(index) != value then
      0
    else
      1 + get_repeated_count array value (index + step) step
  with Invalid_argument _ -> 0

let rec try_move_file_no_fragment file_id file_index file_length file_layout front_index =
  if front_index >= file_index then
    ()
  else if file_layout.(front_index) >= 0 then
    try_move_file_no_fragment file_id file_index file_length file_layout (front_index + 1)
  else
    let free_space_length = get_repeated_count file_layout (-1) front_index 1 in
    if free_space_length < file_length then
      try_move_file_no_fragment file_id file_index file_length file_layout (front_index + free_space_length)
    else begin
      Seq.ints 0
        |> Seq.take file_length
        |> Seq.iter (
          fun i ->
            file_layout.(front_index + i) <- file_id;
            file_layout.(file_index + i) <- (-1));
    end

let rec fix_file_layout_defragment_inner back_index file_layout =
  if back_index < 0 then
    ()
  else if file_layout.(back_index) < 0 then
    let free_space_length = get_repeated_count file_layout (-1) back_index (-1) in
    fix_file_layout_defragment_inner (back_index - free_space_length) file_layout
  else
    let file_id = file_layout.(back_index) in
    let file_length = get_repeated_count file_layout file_id back_index (-1) in
    let new_back_index = back_index - file_length in
    try_move_file_no_fragment file_id (new_back_index + 1) file_length file_layout 0;
    fix_file_layout_defragment_inner new_back_index file_layout
    

let fix_file_layout_defragment file_layout = 
  let array_copy = Array.copy file_layout in
  fix_file_layout_defragment_inner (Array.length file_layout - 1) array_copy;
  array_copy

let get_checksum file_layout =
  file_layout
    |> Array.to_seq
    |> Seq.mapi (fun i x -> i * x)
    |> Seq.filter (fun x -> x >= 0)
    |> Seq.fold_left Int.add 0

let () =
  let f = open_in "input.txt" in
  let disk_map = f 
    |> list_of_lines
    |> List.hd
    |> String.to_seq
    |> Seq.map digit_code_to_int
    |> List.of_seq in
  let file_layout = disk_map
    |> List.mapi (fun i x -> parse_filesystem_block i x)
    |> Array.concat in
  let fixed_layout = fix_file_layout file_layout in
  let result = get_checksum fixed_layout in
  printf "%d\n%!" result;
  let fixed_layout_defragment = fix_file_layout_defragment file_layout in
  let result2 = get_checksum fixed_layout_defragment in
  printf "%d\n"result2