let () = Random.self_init ();

type action = string

type grid_arg = {
  mutable width: int;
  mutable height: int;
}

type qudg_arg = {
  mutable width: float;
  mutable height: float;
  mutable radius: float;
  mutable r1: float;
  mutable p: float;
}

type er_prob = float (*between 0 and 1*)
type ba_m = int (*positive*)

type t = {
  mutable outputFile: string;
  mutable dotUDG: string;
  mutable dotUDGrad: string;
  mutable action: action;

  mutable n : int;
  mutable grid : grid_arg;
  mutable er : er_prob;
  mutable ba : ba_m;
  mutable qudg : qudg_arg;

  mutable silent : bool;
  mutable directed : bool;

  mutable _args : (string * Arg.spec * string) list;
  mutable _man : (string * (string list * action) list) list;

  mutable _others : string list;
  mutable _margin : int;
}

let usage_msg do_print_command tool =
  "gg is an experimental graph generator.\n"^
  (if do_print_command then 
    ("usage: " ^ tool ^ " <graph-kind> [<option>]*\n")
  else
    ("usage: "^tool^" [option]*\n"))

let print_usage output do_print_command tool = 
  Printf.fprintf output "%s%s" (usage_msg do_print_command tool) (
    if (do_print_command) then 
      "use -h to see the available <graph-kind>.\n\n" 
    else "use -h to see available <option>.\n"
  )


let (make_args : unit -> t) =
  fun () ->
  {
    outputFile = "";
    dotUDG = "";
    dotUDGrad = "";
    action = "void";

    n = -1;
    grid = {
      width = 0;
      height = 0;
    };
    er = 0.3;
    ba = 2;
    qudg = {
      width = 10.;
      height = 10.;
      radius = 3.;
      r1 = 2.;
      p = 0.5;
    };


    silent = false;
    directed = false;
    _args  = [];
    _man = [];

    _others  = [];
    _margin  = 18;
  }

let first_line b = (
  try (
    let f = String.index b '\n' in
    String.sub b 0 f
  ) with Not_found -> b
)
let unexpected s = (
  prerr_string ("unexpected argument \""^s^"\"");
  prerr_newline ();
  exit 2 (* unexpected argument *)
)


let printSpec args outChannel action (c, messageList) = (
  List.iter (fun (ml,action_type) ->
    if (action = action_type) then 
      let (m1, oth) = match ml with
       |  h::t -> (h,t)
       |  _ -> ("",[])
      in
      let t2 = String.make args._margin ' ' in
      let cl = 1 + String.length c in
      let t1 = if (cl < args._margin ) then
       String.make (args._margin - cl) ' '
      else
       "\n"^t2
      in
       Printf.fprintf outChannel "  %s%s%s" c t1 m1;
       List.iter (function x -> Printf.fprintf outChannel "\n%s%s" t2 x) oth ;
       Printf.fprintf outChannel "\n" ;
    ) messageList
)

let help args tool = (
  Printf.printf "%s" (usage_msg (args.action = "void") tool);
  (
    if (args.action = "void") then (
      Printf.printf "where <graph-kind> can be:\n";
      List.iter (printSpec args stdout "") [
        ("clique",[(["Generate a clique graph"],"")]);
        ("star",[(["Generate a star graph"],"")]);
        ("ring",[(["Generate a ring graph"],"")]);
        ("grid",[(["Generate a grid graph"],"")]);
        ("HC",[(["Generate a hyper-cube graph"],"")]);
        ("ER",[(["Generate a graph using the Erdos Renyi algo"],"")]);
        ("BA",[(["Generate a graph using the Barabasi–Albert algo"],"")]);
        ("tree",[(["Generate an acyclic graph (tree)"],"")]);
        ("UDG",[(["Generate a graph using the Unit Disc Graph algo"],"")]);
        ("QUDG",[(["Generate a graph using the Quasi UDG algo"],"")]);
      ];
      Printf.printf "\n";
    );
    Printf.printf "<option> can be:\n";
    List.iter (printSpec args stdout "void") (List.rev args._man);
    if (args.action <> "void") then (
      Printf.printf "\n";
      Printf.printf "<option> specific to this <graph-kind>:\n";
      List.iter (printSpec args stdout args.action) (List.rev args._man)
    ) else (
      Printf.printf "Use '%s <graph-kind> -h' to see specific options." tool;
      Printf.printf "\n";
    )
  );
  Printf.printf "\n";
  exit 0 
)


