(***********************************************************************)
(*                                                                     *)
(*                          OCamlP3l                                   *)
(*                                                                     *)
(* (C) 2004-2007                                                       *)
(*             Roberto Di Cosmo (dicosmo@dicosmo.org)                  *)
(*             Zheng Li (zli@lip6.fr)                                  *)
(*             Pierre Weis (Pierre.Weis@inria.fr)                      *)
(*             Francois Clement (Francois.Clement@inria.fr)            *)
(*                                                                     *)
(* Based on original Ocaml P3L System                                  *)
(* (C) 1997 by                                                         *)
(*             Roberto Di Cosmo (dicosmo@ens.fr)                       *)
(*             Marco Danelutto (marcod@di.unipi.it)                    *)
(*             Xavier Leroy  (Xavier.Leroy@inria.fr)                   *)
(*             Susanna Pelagatti (susanna@di.unipi.it)                 *)
(*                                                                     *)
(* This program is free software; you can redistribute it and/or       *)
(* modify it under the terms of the GNU Library General Public License *)
(* as published by the Free Software Foundation; either version 2      *)
(* of the License, or (at your option) any later version.              *)
(*                                                                     *)
(* This program is distributed in the hope that it will be useful,     *)
(* but WITHOUT ANY WARRANTY; without even the implied warranty of      *)
(* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the       *)
(* GNU Library General Public License for more details.                *)
(*                                                                     *)
(***********************************************************************)

1(***********************************************************************)
(*                                                                     *)
(*                   Objective Caml P3L system                         *)
(*                                                                     *)
(* (C) 1997 by                                                         *)
(*             Roberto Di Cosmo (dicosmo@ens.fr)                       *)
(*             Marco Danelutto (marcod@di.unipi.it)                    *)
(*             Xavier Leroy  (Xavier.Leroy@inria.fr)                   *)
(*             Susanna Pelagatti (susanna@di.unipi.it)                 *)
(*                                                                     *)
(* This program is free software; you can redistribute it and/or       *)
(* modify it under the terms of the GNU Library General Public License *)
(* as published by the Free Software Foundation; either version 2      *)
(* of the License, or (at your option) any later version.              *)
(*                                                                     *)
(* This program is distributed in the hope that it will be useful,     *)
(* but WITHOUT ANY WARRANTY; without even the implied warranty of      *)
(* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the       *)
(* GNU Library General Public License for more details.                *)
(*                                                                     *)
(***********************************************************************)

(* $Id: grafp3l.ml,v 1.19 2007-01-23 15:50:04 weis Exp $ *)

open Graphics;;

type color = int;;

type ('a, 'b) tree =
   | Seq of string * color
   | Pipe of ('a, 'b) tree list
   | Farm of ('a, 'b) tree * int * (color * color list)
   | Map of ('a, 'b) tree * int * (color * color list)
   | Reduce of ('a, 'b) tree * int * (color * color list)
   | Loop of ('a,'b) tree * color;;

let seq ?(col=0) (f : unit -> 'a -> 'b) =
 (Obj.magic (Seq ("", col)) : ('a P3lstream.t, 'b P3lstream.t) tree);;

let pipeline (f : ('a, 'b) tree) (g : ('b, 'c) tree) =
  Pipe [ (Obj.magic f : ('a, 'c) tree); (Obj.magic g : ('a, 'c) tree) ];;

let ( ||| ) = pipeline;;

let farm ?(col=0) ?(colv=[]) (f, n) = Farm (f, n, (col, colv));;
let loop ?(col=0) (t, (f : ('a, 'a) tree)) = Loop (f, col);;
let mapvector ?(col=0) ?(colv=[])
    ((f : ('a P3lstream.t, 'b P3lstream.t) tree), n) =
  Map((Obj.magic f : ('a array P3lstream.t, 'b array P3lstream.t) tree),
      n, (col, colv));;
