(* Programmed by Jedidiah Barber *)
(* Licensed under the Sunset License v1.0 *)


module type Algebra =
  sig
    type t
    val zero : t
    val one : t
    val add : t -> t -> t
    val sub : t -> t -> t
    val mul : t -> t -> t
    val div : t -> t -> t
    val to_string : t -> string
  end



module type S =
  sig
    type element
    type wall
    val create : dimx:int -> dimy:int -> init:element array -> wall
    val x_length : wall -> int
    val y_length : wall -> int
    val has_value : wall -> x:int -> y:int -> bool
    val get : wall -> x:int -> y:int -> element
    val get_opt : wall -> x:int -> y:int -> element option
    val last_nonzero_row : wall -> int
    val to_string : wall -> string
    val print : wall -> unit
    val output : Stdlib.out_channel -> wall -> unit
    val pp_print : Stdlib.Format.formatter -> wall -> unit
  end



let (let*?) = Option.bind



module Make (A : Algebra) =
  struct

    type element = A.t
    type wall = element option array array

    let x_length wall =
      Array.length wall

    let y_length wall =
      Array.length wall.(0)

    let has_value wall ~x ~y =
      if x >= 0 && x < x_length wall &&
         y >= 0 && y < y_length wall
      then match wall.(x).(y) with
          None -> false
        | Some _ -> true
      else false

    let get_opt wall ~x ~y =
      if x >= 0 && x < x_length wall &&
         y >= 0 && y < y_length wall
      then wall.(x).(y)
      else None



    let next_nonzero_up wall ~x ~y =
      let rec walker y1 =
        let*? v = get_opt wall ~x ~y:y1 in
        if v <> A.zero then Some y1 else walker (y1-1)
      in walker y

    let next_nonzero_left wall ~x ~y =
      let rec walker x1 =
        let*? v = get_opt wall ~x:x1 ~y in
        if v <> A.zero then Some x1 else walker (x1-1)
      in walker x

    let next_nonzero_right wall ~x ~y =
      let rec walker x1 =
        let*? v = get_opt wall ~x:x1 ~y in
        if v <> A.zero then Some x1 else walker (x1+1)
      in walker x



    let next_valid_left wall ~x ~y =
      let rec walker x1 =
        if x1 < 0
        then None
        else if Option.is_some (get_opt wall ~x:x1 ~y)
          then Some x1
          else walker (x1-1) in
      walker x



    let safe_div a b =
      if b = A.zero then None else Some (A.div a b)

    let multi_mul items =
      let f a b = A.mul a b in
      List.fold_left f A.one items

    let square a =
      A.mul a a

    let rec power ~base ~num ~denom ~times =
      if times <= 0
      then Some base
      else let*? new_base = safe_div (A.mul base num) denom in
        power ~base:new_base ~num ~denom ~times:(times-1)



    (* red x red + green x green = blue squared *)
    (* target to calculate is red on the bottom of the cross *)
    (* current cell x two up + one diag up left x one diag up right = one up squared *)
    let cross_rule wall ~x ~y =
      let*? green_left = get_opt wall ~x:(x-1) ~y:(y-1) in
      let*? green_right = get_opt wall ~x:(x+1) ~y:(y-1) in
      let*? blue = get_opt wall ~x ~y:(y-1) in
      let*? red_top = get_opt wall ~x ~y:(y-2) in
      safe_div (A.sub (square blue) (A.mul green_left green_right)) red_top



    (* outer yellow x inner yellow squared + outer blue x inner blue squared *)
    (* equals outer grey x inner grey squared + outer red x inner red squared *)
    (* inner yellow is up, outer yellow is down *)
    (* inner blue is down, outer blue is up *)
    (* inner grey is left, outer grey is right *)
    (* inner red is right, outer red is left *)
    (* target to calculate is outer yellow *)
    let long_cross_rule wall ~x ~y =
      let*? outer_red = get_opt wall ~x:(x-2) ~y:(y-2) in
      let*? inner_red = get_opt wall ~x:(x+1) ~y:(y-2) in
      let*? red_product = Some (A.mul outer_red (square inner_red)) in
      let*? outer_grey = get_opt wall ~x:(x+2) ~y:(y-2) in
      let*? inner_grey = get_opt wall ~x:(x-1) ~y:(y-2) in
      let*? grey_product = Some (A.mul outer_grey (square inner_grey)) in
      let*? grey_red_sum = Some (A.add grey_product red_product) in
      let*? outer_blue = if y - 4 = -1 then Some A.zero else get_opt wall ~x ~y:(y-4) in
      let*? inner_blue = get_opt wall ~x ~y:(y-1) in
      let*? blue_product = Some (A.mul outer_blue (square inner_blue)) in
      let*? greyred_blue_diff = Some (A.sub grey_red_sum blue_product) in
      let*? inner_yellow = get_opt wall ~x ~y:(y-3) in
      safe_div greyred_blue_diff (square inner_yellow)



    (* All blocks of zeros in the wall come in squares. *)
    let zero_window wall ~x ~y =
      let*? top = next_nonzero_up wall ~x ~y:(y-1) in
      let*? left = next_nonzero_left wall ~x ~y:(top+1) in
      let*? right = next_nonzero_right wall ~x ~y:(top+1) in
      if y - 1 > top && right - left > 0 && right - left + 1 > y - 1 - top
      then Some A.zero
      else None



    (* denote ratio of top sequence from left to right as a/b *)
    (* denote ratio of bottom sequence from right to left as c/d *)
    (* denote ratio of left sequence from top to bottom as e/f *)
    (* denote ratio of right sequence from bottom to top as g/h *)
    (* denote rhs as s *)
    (* s is 1 when even sized window and -1 when odd sized window *)
    (* goal is to calculate c/d *)
    (* then *)
    (* ((a/b)(c/d))/((e/f)(g/h)) = s *)
    (* (a/b)(c/d) = s(e/f)(g/h) *)
    (* (a/b)(c/d) = (seg)/(fh) *)
    (* c/d = ((seg)/(fh))(b/a) *)
    (* c/d = (begs)/(afh) *)
    (* and since that ratio is going from right to left, invert to go left to right *)
    (* thus to transform bottom left corner number into next number... *)
    (* multiply by afh then divide by begs *)
    let horseshoe_rule wall ~x ~y =
      let*? top = next_nonzero_up wall ~x ~y:(y-1) in
      let*? left = next_nonzero_left wall ~x ~y:(top+1) in
      let*? right = next_nonzero_right wall ~x ~y:(top+1) in
      let*? _ = if y - top <> right - left || right - left <= 2 then None else Some 1 in
      let*? a = get_opt wall ~x:(left+1) ~y:top in
      let*? b = get_opt wall ~x:left     ~y:top in
      let*? e = get_opt wall ~x:left     ~y:(top+1) in
      let*? f = get_opt wall ~x:left     ~y:top in
      let*? g = get_opt wall ~x:right    ~y:top in
      let*? h = get_opt wall ~x:right    ~y:(top+1) in
      let*? s = if (right - left - 1) mod 2 = 0 then Some A.one else Some (A.sub A.zero A.one) in
      let*? prev = next_valid_left wall ~x ~y in
      power
        ~base:(Option.get wall.(prev).(y))
        ~num:(A.mul (A.mul a f) h)
        ~denom:(A.mul (A.mul (A.mul b e) g) s)
        ~times:(x - prev)



    (* denote ratio of top sequence from left to right as a/b *)
    (* denote ratio of bottom sequence from right to left as c/d *)
    (* denote ratio of left sequence from top to bottom as e/f *)
    (* denote ratio of right sequence from bottom to top as g/h *)
    (* denote sign of grey/red terms as s *)
    (* goal is to calculate outer yellow *)
    (* and also to put off division until the very end to ensure integer intermediate results *)
    (* then *)
    (* (e/f)(o_bl/i_bl) + s(a/b)(o_gr/i_gr) = (g/h)(o_ye/i_ye) + s(c/d)(o_re/i_re) *)
    (* (e/f)(o_bl/i_bl) + s(a/b)(o_gr/i_gr) - s(c/d)(o_re/i_re) = (g/h)(o_ye/i_ye) *)
    (* i_ye((e/f)(o_bl/i_bl) + s(a/b)(o_gr/i_gr) - s(c/d)(o_re/i_re)) = (g/h)o_ye *)
    (* o_ye = i_ye(h/g)((e/f)(o_bl/i_bl) + s(a/b)(o_gr/i_gr) - s(c/d)(o_re/i_re)) *)
    (* o_ye = i_ye(h/g)((e*o_bl)/(f*i_bl) + s(a*o_gr)/(b*i_gr) - s(c*o_re)/(d*i_re)) *)
    (* o_ye = h*i_ye*(e*o_bl*b*i_gr*d*i_re + s*a*o_gr*f*i_bl*d*i_re - s*c*o_re*f*i_bl*b*i_gr) / *)
    (*        (g*f*i_bl*b*i_gr*d*i_re) *)
    let broken_cross_rule wall ~x ~y =
      let*? top = next_nonzero_up wall ~x ~y:(y-2) in
      let*? left = next_nonzero_left wall ~x ~y:(top+1) in
      let*? right = next_nonzero_right wall ~x ~y:(top+1) in
      let*? bottom = Some (y - 1) in
      let*? _ = if bottom - top <> right - left || right - left <= 2 then None else Some 1 in
      let*? a = get_opt wall ~x:(left+1) ~y:top in
      let*? b = get_opt wall ~x:left     ~y:top in
      let*? c = get_opt wall ~x:left     ~y:bottom in
      let*? d = get_opt wall ~x:(left+1) ~y:bottom in
      let*? e = get_opt wall ~x:left     ~y:(top+1) in
      let*? f = get_opt wall ~x:left     ~y:top in
      let*? g = get_opt wall ~x:right    ~y:top in
      let*? h = get_opt wall ~x:right    ~y:(top+1) in
      let*? s = if (right - x) mod 2 = 0 then Some A.one else Some (A.sub A.zero A.one) in
      let*? inner_yellow = get_opt wall ~x                ~y:bottom in
      let*? inner_blue   = get_opt wall ~x:(right-x+left) ~y:top in
      let*? inner_red    = get_opt wall ~x:right          ~y:(bottom-right-x) in
      let*? inner_grey   = get_opt wall ~x:left           ~y:(right-x+top) in
      let*? outer_blue   = get_opt wall ~x:(right-x+left) ~y:(top-1) in
      let*? outer_red    = get_opt wall ~x:(right+1)      ~y:(bottom-right-x) in
      let*? outer_grey   = get_opt wall ~x:(left-1)       ~y:(right-x+top) in
      let*? term_1 = Some (multi_mul [   e; outer_blue; b; inner_grey; d; inner_red]) in
      let*? term_2 = Some (multi_mul [s; a; outer_grey; f; inner_blue; d; inner_red]) in
      let*? term_3 = Some (multi_mul [s; c; outer_red;  f; inner_blue; b; inner_grey]) in
      let*? term_sum = Some (A.sub (A.add term_1 term_2) term_3) in
      let*? numerator = Some (multi_mul [h; inner_yellow; term_sum]) in
      let*? denominator = Some (multi_mul [g; f; inner_blue; b; inner_grey; d; inner_red]) in
      safe_div numerator denominator



    let rec attempt_order wall ~x ~y funcs =
      if funcs = []
      then None
      else let r = (List.hd funcs) wall ~x ~y in
        if Option.is_some r
        then r
        else attempt_order wall ~x ~y (List.tl funcs)

    let create ~dimx ~dimy ~init =
      let result = Array.make_matrix dimx dimy None in
      for i = 0 to dimx - 1 do
        result.(i).(0) <- Some A.one
      done;
      for i = 0 to min (dimx - 1) ((Array.length init) - 1) do
        result.(i).(1) <- Some init.(i)
      done;
      try
        for j = 2 to dimy - 1 do
          let do_exit = ref true in
          for i = 0 to dimx - 1 do
            let v = attempt_order result ~x:i ~y:j
              [cross_rule; long_cross_rule; zero_window; horseshoe_rule; broken_cross_rule] in
            result.(i).(j) <- v;
            if Option.is_some v && Option.get v <> A.zero
            then do_exit := false
          done;
          if !do_exit then raise Exit
        done;
        result
      with Exit ->
        result



    let get wall ~x ~y =
      if x >= 0 && x < x_length wall &&
         y >= 0 && y < y_length wall
      then match wall.(x).(y) with
          None -> raise Not_found
        | Some v -> v
      else raise Not_found

    let last_nonzero_row wall =
      let rec checker x y =
        if wall.(x).(y) <> Some A.zero && wall.(x).(y) <> None
        then (if y + 1 = y_length wall then y else checker 0 (y + 1))
        else (if x + 1 = x_length wall then y - 1 else checker (x + 1) y)
      in checker 0 0



    let row_string wall ~y =
      let result = ref "" in
      for i = 0 to (x_length wall) - 2 do
        let v = wall.(i).(y) in
        if Option.is_some v
        then result := !result ^ (A.to_string (Option.get v)) ^ ","
        else result := !result ^ ","
      done;
      let v = wall.(x_length wall - 1).(y) in
      if Option.is_some v
      then result := !result ^ (A.to_string (Option.get v));
      !result

    let to_string wall =
      let result = ref "" in
      for j = 0 to (y_length wall) - 1 do
        result := !result ^ (row_string wall ~y:j) ^ (String.make 1 (char_of_int 10))
      done;
      !result

    let print wall =
      print_string (to_string wall)

    let output channel wall =
      output_string channel (to_string wall)

    let pp_print format wall =
      for j = 0 to (y_length wall) - 1 do
        Format.pp_print_string format (row_string wall ~y:j);
        Format.pp_print_break format 8 0
      done

  end