#!/usr/bin/env owl

open Owl

let fname = "test.ppm"

let read_img fname =
  (* naive grayscale binary ppm reading. If you know how to, improve it *)
  let fp = open_in fname in
  let ver = input_line fp in (* version *)

  if ver <> "P6" then (* expect a .ppm file *)
    raise (Invalid_argument ("Unable to read image: " ^ fname))
  else
    ();

  (* This will skip comments. *)
  let rec ignore_comments_then_get_w_h () =
    let maybe_comment = input_line fp in
    if maybe_comment.[0] = '#' then
      ignore_comments_then_get_w_h ()
    else
      (* This line has the width and height in it *)
      maybe_comment
  in

  (* width, height num colors *)
  let w_h_line = ignore_comments_then_get_w_h () in
  let num_col_line = input_line fp in
  let w, h = Scanf.sscanf w_h_line "%d %d" (fun w h -> w, h) in
  let num_col = Scanf.sscanf num_col_line "%d" (fun n -> n) in

  let img   = String.make (w * h * 3) ' ' in
  (*
  let imf_r = Array.make_matrix w h 0.0 in
  let imf_g = Array.make_matrix w h 0.0 in
  let imf_b = Array.make_matrix w h 0.0 in
  *)

  let imf_o = Array.make_matrix (w * 3) h 0.0 in

  (* Note that under 32bit OCaml, this will only work when reading strings up
     to ~16 megabytes. *)
  really_input fp img 0 (w * h * 3); (* this is indeed the end of file *)

  close_in fp;

  (*
  for i = 0 to w - 1 do
    for j = 0 to h - 1 do
      (*
      imf_r.(i).(j) <-
        (* flip image up-down *)
        float_of_int (int_of_char (img.[(h - 1 - j ) * w + (3*i + 0)]));
      imf_g.(i).(j) <-
        (* flip image up-down *)
        float_of_int (int_of_char (img.[(h - 1 - j ) * w + (3*i + 1)]));
      imf_b.(i).(j) <-
        (* flip image up-down *)
        float_of_int (int_of_char (img.[(h - 1 - j ) * w + (3*i + 2)]));
      *)
    done
  done;
  *)

  let ww = 3 * w  in
  for i = 0 to ww - 1 do
    for j = 0 to h - 1 do
      imf_o.(i).(j) <- float_of_int (int_of_char (img.[(h - 1 - j ) * ww + i ]));
    done
  done;

  (* imf_o, imf_r, imf_g, imf_b, w, h, num_col *)
  imf_o, w, h, num_col

let img_to_owl fname =
  let img, w, h, num_col = read_img fname in
  let m =  Mat.of_arrays img in
  let m = Mat.rotate m 270 in
  (* r,g, b: Mat of size h * w *)
  let r = Mat.get_slice_simple [[];[0;-1;3]] m in
  let g = Mat.get_slice_simple [[];[1;-1;3]] m in
  let b = Mat.get_slice_simple [[];[2;-1;3]] m in

  (* combine 3 channel matrices into one ndarray *)
  let r' = Mat.to_ndarray r in
  let g' = Mat.to_ndarray g in
  let b' = Mat.to_ndarray b in

  let r' = Arr.reshape r' [|h;w;1|] in
  let g' = Arr.reshape g' [|h;w;1|] in
  let b' = Arr.reshape b' [|h;w;1|] in

  let img = Arr.zeros [|h;w;3|] in
  Arr.set_slice  [R []; R []; I 0] img r';
  Arr.set_slice  [R []; R []; I 1] img g';
  Arr.set_slice  [R []; R []; I 2] img b';
  r, g, b, img

(*
let _ =
  let r, _, _, _ = img_to_owl fname in
  Plot.image r
*)