open Angstrom
open Base
open Core

type tile = Ash | Rocks [@@deriving sexp, show, eq]

(* A field consists of a tile array of size rows * cols *)
type field = { rows : int; cols : int; tiles : tile array }
[@@deriving sexp, show]

(* convert indices and access field. *)

let ix_of_rc field r c = (r * field.cols) + c
let rc_of_ix field ix = (ix / field.cols, ix mod field.cols)
let tile_at field r c = field.tiles.(ix_of_rc field r c)
let get_column field c = Array.init field.rows ~f:(fun r -> tile_at field r c)
let get_row field r = Array.init field.cols ~f:(fun c -> tile_at field r c)

module Parse = struct
  let maybe p = p >>| (fun _ -> ()) <|> return ()
  let parse_tile = char '.' >>| (fun _ -> Rocks) <|> (char '#' >>| fun _ -> Ash)

  let parse_field =
    let parse_row = many1 parse_tile in
    let parse_rows = sep_by1 (char '\n') parse_row in
    let parse_field rows =
      let cols = List.hd_exn rows |> List.length in
      let tiles = List.concat rows |> Array.of_list in
      { rows = List.length rows; cols; tiles }
    in
    parse_rows >>| parse_field

  let parse_fields = sep_by1 (string "\n\n") parse_field <* maybe (char '\n')
end

module Part1 = struct
  let check_mirror getter max field ix =
    let rec check_with_off field ix off =
      if ix - off - 1 < 0 || ix + off >= max then off > 0 && off < max
      else
        let left = getter field (ix - off - 1)
        and right = getter field (ix + off) in
        Array.equal equal_tile left right && check_with_off field ix (off + 1)
    in
    check_with_off field ix 0

  (* check if mirror axis exists between columns c and c+1 (one-based) *)
  let check_mirror_v field c = check_mirror get_column field.cols field c

  (* check if mirror axis exists between rows r and r+1 (one-based) *)
  let check_mirror_h field r = check_mirror get_row field.rows field r

  let check_any checker max field =
    let indices = List.range 1 (max + 1) in
    List.find indices ~f:(fun ix -> checker field ix)

  (* check if mirror axis exists between any rows r and r+1 (one-based),
     and return r. *)
  let check_any_h field = check_any check_mirror_h field.rows field

  (* check if mirror axis exists between any columns c and c+1 (one-based),
     and return c. *)
  let check_any_v field = check_any check_mirror_v field.cols field

  (* a symmetry may be either horizontal or vertical or both.
     The integer signifies the row/column to the left/top of the
     symmetry line. *)
  type symmetry = { horizontal : int option; vertical : int option }
  [@@deriving sexp, show]

  let check field =
    let horizontal = check_any_h field in
    let vertical = check_any_v field in
    { horizontal; vertical }

  (* check symmetry for all fields *)
  let check_all fields = List.map fields ~f:check

  (* calculate score for a single symmetry according to the puzzle rules. *)
  let symmetry_to_score { horizontal; vertical } =
    let h = Option.value ~default:0 horizontal in
    let v = Option.value ~default:0 vertical in
    v + (100 * h)

  (* calculate score for each symmetry and sum them up *)
  let symmetries_to_score symmetries =
    List.map symmetries ~f:symmetry_to_score |> List.fold ~init:0 ~f:( + )
end

module Part2 = struct
  (* throws on index error, ensure that there is only one differing element *)
  let rec find_different_pos ?(ix = 0) a b =
    if equal_tile a.(ix) b.(ix) then ix else find_different_pos a b ~ix:(ix + 1)

  (* tracking is a bit more complex than it seems: when comparing two rows/columns,
       the result is either No smudge: we can continue looking; One smudge: we found the smudge
       but there can't be a second one; or more than one difference, in which case we're not
     looking at a valid symmetry axis. *)
  type continue = No_smudge | One_smudge of int | Symmetry_fail

  (* find the difference between two rows/columns, and detect a smudge if present. *)
  let find_diff a b =
    match
      Array.fold2_exn a b ~init:0 ~f:(fun acc a b ->
          if equal_tile a b then acc else acc + 1)
    with
    | 0 -> No_smudge
    | 1 ->
        let ix = find_different_pos a b in
        One_smudge ix
    | _ -> Symmetry_fail

  let find_smudge getter max field ix =
    let rec find_with_off field ix off =
      if ix - off - 1 < 0 || ix + off >= max then No_smudge
      else
        let left = getter field (ix - off - 1)
        and right = getter field (ix + off)
        and rest = find_with_off field ix (off + 1) in
        match find_diff left right with
        | One_smudge ix -> (
            match rest with No_smudge -> One_smudge ix | _ -> Symmetry_fail)
        | No_smudge -> rest
        | Symmetry_fail -> Symmetry_fail
    in
    find_with_off field ix 0

  (* find a smudge in column c (one-based). *)
  let find_smudge_v field c = find_smudge get_column field.cols field c

  (* find a smudge in row r (one-based). *)
  let find_smudge_h field r = find_smudge get_row field.rows field r

  let find_any_smudge max smudgefinder field =
    let open Option in
    let indices = List.range 1 (max + 1) in
    List.find_map indices ~f:(fun ix ->
        match smudgefinder field ix with One_smudge _ -> Some ix | _ -> None)

  (* find a new vertical symmetry axis in a field, ignoring exactly one smudge. *)
  let find_any_smudge_v field = find_any_smudge field.cols find_smudge_v field

  (* find a new horizontal symmetry axis in a field, ignoring exactly one smudge. *)
  let find_any_smudge_h field = find_any_smudge field.rows find_smudge_h field

  (* find a new symmetry axis in a field, ignoring exactly one smudge. *)
  let find_new_symmetry field =
    let open Part1 in
    match (find_any_smudge_h field, find_any_smudge_v field) with
    | Some h, Some v -> { horizontal = Some h; vertical = Some v }
    | Some h, None -> { horizontal = Some h; vertical = None }
    | None, Some v -> { horizontal = None; vertical = Some v }
    | None, None -> assert false

  (* find all new symmetries in a list of fields. *)
  let find_all_symmetries fields = List.map fields ~f:find_new_symmetry
end

let () =
  let inp = In_channel.(input_all stdin) in
  let fields =
    Angstrom.parse_string ~consume:All Parse.parse_fields inp
    |> Result.ok_or_failwith
  in
  let symmetries = Part1.check_all fields in
  let score = Part1.symmetries_to_score symmetries in
  let new_symmetries = Part2.find_all_symmetries fields in
  let new_score = Part1.symmetries_to_score new_symmetries in
  Out_channel.(
    printf "score = %d\n" score;
    printf "new score = %d\n" new_score)