let (mkopt : t -> string list -> ?arg:string -> Arg.spec ->
     (string list * action) list -> unit) =
  fun opt ol ?(arg="") se ml ->
    let add_option o = opt._args <- (o, se, "")::opt._args in
     List.iter add_option ol ;
     let col1 = (String.concat ", " ol)^arg in
     opt._man <- (col1, ml)::opt._man

(*** User Options Tab **)
let (mkoptab : string array -> t -> unit) =
  fun argv args -> (

    mkopt args ["--output";"-o"] ~arg:" <string>"
      (Arg.String (fun s -> args.outputFile <- s))
      [(["Redirect stdout into a file"],"void")];

    let msg = "Set the node number in the graph" in
    mkopt args  ["--nodes-number";"-n"] ~arg:" <int>"
      (Arg.Int (fun n -> match args.action with
      | "grid" | "HC" | "void" -> unexpected "-n"
      | _ -> args.n <- n ))
      [([msg],"clique");([msg],"star");([msg],"ring");
      ([msg],"ER");([msg],"BA");([msg],"tree");([msg],"UDG");([msg],"QUDG")];

    mkopt args  ["--dimension";"-d"] ~arg:" <int>"
      (Arg.Int (fun n -> match args.action with
      | "HC"-> args.n <- n
      | _ -> unexpected "-d"))
      [(["Set the hyper-cube dimension"],"HC")];

    mkopt args  ["--width";"-w"] ~arg:" <float>"
      (Arg.Float (fun w -> match args.action with
      | "grid" -> args.grid.width <- (int_of_float w)
      | "UDG" | "QUDG" ->  args.qudg.width <- w
      | _ -> unexpected "-w" ))
      [(["Set the grid's width \n"],"grid");
      (["Set the UDG's terrain width";"10 by default.\n"],"UDG");
      (["Set the QUDG's terrain width";"10 by default.\n"],"QUDG")];

    mkopt args  ["--height";"-he"] ~arg:" <float>"
      (Arg.Float (fun h -> match args.action with
      | "grid" -> args.grid.height <- (int_of_float h)
      | "UDG" | "QUDG" ->  args.qudg.height <- h
      | _ -> unexpected "-he"))
      [(["Set the grid's height\n"],"grid");
      (["Set the UDG's terrain height";"10 by default.\n"],"UDG");
      (["Set the QUDG's terrain height";"10 by default.\n"],"QUDG")];


    mkopt args ["--edge-probability";"-p"]~arg:" <float>"
      (Arg.Float (fun p -> match args.action with
      | "ER" ->  args.er <- p
      | _ -> unexpected "-p"))
      [(["Set the edge appearing probability.";
        "Must be between 0 and 1, and is set to 0.3 by default\n"],"ER")];


    mkopt args ["--";"-m"]~arg:" <int>"
      (Arg.Int (fun m -> match args.action with
      | "BA" ->  args.ba <- m
      | _ -> unexpected "-m"))
      [(["Set the number of edge generated per additional node";
        "(2 by default)\n"],"BA")];

    mkopt args ["--radius";"-r"]~arg:" <float>"
      (Arg.Float (fun r -> match args.action with
      | "UDG" ->  args.qudg.radius <- r
      | _ -> unexpected "-r"))
      [(["Set the Unit Disc's radius around all nodes.";
         "3 by default.\n"],"UDG")];

    mkopt args ["--first-radius";"-r0"]~arg:" <float>"
      (Arg.Float (fun r -> match args.action with
      | "QUDG" ->  args.qudg.radius <- r
      | _ -> unexpected "-r0"))
      [(["Set the first radius around all nodes.";
         "3 by default.\n"],"QUDG")];

    mkopt args ["--second-radius";"-r1"]~arg:" <float>"
      (Arg.Float (fun r -> match args.action with
      | "QUDG" ->  args.qudg.r1 <- r
      | _ -> unexpected "-r1"))
      [(["Set the second radius around all nodes.";
         "2 by default.\n"],"QUDG")];

    mkopt args ["--probability"]~arg:" <float>"
      (Arg.Float (fun p -> args.qudg.p <- p))
      [(["Sets the probability of the nodes being neighbors when they are ";
           "inside the second";
         "radius, but not the first one.\n"
        ], "QUDG")];
    
    let msg = ["Create a DOT file to visualize the UDG plan.";
        "When it transformed into a PDF that takes the positioning tags into account"; 
        "(like 'neato' command from GraphViz), each node is visible at the coordinates";
        "where they were placed during execution.\n"] in
    mkopt args ["--dot-udg";"-du"]~arg:" <file>"
      (Arg.String (fun f -> match args.action with
                            | "UDG" | "QUDG" -> args.dotUDG <- f
                            | _ -> unexpected "-du"))
      [(msg,"UDG");(msg,"QUDG")];

    mkopt args ["--dot-udg-radius";"-dur"]~arg:" <file>"
      (Arg.String (fun f -> match args.action with
                            | "UDG" | "QUDG" -> args.dotUDGrad <- f
                            | _ -> unexpected "-dur"))
      [(["Create a DOT file to visualize the UDG plan.";
        "Same as the option '-du', but with the radius being also displayed.\n"],"UDG");
      (["Create a DOT file to visualize the UDG plan.";
        "Same as the option '-du', but with the two radiuses being ";
           "also displayed.\n"],"QUDG")];
      
    mkopt args ["--silent";"-s"]
      (Arg.Unit (fun () -> args.silent <- true))
      [(["be quiet"],"void")];

    mkopt args ["--directed";"-dir"]
      (Arg.Unit (fun () -> args.directed <- true))
      [(["Generate a directed graph."],"void")];

    mkopt args ["--help";"-h"]
      (Arg.Unit (fun () -> help args ((argv.(0))^(if args.action = "void" then "" 
        else " "^args.action))))
      [(["Print this help\n"],"void")];
  )

