(* Copyright 2004 Lodewijk Voge / Stichting Wireless Leiden,
 *                All Rights Reserved.
 * License: http://www.wirelessleiden.nl/LICENSE. *)

(* \abstract{Experimental channel assigning program using an evolutionary algorithm.} *)
(*s Introduction
  The scoring function is simple:
  \begin{itemize}
  \item the score of a combination of two interfaces is 1 if there's two or
      more channels between them, -1 if there's 1 channel between them, -5 if
      there's no channels between them (they're on adjacent channels) and -10
      if they're on the same channel
  \item the score for a node is the sum of the scores of all the combinations
      of interfaces for that node, minus any node-specific penalties (eg.
      channels 7, 8 and 9 are unusable and score -1000 on node $x$), scaled by
      the number of interfaces to give more weight to larger nodes (the
      assumption being that larger nodes are more important nodes)
  \item the total score of the network is the sum of the score of the nodes,
      minus any network-wide penalties (eg. the omni's for node $x$ and node
      $y$ can see eachother, so they should be apart)
  \end{itemize}

   Supported special cases:
  \begin{itemize}
  \item on node $x$, interfaces $y$ and $z$ are well separated and can live
       on the same channel
  \item the link between $x$ on $y$ and $z$ on $w$ is very important, make
        sure it is well separated on both ends of the link
  \item interface $x$ on node $y$ will bother interface $z$ on node $w$ and
        need to be on separate channels
  \item on node $x$, channels $y$, $z$ and $w$ are not usable
  \item node $x$ $y$ and $z$ should be treated as one node
  \item if at all possible, do not put accesspoints above channel 11
  \end{itemize}
*)
(*s Installation

 \begin{itemize}
 \item install an O'Caml compiler. /usr/ports/lang/ocaml/ in FreeBSD, ocaml in
       Debian.
 \item compile with
\begin{verbatim}
    $ ocamlopt -o foo str.cmxa channelea.ml
\end{verbatim}
 \item run with
\begin{verbatim}
   $ ./foo f
\end{verbatim}
   where f is the result of running prepconf.py on a file with a list of 
   paths to the wleiden.conf's to consider.
 \end{itemize}

*)
open Pxp_yacc
open Pxp_types

module StringMap = Map.Make(String)

(*s Commandline settable values *)
let fname = ref "nodes.xml"
let numthreads = ref 1
let special_conf = ref "special.conf"

(*s Constants *)
(* How large a population should the system maintain? *)
let population_size = 20
(* How long should the system iterate after an improvement? *)
and max_stagnant_iterations = 200
(* How many little iterations before joining threads? *)
and little_iterations = 10
(* What is the chance for an ESSID-to-channel assignment to mutate to a 
   random channel? *)
and mutation_rate = 0.1
(* The most basic score table *)
and scoretable = [ (<=) 2,  1;
		   (==) 2, -30;
		   (==) 1, -70;
		   (==) 0, -100 ]

(*s Type definitions *)
(* Note that Caml has trouble with mutually recursive data structures. You
   can define them, you just can't ever instantiate them. This is why the
   fields in wi are all loose references by way of strings *)
type wi = {
	wi_name: string;
	wi_nodename: string;
	wi_essid: string
}
type group = {
	group_essid: string;
	mutable group_wis: wi list
}
type node = {
	node_name: string;
	node_wis: wi list
}
(* A configuration is an assignment of groups, identified by essid, to a
   channel, plus a score. The code should be careful not to use the score
   between mutating an re-evaluating. *)
type configuration = {
	mutable score: int;
	mutable conf: int StringMap.t
}
type 'a maybe = Nothing | Just of 'a

(*s Globals *)

(* The global map from node name to node struct *)
let nodes = ref StringMap.empty
(* The global map from essid to group struct *)
let groups = ref StringMap.empty

