open Angstrom
open Base
open Core

module Field = struct
  (* a galaxy is a pair of row and column *)
  type galaxy = int * int [@@deriving sexp]

  (* a field is a list of galaxies *)
  type t = { rows : int; cols : int; galaxies : galaxy list } [@@deriving sexp]

  (* convert a field to a string *)
  let string_of_t t = Sexp.to_string_hum (sexp_of_t t)
end

module Parse = struct
  type tile = Galaxy | Empty

  exception Tile_error of char

  let parse_tile = function
    | '.' -> Empty
    | '#' -> Galaxy
    | c -> raise (Tile_error c)

  (* parse a tile into a galaxy *)
  let tile_pos row col c : Field.galaxy option =
    match parse_tile c with Empty -> None | Galaxy -> Some (row, col)

  (* parse a row into a list of galaxies *)
  let parse_row ix row =
    let lc = String.to_list row in
    List.filter_mapi ~f:(tile_pos ix) lc

  (* parse the input into a field *)
  let parse_field input =
    let lines = String.split_lines input in
    let rows, cols = (List.length lines, String.length (List.hd_exn lines))
    and galaxiess = List.mapi ~f:parse_row lines in
    let galaxies = List.concat galaxiess in
    Field.{ rows; cols; galaxies }
end

(* By representing the fields sparsely,
   we got the second part of the puzzle essentially for free. *)
module Part1 = struct
  type row = int
  type col = int

  (* check if given row is empty *)
  let row_is_empty (field : Field.t) (row : row) =
    assert (row < field.rows);
    let f (r, _) = Int.(r = row) in
    not (List.exists field.galaxies ~f)

  (* check if given column is empty *)
  let col_is_empty (field : Field.t) (col : col) =
    assert (col < field.cols);
    let f (_, c) = Int.(c = col) in
    not (List.exists field.galaxies ~f)

  (* generate a list of empty rows in the field *)
  let empty_rows (field : Field.t) =
    let rows = Sequence.range 0 field.rows and empty = row_is_empty field in
    Sequence.to_list (Sequence.filter ~f:empty rows)

  (* generate a list of empty columns in the field *)
  let empty_cols (field : Field.t) =
    let cols = Sequence.range 0 field.cols and empty = col_is_empty field in
    Sequence.to_list (Sequence.filter ~f:empty cols)

  (* calculate the offset for a given row/column based on the
     list of empty columns (before it) *)
  let count_offset empty_ixs ix =
    (* PART 1 / PART 2: change `replace` from 2 to 1000000 *)
    let replace = 2 in
    let multiplier = replace - 1 and f e = ix > e in
    multiplier * List.count empty_ixs ~f

  (* update the field by expanding empty rows *)
  let update_rows_with_empty (field : Field.t) (empty : row list) =
    let rows = field.rows + List.length empty
    and f (r, c) = (r + count_offset empty r, c) in
    let galaxies = List.map ~f field.galaxies in
    { field with galaxies; rows }

  (* update the field by expanding empty columns *)
  let update_cols_with_empty (field : Field.t) (empty : col list) =
    let cols = field.cols + List.length empty
    and f (r, c) = (r, c + count_offset empty c) in
    let galaxies = List.map ~f field.galaxies in
    { field with galaxies; cols }

  (* update the field by expanding empty rows and columns *)
  let update_field (field : Field.t) =
    let empty_rows = empty_rows field and empty_cols = empty_cols field in
    let field' = update_rows_with_empty field empty_rows in
    let field'' = update_cols_with_empty field' empty_cols in
    field''

  (* generate all pairs of galaxies *)
  let all_pairs (field : Field.t) : (Field.galaxy * Field.galaxy) Sequence.t =
    let g = Sequence.of_list field.galaxies in
    let cp = Sequence.cartesian_product g g
    and f ((a, b), (c, d)) = Int.(a < c || (a = c && b < d)) in
    Sequence.filter cp ~f

  (* convert a list of pairs to a string *)
  let string_of_pairs p =
    let f (g1, g2) =
      Sexp.List [ Field.sexp_of_galaxy g1; Field.sexp_of_galaxy g2 ]
    in
    Sexp.to_string_hum (List.sexp_of_t f p)

  (* calculate the manhattan distance between two galaxies *)
  let distance (r1, c1) (r2, c2) =
    let d a b = Int.(abs (a - b)) in
    d r1 r2 + d c1 c2

  (* calculate the cumulative manhattan distance of all pairs *)
  let cumulative_distance pairs =
    let f a (g1, g2) = a + distance g1 g2 in
    Sequence.fold pairs ~init:0 ~f
end

let () =
  let inp = In_channel.(input_all stdin) in
  let t = Parse.parse_field inp in
  let t' = Part1.update_field t in
  let pairs = Part1.all_pairs t' in
  (*let pairs_str = Part1.string_of_pairs (Sequence.to_list pairs) in*)
  let dist = Part1.cumulative_distance pairs in
  (*Out_channel.printf "original: %s\n" (Field.string_of_t t);
    Out_channel.printf "updated: %s\n" (Field.string_of_t t');
    Out_channel.printf "pairs with total dist %d: %s\n" dist pairs_str*)
  Out_channel.printf "%d pairs with total dist %d\n" (Sequence.length pairs)
    dist