let reducevector ?(col=0) ?(colv=[]) 
    ((f: (('a * 'a) P3lstream.t, 'a P3lstream.t) tree), n) = 
  Reduce((Obj.magic f : ('a array P3lstream.t, 'a P3lstream.t) tree), n,
         (col, colv));;

(* estimate of actual size, expressed in number of "unit balls",
   counting seq as 1 *)

let defaultmaxdisplay = ref 5;;

let rec evalsize ?(maxdisplay=defaultmaxdisplay) = function
  | Seq _ -> (1, 1)
  | Pipe tl ->
    let sl = List.map evalsize tl in
    List.fold_left
      (fun (sumx, maxy) (x, y) -> (x + sumx, max y maxy))
      (0, 0) sl
  | Farm (f, n, _)
  | Map (f, n, _) ->
    let x, y = evalsize f in
    (x + 2, (if n > !maxdisplay then 3 else n) * y)
  | Reduce (f,n, _) ->
      let x, y = evalsize f in
      (x + 2, (if n > !maxdisplay then 3 else n) * y)
  | Loop (f, _) ->
    let x, y = evalsize f in
    (x + 2, y + 1)  (* y + 1 for the feedback loop *)

let get_color col = function
  | 0 -> col
  | color -> color;;

let rec nvlist n c cl =
  match n, cl with
      (0, _) -> []
    | (_, []) -> c :: nvlist (n - 1) c []
    | (_, hd :: tl) -> get_color c hd :: nvlist (n - 1) c tl;;

let ht = Hashtbl.create 16;;

let add_tbl tbl col num =
  let tmp = try Hashtbl.find tbl col with Not_found -> 0 in
  Hashtbl.replace tbl col (tmp + num);;

let rec evalcol col =
  let color = get_color col in
  function
  | Seq (_, c) -> add_tbl ht (color c) 1
  | Pipe tl -> List.iter (evalcol col) tl
  | Farm (f, n, (c, cl))
  | Map (f, n, (c, cl)) ->
      add_tbl ht (color c) 2;
      List.iter (fun x -> evalcol x f) (nvlist n (color c) cl)
  | Reduce (f,n, (c,cl)) ->
      add_tbl ht (color c) 2;
      List.iter (fun x -> evalcol x f) (nvlist n (color c) cl)
  | Loop (f, c) ->
    add_tbl ht (color c) 1;
    evalcol col f;;

let rec draw_elem n col x0 y0 x1 y1 r =
  let name = if col = 0 then n else (n ^ " C" ^ string_of_int col) in
  let xc = (x1 - x0) / 2 in
  let yc = (y1 - y0) / 2 in
  let lname = fst (text_size name) in
  begin
    draw_circle (xc + x0) (y0 + yc) r ;
    moveto (xc + x0 - (lname / 2)) (y0 + yc);
    draw_string name
  end

and draw_seq name col x0 y0 x1 y1 r =
  draw_elem name col x0 y0 x1 y1 r;
  let xc = (x1 - x0) / 2 in
  let yc = (y0 + (y1 - y0) / 2) in
  begin
    moveto x0 yc;
    lineto (x0 + xc - r) yc;
    moveto (x1 - xc + r) yc;
    lineto x1 yc;
  end

and draw_pipe e col x0 y0 x1 y1 r =
  let (dxl, _) = List.split(List.map evalsize e) in
  let dxll =
    let sum = ref 0 in
      List.fold_left (fun l x -> (sum:=!sum + x;l@[(!sum)])) [] dxl in
  let dx = (x1 - x0) / (List.hd (List.rev dxll)) in
  let combl = List.combine dxl dxll in
  List.iter2
    (fun (a, b) x ->
       draw_expr x col (x0 + dx * (b - a)) y0 (x0 + dx * b) y1 r)
    combl e

and draw_farm_like ?(nc="FarmC") ?(ne="FarmE")
      ?(maxdisplay=defaultmaxdisplay) e c cl n x0 y0 x1 y1 r =
  let truen = n in
  let (draw_n, n) = if n > !maxdisplay then (true, 3) else (false, n) in
  let (nx, ny) = evalsize e in
  let dy = (y1 - y0) / n in
  let dx = (x1 - x0) / (nx + 2) in
  let x00 = x0 + dx in
  let x0c = x0 + dx / 2 + r in
  let x11 = x1 - dx in
  let x1c = x1 - dx / 2 - r in
  let yc = y0 + (y1 - y0) / 2 in
(* Printf.printf
     "Farm: n =%d;nx=%d;ny=%d;dx=%d;dy=%d;x0=%d;x1=%d\n"
     n nx ny dx dy x0 x1; *)
  begin
    for i = 0 to n - 1 do
      begin
        (* draw the subnet, or the omissis if too big *)
        (match draw_n, i with
         | true, 1 -> draw_seq (Printf.sprintf "* %d " truen) 0
         | true, 2 -> draw_expr e (List.nth cl (truen - 1))
         | _ ->
           draw_expr e (List.nth cl i))
           x00 (y0 + (n - i - 1) * dy) x11 (y0 + (n - i) * dy) r;
        (* connect emitter and collector to this subnet *)
        moveto x0c yc;lineto x00 ((y0 + i * dy) + dy / 2);
        moveto x1c yc;lineto x11 ((y0 + i * dy) + dy / 2)
      end
    done;
    draw_elem nc c x11 y0 x1  y1 r;
    moveto x0 yc; lineto (x0c - 2 * r) yc;
    draw_elem ne c x0  y0 x00 y1 r;
    moveto x1 yc; lineto (x1c + 2 * r) yc;
  end

and draw_loop e col x0 y0 x1 y1 r =
  let (nx, ny) = evalsize e in
  let dx = (x1 - x0) / (nx + 2) in
  begin
    (* the 2 loop nodes *)
    draw_seq "Li" col x0 y0 (x0 + dx) y1 r;
    draw_seq "Lo" col (x1 - dx) y0 x1 y1 r;
    (* the lines to the inner expression *)
    moveto (x0 + dx / 2 + r) (y0 + (y1 - y0) / 2);
    lineto (x0 + dx) (y0 + (y1 - y0) / 2);
    moveto (x1 - dx / 2 - r) (y0 + (y1 - y0) / 2);
    lineto (x1 - dx) (y0 + (y1 - y0) / 2);
    (* the inner expression *)
    draw_expr e col (x0 + dx) y0 (x1 - dx) y1 r;
    (* the feedback loop *)
    moveto (x1 - (dx / 2)) (y0 + (y1 - y0) / 2 - r);
    lineto (x1 - (dx / 2)) (y0 + (y1 - y0) / 2 - (y1 - y0) / 4 - (y1 - y0) / 5);
    lineto (x0 + (dx / 2)) (y0 + (y1 - y0) / 2 - (y1 - y0) / 4 - (y1 - y0) / 5);
    lineto (x0 + (dx / 2)) (y0 + (y1 - y0) / 2 - r)
  end

and draw_expr e col x0 y0 x1 y1 r =
  (* moveto x0 y0; lineto x0 y1; lineto x1 y1; lineto x1 y0; lineto x0 y0; *)
  let color = get_color col in
  match e with
  | Seq (s, c) -> draw_seq s (color c) x0 y0 x1 y1 r
  | Pipe pl -> draw_pipe pl col x0 y0 x1 y1 r
  | Farm (e, n, (c, cl)) ->
    draw_farm_like e (color c) (nvlist n (color c) cl) n x0 y0 x1 y1 r
  | Map (e, n, (c, cl)) ->
    draw_farm_like ~nc:"MapC" ~ne:"MapE" e (color c)
      (nvlist n (color c) cl) n x0 y0 x1 y1 r
  | Reduce (e, n, (c, cl)) ->
    draw_farm_like ~nc:"RedC" ~ne:"RedE" e (color c)
      (nvlist n (color c) cl) n x0 y0 x1 y1 r
  | Loop (e, c) -> draw_loop e (color c) x0 y0 x1 y1 r;;

let draw_main sx sy =
  (* draw box *)
  moveto 2 2; lineto 2 (sy - 4); lineto sx (sy - 4); lineto sx 2; lineto 2 2;
  (* draw name *)
  moveto 6 (sy - 20); draw_string "pardo";;

(* Handling parfuns *)

let parfuns = ref [];;

let (parfun, getparfuns) =
  (fun (e : unit -> ('a, 'b) tree) ->
     parfuns := (Obj.magic e : (unit -> (unit, unit) tree)) :: !parfuns;
     (fun (x : 'a) -> (Obj.magic x : 'b))),
  (fun () -> !parfuns);;

let setsize, getwidth, getheight =
  let picwidth = ref 1024 in
  let picheight = ref 768 in
  (fun s ->
    let pos = String.index s 'x' in
    let len = String.length s in
    picwidth := int_of_string (String.sub s 0 pos);
    picheight := int_of_string (String.sub s (pos + 1) (len - pos - 1))),
  (fun () -> !picwidth),
  (fun () -> !picheight);;

let spec_list = [
  "-size", Arg.String setsize, "Specify the pic sizes with int*int";
  "-maxpara", Arg.Int (fun x -> defaultmaxdisplay:= x),
    "Define the max para-structures shown for \
     farm / mapvector / reducevector with int";
];;

List.iter (fun (nm, act, man) -> Command_options.add nm act man) spec_list;;

let pardo expr =
  Arg.parse
    (Command_options.all ())
    (fun s -> prerr_endline ("Unknown anonymous argument " ^ s); exit 2)
    ("Usage: " ^ Sys.argv.(0) ^ " [options]");
  (* first, compute global x and y size in basic units,
     and add to the skeletons info on their x, y size and y placement *)
  let (skels, (xsize, ysize)) =
    List.fold_left
      (fun (acc, (x, y)) ss ->
       let s = ss () in
       let x', y' = evalcol 0 s; evalsize s in
       (s, (x', y'), y + y') :: acc, (max x x', y + y'))
      ([], (0, 0))
      (getparfuns ()) in

  open_graph (Printf.sprintf " %ix%i" (getwidth ()) (getheight ()));
  let main_width = 80 in
  let sx = size_x () - main_width and sy = size_y() in
  (* avoid division by 0 *)
  let dy = size_y() / (max ysize 1) and dx = sx / (max xsize 1) in
  let dxy = min dx dy in (* this is the actual pixel size used to draw a seq *)
  (* draw the main "pardo" sequential box *)
  draw_main main_width sy;
  (* draw each skeleton in a box which is x * dxy wide, y * dxy tall
     and placed at (y' - y) * dxy position *)
  (* y' is the absolute position in basic units
     of the top corners of the box *)
  List.iter
    (fun (s, (x, y), y') ->
       draw_expr s 0 main_width ((y' - y) * dxy) (main_width + x * dxy)
         (y' * dxy) (dxy / 5))
    skels; (* dxy / 5 is used as the ball radius *)
  Hashtbl.iter
    (fun a b ->
       if a = 0
       then Printf.printf "%d virtual processors of unspecified color.\n" b
       else Printf.printf "%d virtual processors of color %d.\n" b a)
    ht;
  print_string "Enter a newline to stop ...\n";
  ignore (read_line ());;