(* Now the hashes for the special cases *)
(* Map from nodename to a list of winame's indicating the wi's that
   don't interfere with eachother for that node. *)
let nointerference = ref StringMap.empty
(* List of (nodename1, winame1, nodename2, winame2) tuples indicating a very
   important link that should be well-separated on both ends *)
let importantlinks = ref []
(* Hash mapping from nodename to a list of unusable channels for that node *)
let unusable = ref StringMap.empty
(* List of (nodename1, winame1, nodename2, winame2) tuples indicating two
   interfering interfaces on different nodes *)
let interference = ref []
(* nodename -> (winame, channel) *)
let pinned = Hashtbl.create 4

(* Run the given diff against the given scoretable and return the score *)
let rec runtable t diff =
     match t with
	[]		-> assert false
     | (cond, s)::xs	-> if (cond diff) then s
			   else runtable xs diff

(*s Convenience functions *)

(* Function composition. *)
let compose f g = fun x -> f(g(x))
let ($) = compose
(* Turn two individual values into a tuple *)
let maketuple a b = (a, b)
(* Shorthand for List.hd *)
let head = List.hd
(* Shorthand for List.tail *)
let tail = List.tl
let even x = (x mod 2) == 0
let just x = match x with
	       Nothing -> assert false
	     | Just s -> s
(* Given a map, return all the keys as a list *)
let keys t = StringMap.fold (fun k d a -> k::a) t []
(* Copy one array into the other *)
let copyarray src dest = Array.blit src 0 dest 0 (Array.length src)
(* get the given key from the given hash. if it isn't there, return an empty list *)
let defget map key = try StringMap.find key map
		      with Not_found -> []
(* *)
let rec fold_n_times f n a = match n with
	  0 -> a
	| i -> fold_n_times f (i - 1) (f a)

(* Is the given element in the given list? uses compare, so it works on
   strings as well *)
let in_list l e = try
			let _ = List.find ((==) 0 $ compare e) l in
			true
		  with Not_found -> false

(* Given a list, return a list of pairs with all possible combinations of 
   items from the given list *)
let rec combinations l =
	match l with
	  []	-> []
	| x::xs	-> (List.map (maketuple x) xs)@(combinations xs)

(* Given a configuration and two wi's, return the score *)
let wi_score c unusable nointerference wi1 wi2 =
	let channel1 = c wi1.wi_essid in
	let channel2 = c wi2.wi_essid in
	let diff = abs (channel1 - channel2) in
	let is_unusable = in_list unusable in
	if (is_unusable channel1) || (is_unusable channel2) then -10000
	else if (in_list nointerference wi1.wi_name) &&
		(in_list nointerference wi2.wi_name) then 1
	else runtable scoretable diff

(* Given a configuration and a node, return the score. this is simply the sum
   of the scores of all the combinations of interfaces, written down as a fold
   for efficiency *)
let node_score c n =
	let nointerference_ = defget !nointerference n.node_name in
	let unusable_ = defget !unusable n.node_name in
	let f a (wi1, wi2) = a + (wi_score c unusable_ nointerference_ wi1 wi2) in
	let base_score = List.fold_left f 0 (combinations n.node_wis) in
	base_score * (List.length n.node_wis)

(* Score the given pair of interferent interfaces against the given
   configuration *)
let score_interference c (nodename1, winame1, nodename2, winame2) = 
	let node1 = StringMap.find nodename1 !nodes in
	let node2 = StringMap.find nodename2 !nodes in
	let f winame = fun wi -> (compare wi.wi_name winame) == 0 in
	let wi1 = List.find (f winame1) node1.node_wis in
	let wi2 = List.find (f winame2) node2.node_wis in
	let channel1 = c wi1.wi_essid in
	let channel2 = c wi2.wi_essid in
	let diff = abs (channel1 - channel2) in
	let res = runtable scoretable diff in
	res

(* Given a list of nodes and a configuration, return the score for the whole
   configuration. This is the sum of the scores for all nodes, plus the sum
   of the scores for all user-specified interferent pairs of interfaces. *)
let score_configuration ns c =
	let mapper n = StringMap.find n c in
	let f1 _ n a = a + (node_score mapper n) in
	let nodescores = StringMap.fold f1 ns 0 in
	let f2 a i = a + (score_interference mapper i) in
	let interference_score = List.fold_left f2 0 !interference in
	nodescores + interference_score

(* Return a random configuration. For some reason, if this function accesses
   the global 'groups' hash instead of getting it passed in from above, that
   hash is empty. *)
let random_configuration groups evaluate =
	let c = StringMap.fold (fun essid _ map ->
			let channel = 1 + (Random.int 12) in
			StringMap.add essid channel map) groups StringMap.empty in
	{ score = evaluate c; conf = c }

let print_conf conf = 
	let channel wi = string_of_int (StringMap.find wi.wi_essid conf) in
	let print_wi wi = wi.wi_name ^ ": " ^ (channel wi) in
	let wis node = List.fold_left (fun a wi -> a ^ " " ^ (print_wi wi))
				      "" node.node_wis in
	StringMap.iter (fun name node ->
		print_string (name ^ ": " ^ wis node ^ "\n")) !nodes

(* n-point crossover operator. pick n points along the length of the parents, 
   produce a child by copying from one parent, switching parents when hitting a
   chosen crossover point *)
let crossover n c1 c2 = 
	let keys1 = keys (c1.conf) in
	let numkeys1 = List.length keys1 in
	let pickpoint i = (if even i then c1.conf else c2.conf),
			  (if i < n then (Random.int numkeys1) else numkeys1) in
	let crosspoints = Array.init (n + 1) pickpoint in
	let result = ref StringMap.empty in
	let i = ref 0 in
	Array.sort (fun a b -> compare (snd a) (snd b)) crosspoints;
	Array.iter (fun (h, p) ->
		while !i < p do
			let key = List.nth keys1 !i in
	
			let d = StringMap.find key h in
			result := StringMap.add key d !result;
			incr i
		done) crosspoints;
	assert ((List.length (keys !result)) == (List.length keys1));
	{ score = 0; conf = !result }

(* Generalized evolutionary algorithm driver. 
      initialize: () -> configuration array
      recombine:
      mutate: configuration array -> configuration array
      evaluate: configuration array -> configuration array
      select: configuration array -> configuration array
      
    and the result is the best configuration found *)
let evolutionary_algorithm numthreads initialize recombine mutate evaluate select = 
	let orig_population = (evaluate $ initialize) () in
	let last_high_score = ref orig_population.(0).score in
	let iterations_since_new_high_score = ref 0 in
	let generation = ref 0 in
	(*let iterate = recombine $ mutate $ evaluate $ select in*)
	let iterate = select $ evaluate $ mutate $ recombine in 

	let population = ref orig_population in
	let mutex = Mutex.create () in

	while !iterations_since_new_high_score < max_stagnant_iterations do
		
		let threads = Array.init numthreads (fun i ->
			Thread.create (fun _ ->
				let pop = fold_n_times iterate
						       little_iterations
						       !population in
				Mutex.lock mutex;
				if pop.(0).score > !population.(0).score then
				  population := pop;
				Mutex.unlock mutex) ()) in
		Array.iter Thread.join threads;

		let high_score = !population.(0).score in
		if high_score > !last_high_score then begin
			last_high_score := high_score;
			iterations_since_new_high_score := 0
		end;
		assert (!last_high_score >= high_score);
		print_string "bigiteration ";
		print_int !generation;
		print_string ": highscore ";
		print_int !last_high_score;
		print_newline();
		incr iterations_since_new_high_score;
		incr generation
	done;
	!population.(0)

(*s Parsing *)

(* Given a filename, return a list of all the lines in the file with the given
   filename. Don't count on the order of the lines in the result. *)
let snarf_lines fname =
	let infile = open_in fname in
	let result = ref [] in
	try
		while true do
			result := (input_line infile)::!result
		done;
		!result	(* never gets here *)
	with End_of_file -> !result

(* Read the main input from the given filename *)
let parse_file fname =
	let getattr e n =
		match e#attribute n with
		  Pxp_types.Value s -> s
		| _ -> "" in
	let do_element n e =
		match e#node_type with
		  Pxp_document.T_element "wireless" ->
			let nodename = getattr n "name" in
		  	let essid = getattr e "essid" in
			let new_wi = { wi_name = getattr e "iface";
				       wi_nodename = nodename;
				       wi_essid = essid; } in
			let _ = try
					let group = StringMap.find essid !groups in
					group.group_wis <- new_wi::group.group_wis;
				with Not_found ->
					let group = { group_essid = essid;
						      group_wis = [ new_wi ] } in
					groups := StringMap.add essid group !groups in
			Just new_wi;
		  | _ -> Nothing in
	let do_node map n = 
		if getattr n "status" = "up" then begin
			let nodename = String.lowercase (getattr n "name") in
			let wis = ref [] in
			n#iter_nodes (fun e -> 
				match (do_element n e) with
				  Just wi -> wis := wi::(!wis)
				| Nothing -> ());
			let sorted_wis = List.sort compare (!wis) in
			let node = { node_name = nodename; node_wis = sorted_wis } in
			StringMap.add nodename node map
		end else map in
	let d = Pxp_yacc.parse_document_entity Pxp_yacc.default_config (Pxp_yacc.from_file fname) Pxp_yacc.default_spec in
	nodes := List.fold_left do_node StringMap.empty
			(Pxp_document.find_all_elements "node" d#root)

(* The parsers for the special case components *)

(* The first field is the nodename, the rest are interface names *)
let parse_nointerference fs = 
	nointerference := StringMap.add (head fs) (tail fs) !nointerference

(* Read four fields from the given list and add them as a tuple to the given
   list reference *)
let parse_quadruplet l fs = 
	let f = List.nth fs in
	l := (f 0, f 1, f 2, f 3)::!l

(* The first field is the nodename, the rest are channels.*)
let parse_unusable fs =
	let channels = List.map int_of_string (tail fs) in
	unusable := StringMap.add (head fs) channels !unusable

(* The first field is the supernode name, the rest are the names of the
   subnodes. Construct a new node for the supernode and remove the subnodes
   from the nodes map *)
let parse_supernode fs =
	let nodename = head fs in
	let subnodenames = tail fs in
	let subnodes = List.map (fun n -> StringMap.find n !nodes) subnodenames in
	nodes := List.fold_left (fun map n ->
			StringMap.remove n map) !nodes subnodenames;
	let prefixed_wis n = List.map (fun w -> { w with wi_name = n.node_name ^ "." ^ w.wi_name}) n.node_wis in
	let wis = List.fold_left (fun a s -> a@(prefixed_wis s)) [] subnodes in
	let node = { node_name = nodename; node_wis = wis } in
	nodes := StringMap.add nodename node !nodes

let parse_pinned fs = 
	let field = List.nth fs in
	let nodename = field 0 in
	let winame = field 1 in
	let channel = int_of_string (field 2) in
	Hashtbl.add pinned nodename (winame, channel)

let parse_special_conf fname =
	let spacere = Str.regexp " " in
	let functable = [ "nointerference", parse_nointerference;
			  "important", parse_quadruplet importantlinks;
			  "interference", parse_quadruplet interference;
			  "unusable", parse_unusable;
			  "supernode", parse_supernode ] in
	let do_line fs = (List.assoc (head fs) functable) (tail fs) in
	List.iter (do_line $ Str.split spacere) (snarf_lines fname)

(*s Main *)

let argopts = [
	"-p", Arg.Set_int numthreads, "Number of threads";
	"-c", Arg.Set_string special_conf, "Special configfile"
]

let _ = 
	Arg.parse argopts (fun s -> fname := s) "channelea";
	parse_file !fname;
	parse_special_conf !special_conf;

	Random.self_init();
	let evaluate_map = score_configuration !nodes in
	let initialize () = Array.init population_size (fun _ -> random_configuration !groups evaluate_map) in
	let recombine pop = pop in
	(*
		let numoffspring = Random.int population\_size in
		let children = Array.init numoffspring (fun \_ -> 
			let father = pop.(Random.int population\_size) in
			let mother = pop.(Random.int population\_size) in
			crossover 2 father mother) in
		Array.append pop children in *)
	let maxchannel essid =
		let group = StringMap.find essid !groups in
		if (List.length group.group_wis) == 1 then 11
		else 13 in
	let mutate_conf conf =
		StringMap.mapi (fun essid channel ->
			let f = Random.float 1.0 in
			if f < mutation_rate then
				1 + (Random.int (maxchannel essid))
			else channel) conf in
	let mutate population =
		let mutants = Array.map (fun c -> let conf = mutate_conf c.conf in
						  { score = evaluate_map conf;
						    conf = conf}) population in
		Array.append population mutants in
	let evaluate population =
		Array.iter (fun c -> c.score <- evaluate_map c.conf) population;
		population in
	let select p = 
		Array.sort (fun a b -> compare b.score a.score) p;
		(*shuffle p;*)
		Array.sub p 0 population_size in
	let best = evolutionary_algorithm !numthreads initialize recombine mutate evaluate select in
	print_conf best.conf