(* a few constants *)
let population_size = 20
and max_stagnant_iterations = 10000
and mutation_rate = 0.05;;

(* the 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;
};;

let nodes = Hashtbl.create 4;;
let groups = Hashtbl.create 4;;

(* some convenience functions *)
let compose f g = fun x -> f(g(x));;
let ($) = compose;;
let maketuple a b = (a, b);;
let head = List.hd;;
let tail = List.tl;;
(* given a hashtable, return all the keys as a list *)
let keys t = Hashtbl.fold (fun k d a -> k::a) t [];;
(* given a hashtable, return all the values as a list *)
let values t = Hashtbl.fold (fun k d a -> d::a) t [];;
let copyarray src dest = Array.blit src 0 dest 0 (Array.length src);;

(* 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 wi1 wi2 =
	let scoretable = [ ((<=) 2,  1);
			   ((==) 2, -1);
			   ((==) 1, -5);
			   ((==) 0, -10) ] in
	let channel1 = c wi1.wi_essid in
	let channel2 = c wi2.wi_essid in
	let diff = abs (channel1 - channel2) in
	let rec runtable t = match t with
				[]		-> assert false
			     | (cond, s)::xs	-> if (cond diff) then s
						   else runtable xs in
	runtable scoretable;;

(* 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 foldfunc acc (wi1, wi2) = acc + (wi_score c wi1 wi2) in
	let base_score = List.fold_left foldfunc 0 (combinations n.node_wis) in
	base_score * (List.length n.node_wis);;

let score_configuration c (ns: node list) =
	let foldfunc acc n = acc + (node_score (Hashtbl.find c) n) in
	let nodescores = List.fold_left foldfunc 0 ns in
	nodescores;;

(* given a filename, return a list of all the lines in the file with the given
   filename *)
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

let parse_pair nodename (wname, gname) = 
	let new_wi = { wi_name = wname; wi_nodename = nodename; wi_essid = gname} in
	let foo = try
			let group = Hashtbl.find groups gname in
			group.group_wis <- new_wi::group.group_wis;
		  with Not_found ->
			let group = { group_essid = gname; group_wis = [ new_wi ] } in
			Hashtbl.add groups gname group in
	new_wi

let parse_fields fields = 
	let nodename = head fields in
	let rec makepairs l = match l with
				[]		-> []
			      | x::[]		-> assert false
			      | a::b::xs	-> (a, b)::(makepairs xs) in
	let wis = List.map (parse_pair nodename) (makepairs (tail fields)) in
	let sorted_wis = List.sort compare wis in
	let node = { node_name = nodename; node_wis = sorted_wis } in
	Hashtbl.add nodes nodename node

let parse_file fname =
	let spacere = Str.regexp " " in
	List.iter (parse_fields $ (Str.split spacere)) (snarf_lines fname)
	;;

(** 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 =
	let conf = Hashtbl.create 30 in
	Hashtbl.iter (fun k _ -> Hashtbl.add conf k (1 + (Random.int 12))) groups;
	conf

(* Mutate the configuration in the given population at the given offset *)
let mutate p i = 
	Hashtbl.iter (fun essid _ -> let f = Random.float 1.0 in
				     let group = Hashtbl.find groups essid in
				     let maxchannel = if (List.length group.group_wis) == 1 then 11
						      else 13 in
				     if (f < mutation_rate) then
					     Hashtbl.replace p.(i) essid (1 + (Random.int maxchannel))) p.(i);;

let print_conf conf = 
	let print_wi wi = wi.wi_name ^ ": " ^ (string_of_int (Hashtbl.find conf wi.wi_essid)) in
	let wis node = List.fold_left (fun acc wi -> acc ^ " " ^ (print_wi wi)) "" node.node_wis in
	let sorted_nodes = List.sort (fun a b -> compare (a.node_name) (b.node_name)) (values nodes) in
	List.iter (fun n -> print_string (n.node_name ^ ":" ^ (wis n) ^ "\n")) sorted_nodes

let parse_nodeclusters fname = 
	let spacere = Str.regexp " " in
	(* handle a row of fields. the first field is the supernode name, the
	   rest are the subnode names. create a new node for the supernode,
	   stuff all the wi's of the subnodes under it, names prefixed with
	   their original node's names for clarity, and remove the subnodes
	   from the nodes hashtable *)
	let do_fields f = let nodename = head f in
			  let subnodenames = tail f in
			  let subnodes = List.map (Hashtbl.find nodes) subnodenames in
			  List.iter (Hashtbl.remove 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
			  Hashtbl.add nodes nodename node in
	List.iter (do_fields $ (Str.split spacere)) (snarf_lines fname);;

let main = 
	parse_file Sys.argv.(1);
	parse_nodeclusters "coupled.conf";
	Random.self_init();
	let population = Array.init population_size (fun _ -> random_configuration groups) in
	let last_high_score = ref (-1000000) in
	let iterations_since_new_high_score = ref 0 in
	let generation = ref 0 in
	let all_nodes = values nodes in
	while !iterations_since_new_high_score < max_stagnant_iterations do
		(* mutate the population *)
		for i = 0 to (population_size / 2 - 1) do
			let i2 = i + population_size / 2 in
			population.(i2) <- Hashtbl.copy population.(i);
			mutate population (i2)
		done;
		(* sort the populations according to score. to do this, make
		   a list of (score, solution) tuples *)
		let score_pop = Array.map (fun c -> ((score_configuration c all_nodes), c)) population in
		(* sort on the first field *)
		Array.sort (fun x y -> compare (fst y) (fst x)) score_pop;
		(* extract the, now sorted, configuration and put it into population *)
		copyarray (Array.map snd score_pop) population;
		(* now look at the best score and update the highscore if
		   necessary *)
		let high_score = fst score_pop.(0) 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);
		if (!generation mod 10) == 0 then begin
			print_int !generation;
			print_string ": ";
			print_int !last_high_score;
			print_newline();
		end;
		incr iterations_since_new_high_score;
		incr generation
	done;
	print_conf population.(0);;

main