(* all unrecognized options are accumulated *)
let (add_other : t -> string -> unit) =
  fun opt s ->
    opt._others <- s::opt._others

let current = ref 1;;

let parse argv = (
  let save_current = !current in
  let args = make_args () in
  mkoptab argv args;
  try (
    (if (Array.length argv) = 1 then (
        print_usage stdout true argv.(0);
        exit 1 (* no command or invalid command *)));
    let possible_actions =
      ["clique";"star";"ring";"grid";"HC";"ER";"BA";"tree";"UDG";"QUDG"]
    in
    (
      if (List.mem argv.(1) possible_actions) then args.action <- argv.(1)
      else
      if (List.mem argv.(1) ["-h";"--help";"-help"]) then help args (argv.(0)) else
        (Printf.fprintf stderr "*** Error when calling '%s %s': No such command\n\n"
           (argv.(0)) (argv.(1));
         (print_usage stderr true argv.(0));
         exit 1 (* no command or invalid command *))
    );

    Arg.parse_argv ~current:current argv args._args (add_other args)
      (usage_msg false (argv.(0)^argv.(1)));
    current := save_current;
    args._others <- List.rev args._others;
    args
  )
  with
  | Arg.Bad msg ->
    Printf.fprintf stderr "*** Error when calling '%s': %s\n" (argv.(0))
      (first_line msg); 
    (print_usage stderr true argv.(0)); exit 3 (* bad argument *); 
  | Arg.Help _msg -> 
    help args argv.(0)
)