open Base
open Core

module Roof = struct
  type t = { a : bool array; m : int; n : int }

  let random_bool ?(chance = 30) () = Random.int 100 < chance

  let create m n =
    { a = Array.init (n * m) ~f:(fun _i -> random_bool ()); m; n }

  let string_of_t arr =
    let char_of_bool = function true -> " X " | false -> " O " in
    let f ix a e =
      if Int.((ix + 1) % arr.n = 1) then
        Printf.sprintf "\n%s" (char_of_bool e) :: a
      else char_of_bool e :: a
    in
    let parts = List.rev @@ ("\n" :: Array.foldi ~init:[] ~f arr.a) in
    let parts = Printf.sprintf "%d x %d roof\n" arr.m arr.n :: parts in
    String.concat parts

  type cover = { rows : bool array; cols : bool array } [@@deriving sexp]

  (* Memory for covered rows/columns. Runtime and memory efficient: no copies
     necessary, just some bits set temporarily. *)
  let cover_for_roof roof =
    {
      rows = Array.create ~len:roof.m false;
      cols = Array.create ~len:roof.n false;
    }

  (* Execute function f while row m is covered in cover.
     Restore previous state after f returns. Return f's result. *)
  let with_row_covered m cover f =
    let old = cover.rows.(m) in
    cover.rows.(m) <- true;
    let r = f cover in
    cover.rows.(m) <- old;
    r

  let with_col_covered m cover f =
    let old = cover.cols.(m) in
    cover.cols.(m) <- true;
    let r = f cover in
    cover.cols.(m) <- old;
    r

  let is_covered cover x y = cover.rows.(x) || cover.cols.(y)

  let covers cover =
    let f a e = if e then a + 1 else a in
    let rr = Array.fold ~init:0 ~f cover.rows in
    let rc = Array.fold ~init:0 ~f cover.cols in
    rr + rc

  let coords_from_ix roof ix =
    let x = ix / roof.m in
    let y = ix % roof.m in
    (x, y)

  let first_open_leak roof cover =
    let f ix e =
      let x, y = coords_from_ix roof ix in
      match (e, is_covered cover x y) with
      | true, false -> Some (x, y)
      | _ -> None
    in
    Array.find_mapi roof.a ~f
end

(* Recursive optimal solution guaranteed *)
let rec solve roof cover =
  match Roof.first_open_leak roof cover with
  | None -> Roof.covers cover
  | Some (lx, ly) ->
      let row_covered = Roof.with_row_covered lx cover (solve roof) in
      let col_covered = Roof.with_col_covered ly cover (solve roof) in
      Int.min row_covered col_covered

let solve0 roof =
  let cover = Roof.cover_for_roof roof in
  solve roof cover

let () =
  let roof = Roof.create 5 5 in
  let min = solve0 roof in
  let open Out_channel in
  printf "min: %d\n" min;
  output_string stdout "\n";
  output_string stdout (Roof.string_of_t roof)