(* This piece of code thresholds an image by some rule *)
(* Todo:
    - Code for coloring corners. Easy as hell.
    - Cut out and push to a matrix.
    - Consider changing color channels. RGB -> YIQ. May be there is
      something to exploit here.
    - Consider DCT'ing each block. Would be good for the matching algo.


   Larger items:
    - Edge matching is probably weak. Use sinusoids from the DCT. Will make
      much better probability calculations.
    - Consider how to rebuild in 2d. An 1d reassembly is there already
      based on a vertex-cover problem selection. So see if this can be
      adapted to a 2d world. If not, there are some papers on other methods.
*)
open Images;;
open OImages;;

let isWhite p = (p.r = 255 && p.g = 255 && p.b = 255);;

let blob_extract (bmp : rgb24) =
  let width = bmp#width in
  let height = bmp#height in
  let rgb24 = new rgb24 width height in
  let current_num = ref 0 in
  let f x y =
    let testCoord x y =
      try
	let point = rgb24#get x y in
	  if point.r != 0
	  then Some (point.r + point.g + point.b)
	  else None
      with
	  Out_of_image -> None in
    let color v =
      let g = v - 255 in
      let b = v - (2 * 255) in
	{ r = if v > 255 then 255 else v;
	  g = if g < 0 then 0 else if g > 255 then 255 else g;
	  b = if b < 0 then 0 else if b > 255 then 255 else b } in
    let point = bmp#get x y in
      if (isWhite point) then
	(* Note this is a maybe monad *)
	match List.fold_left
	  (fun res (x, y) ->
	    match res with
	      | Some v -> Some v
	      | None -> match testCoord x y with
		  | Some v -> Some v
		  | None -> None)
	  None
	  [(x + 1, y - 1); (x, y - 1); (x - 1, y - 1); (x - 1, y)] with
	    | None ->
		current_num := !current_num +1;
		color !current_num
	    | Some v ->
		color v
      else
	{r = 0; g = 0; b = 0} in
  let clear_image rgb =
    for x = 0 to width - 1 do
      for y = 0 to height - 1 do
	rgb24#unsafe_set x y {r = 0; g = 0; b = 0}
      done
    done in
    clear_image rgb24;
    for x = 0 to width - 1 do
      for y = 0 to height - 1 do
	rgb24#unsafe_set x y (f x y)
      done
    done;
    print_string ((string_of_int !current_num) ^ "\n");
    rgb24;;

let threshold (bmp : rgb24) =
  let width = bmp#width in
  let height = bmp#height in
  let rgb24 = new rgb24 width height in

  let f x y =
    let point = bmp#get x y in
      if (point.r = 255
          && point.g = 255
          && point.b = 255)
      then
	{ r = 0; g = 0; b = 0}
      else
	{ r = 255; g = 255; b = 255} in

    for x = 0 to width -1 do
      for y = 0 to height -1 do
	rgb24#unsafe_set x y (f x y)
      done
    done;

    rgb24;;

let files = ref [] in
  Arg.parse [] (fun s -> files := s :: !files) "threshold files";
  let files = List.rev !files in
    List.iter (fun file ->
      try
	let rgb =
	  let oimage = OImages.load file [] in
	    match OImages.tag oimage with
	      | Index8 img ->
		  let rgb = img#to_rgb24 in
		    img#destroy;
		    rgb
	      | Index16 img ->
		  let rgb = img#to_rgb24 in
		    img#destroy;
		    rgb
	      | Rgb24 img -> img
	      | _ -> raise (Invalid_argument "not supported") in
	let rgb' = blob_extract (threshold rgb) in

	let get_extension s =
	  let dotpos = String.rindex s '.' in
	    String.sub s (dotpos + 1) (String.length s - dotpos - 1) in

	let ext = get_extension file in
	let body = String.sub file 0
	  (String.length file - String.length ext - 1) in
	let outfile = body ^ ".thresh.bmp" in
	  rgb'#save outfile (Some Bmp) [];
      with _ -> ()) files;;