open Angstrom
open Base
open Core

type tile = Round | Cube | Empty [@@deriving eq, sexp]

let string_of_tile = function Round -> "O" | Cube -> "#" | Empty -> "."

type field = { rows : int; cols : int; fields : tile array }
[@@deriving sexp, eq]

let rc_to_ix field r c = (r * field.cols) + c
let field_get field r c = field.fields.(rc_to_ix field r c)
let field_set field r c v = field.fields.(rc_to_ix field r c) <- v

let field_column field c =
  Array.init field.rows ~f:(fun r -> field_get field r c)

let field_row field r =
  Array.sub field.fields ~pos:(r * field.cols) ~len:field.cols

let field_of_tiles r c tiles =
  { rows = r; cols = c; fields = Array.of_list tiles }

let string_of_field field =
  let rows = Sequence.range 0 field.rows in
  let row_strings =
    Sequence.map rows ~f:(fun r ->
        let row = field_row field r in
        Array.map row ~f:string_of_tile |> Array.to_list |> String.concat)
  in
  Sequence.to_list row_strings |> String.concat ~sep:"\n"

let print_field ?(ch = Out_channel.stdout) field =
  Out_channel.(
    output_string ch (string_of_field field);
    output_string ch "\n\n")

module Parse = struct
  let parse_tile = function
    | 'O' -> Round
    | '#' -> Cube
    | '.' -> Empty
    | _ -> failwith "invalid tile"

  let parse_row s = String.to_list s |> List.map ~f:parse_tile

  let parse_field s =
    let lines = String.split_lines s in
    let rows = List.length lines and cols = String.length (List.hd_exn lines) in
    let tile_rows = List.map lines ~f:parse_row in
    let tiles = List.concat tile_rows in
    field_of_tiles rows cols tiles
end

module Part1 = struct
  type load_state = { sum : int; last_stop : int; since : int }

  let calculate_column_load rows col =
    let f ix ({ sum; last_stop; since } as a) = function
      | Empty -> a
      | Cube -> { a with last_stop = ix; since = 0 }
      | Round ->
          {
            a with
            sum = sum + (rows - last_stop - since - 1);
            since = since + 1;
          }
    in
    let result =
      Array.foldi col ~init:{ sum = 0; last_stop = -1; since = 0 } ~f
    in
    result.sum

  let calculate_column_loads field =
    Array.init field.cols ~f:(fun c ->
        calculate_column_load field.rows (field_column field c))

  let calculate_load field =
    let column_loads = calculate_column_loads field in
    Array.fold column_loads ~init:0 ~f:Int.( + )
end

module Part2 = struct
  (* with round boulders already aligned, we can use a simplified algorithm. *)
  let simplified_column_load field col =
    let f ix sum item =
      match item with Round -> sum + (field.rows - ix) | _ -> sum
    in
    Array.foldi (field_column field col) ~init:0 ~f

  let calculate_column_loads field =
    let cols = Sequence.range 0 field.cols in
    Sequence.map cols ~f:(simplified_column_load field) |> Sequence.to_array

  (* Calculate load on northern pole *)
  let simplified_load field =
    let column_loads = calculate_column_loads field in
    Array.fold column_loads ~init:0 ~f:Int.( + )

  type state = { last_stop : int }

  (* shift a row or column rightwards, in-place. The getter/setter functions
     determine whether a row or column is shifted, and in which direction *)
  let shifter (getter : int -> tile) (setter : int -> tile -> unit) (max : int)
      =
    let f st ix =
      match getter ix with
      | Empty -> st
      | Cube -> { last_stop = ix }
      | Round ->
          setter ix Empty;
          setter (st.last_stop + 1) Round;
          { last_stop = st.last_stop + 1 }
    in
    let ixs = Sequence.range 0 max in
    let _ = Sequence.fold ixs ~init:{ last_stop = -1 } ~f in
    ()

  let shift_row_left field r =
    let getter ix = field_get field r ix and setter ix = field_set field r ix in
    shifter getter setter field.cols

  let shift_left field =
    let rows = Sequence.range 0 field.rows in
    Sequence.iter rows ~f:(shift_row_left field)

  let shift_row_right field r =
    let getter ix = field_get field r (field.cols - 1 - ix)
    and setter ix = field_set field r (field.cols - 1 - ix) in
    shifter getter setter field.cols

  let shift_right field =
    let rows = Sequence.range 0 field.rows in
    Sequence.iter rows ~f:(shift_row_right field)

  let shift_column_up field c =
    let getter ix = field_get field ix c and setter ix = field_set field ix c in
    shifter getter setter field.rows

  let shift_up field =
    let cols = Sequence.range 0 field.cols in
    Sequence.iter cols ~f:(shift_column_up field)

  let shift_column_down field c =
    let getter ix = field_get field (field.rows - 1 - ix) c
    and setter ix = field_set field (field.rows - 1 - ix) c in
    shifter getter setter field.rows

  let shift_down field =
    let cols = Sequence.range 0 field.cols in
    Sequence.iter cols ~f:(shift_column_down field)

  (* Run one cycle of the field. *)
  let cycle field =
    shift_up field;
    shift_left field;
    shift_down field;
    shift_right field

  (* Run the cycle function n times, printing the field after each 1000 completed cycles. *)
  let cycle_n field n =
    let rng = Sequence.range 0 n in
    Sequence.iter rng ~f:(fun i ->
        cycle field;
        if Int.(i < 10) then print_field field;
        (* opportunity for stupid off-by-one error: if i = 9, then 10 cycles are completed. *)
        if Int.((i + 1) % 1000 = 0) then
          Out_channel.printf "load? %d -> %d\n" i (simplified_load field);
        Out_channel.flush stdout)

  (* idea: the load numbers repeat cyclically. For the test input, with a period of 7.
     for my puzzle input, with a period of 9, when checked after every 1000 completed cycles.
     So we only need to observe the load number cycle, and then calculate back from
     the desired cycle number 1e9 back to the load number expected there:

        load_number = load_sequence[ (1e9 / 1000) * period ]

      For 1e9, the index is = 1, thus the solution is 84238. For the test input, similarly, the index is 1
      and the solution is 64.
  *)
  let test_load_sequence = [ 65; 64; 65; 69; 69; 68; 63 ]

  let load_sequence =
    [ 84239; 84328; 84202; 84342; 84210; 84332; 84237; 84299; 84276 ]
end

let () =
  let inp = In_channel.(input_all stdin) in
  let field = Parse.parse_field inp in
  let load = Part1.calculate_load field in
  Out_channel.printf "load: %d\n" load;
  print_field field;
  Part2.cycle_n field 1000000;
  let load2 = Part2.simplified_load field in
  print_field field;
  Out_channel.printf "load after: %d\n" load2