(** Extensions to the standard List module *)

(** Analogous to Array.iteri, iterate with index number *)
let iteri f lst =
  let rec loop n lst =
    match lst with
      [] -> ()
    | a :: b ->
	f n a;
	loop (n+1) b in
  loop 0 lst

(** Analogous to Array.mapi *)
let mapi f lst =
  let rec loop n lst res =
    match lst with
      [] -> (List.rev res)
    | a :: b -> loop (n+1) b ((f n a) :: res) in
  loop 0 lst []

(** Return the position of the first element satisfying PRED.
  May raise Not_found. *)
let index pred lst =
  let rec loop n tl =
    match tl with
      [] -> raise Not_found
    | x :: etc ->
	if pred x then n else loop (n+1) (List.tl tl) in
  loop 0 lst

let has = List.exists

(** Drop the first n elements of lst.  May raise Failure "tl". *)
let drop n lst =
  let rec loop lst n =
    if n = 0
    then lst
    else loop (List.tl lst) (n-1) in
  loop lst n

(** Return the first n elements of lst.  May raise Failure "hd". *)
let take n lst =
  let rec loop lst n result =
    if n = 0
    then List.rev result
    else loop (List.tl lst) (n-1) (List.hd lst :: result) in
  loop lst n []

(** Return a pair with the result of take and drop. *)
let takedrop n lst =
  let rec loop lst n result =
    if n = 0
    then (List.rev result), lst
    else loop (List.tl lst) (n-1) (List.hd lst :: result) in
  loop lst n []

(** Return the common prefix of two lists.
  
  prefix (=) \[1;2;3;5\] \[1;2;4;5\] -> \[1;2\] *)
let prefix eq l1 l2 =
  let rec loop prefix suff1 suff2 =
    match suff1, suff2 with
      h1 :: t1, h2 :: t2 when eq h1 h2 -> loop (h1 :: prefix) t1 t2
    | _, _ -> List.rev prefix in
  loop [] l1 l2

(** O(n^2) function to return the diff of two lists.  The resulting list has
  the same order as the original.

  diff (=) \[3;4;1;5;2\] \[1;3;5\] -> \[4;2\] *)
let diff eq a b =
  List.rev
    (List.fold_left
       (fun result a -> if List.exists (eq a) b then result else a :: result)
       [] a)

(** O(n) function to return the diff of two sorted lists.

  diff_sorted (-) \[1;2;3;4;5\] \[1;3;5\] -> \[2;4\] *)
let rec diff_sorted cmp a b =
  match a, b with
    [], b -> []
  | a, [] -> a
  | ahd :: atl, bhd :: btl ->
      let s = cmp ahd bhd in
      if s == 0 
      then diff_sorted cmp atl b
      else if s < 0
      then ahd :: (diff_sorted cmp atl b)
      else diff_sorted cmp a btl

let map3 fn a b c =
  List.map2 (fun a (b, c) -> fn a b c)
    a (List.combine b c)

let map4 fn a b c d =
  List.map2 (fun a (b, (c, d)) -> fn a b c d)
    a (List.combine b (List.combine c d))

let map5 fn a b c d e =
  List.map2 (fun a (b, (c, (d, e))) -> fn a b c d e)
    a (List.combine b (List.combine c (List.combine d e)))

let map6 fn a b c d e f =
  List.map2 (fun a (b, (c, (d, (e, f)))) -> fn a b c d e f)
    a (List.combine b (List.combine c (List.combine d (List.combine e f))))

let map7 fn a b c d e f g =
  List.map2 (fun a (b, (c, (d, (e, (f, g))))) -> fn a b c d e f g)
    a (List.combine b (List.combine c (List.combine d 
					 (List.combine e
					    (List.combine f g)))))

let map8 fn a b c d e f g h =
  List.map2 (fun a (b, (c, (d, (e, (f, (g, h)))))) -> fn a b c d e f g h)
    a (List.combine b (List.combine c (List.combine d
					 (List.combine e
					    (List.combine f
					       (List.combine g h))))))

(** Segment a list into a list of lists where each element satisfying
  PRED is the first element of the resulting sublist.  If there are
  elements at the beginning of LST that do not match PRED, they will
  appear in the first element of the resulting list.

  segment (fun x -> x = 0) \[1;0;2;3;0;0;5;6\] ->
  \[\[1\]; \[0;2;3\]; \[0\]; \[0;5;6\]\] *)
let segment pred lst =
  List.rev
    (List.map List.rev
       (List.fold_left
	  (fun result item ->
	     match result with
	       [] ->
		 [[item]]
	     | hd :: tl ->
		 if pred item
		 then [item] :: hd :: tl
		 else (item :: hd) :: tl)
	  [] lst))

(** Like segment, but elements that satisfy p are removed, and any
  resulting empty lists are also removed.

  segment_remove (fun x -> x = 0) \[1;0;2;3;0;0;5;6\] ->
  \[\[1\]; \[2;3\]; \[5;6\]\] *)
let segment_remove p lst =
  let result = segment p lst in
  let result = List.map (List.filter (fun x -> not (p x))) result in
  List.filter (fun x -> x <> []) result

(** Assign counting numbers to the elements of a list:
  enumerate \['a';'b';'a'\] -> \[(1, 'a'); (2, 'b'); (3, 'a')\] *)
let enumerate lst =
  let result = ref [] in
  ignore (List.fold_left
	    (fun n item -> result := (n, item) :: !result; n + 1)
	    1 lst);
  List.rev !result

(** Group a list into runs of equal elements wrt cmp.

  group compare \[1;1;1;2;2;3;4;4;5\] ->
	\[\[1;1;1\]; \[2;2\]; \[3\]; \[4;4\]; \[5\]\] *)
let group cmp lst =
  List.fold_left
    (fun result next ->
       match result with
	 (item :: tail) :: groups ->
	   if cmp item next = 0
	   then (next :: item :: tail) :: groups
	   else [next] :: (item :: tail) :: groups
       | _ -> [next] :: result)
    [] (List.rev lst)

(** Given a list of ('a, 'b list) pairs, distribute the value of
  the head over the elements of the tail.

  distribute \[1, \['a'; 'b'\]; 2, \['c'; 'd'\]\] ->
  	\[(1, 'a'); (1, 'b'); (2, 'c'); (2, 'd')\] *)
let distribute lst =
  List.fold_left
    (fun result (a, b) -> (List.map (fun x -> a, x) b) @ result)
    [] (List.rev lst)

(** Take a list, sort and group it over cmp, and then pass each group
    to f.  F should return a list of the same length as its argument.
    The elements returned by f are then restored to the original order
    of the argument list.

    Example: Uniquify a set of strings by adding a numerical suffix.

    let add_suffix n s = s ^ "-" ^ string_of_int (n+1) in
    let f =
      function
        [a] -> [a]
      | lst -> List2.mapi add_suffix lst in
    map_equiv compare f ["a";"b";"c";"b";"c";"d";"c";"b"] ->
      ["a"; "b-1"; "c-1"; "b-2"; "c-2"; "d"; "c-3"; "b-3"]

   If the EXIT_ON_UNIQUE flag is set an Exit exception will be raised
   if all the resulting equivalence groups are singletons.  This is
   handy when you are trying to uniquify a set by applying one
   mapping function after another. *)
let map_equiv ?(exit_on_unique=false) cmp f lst =
  let enum = mapi (fun i x -> (i, x)) lst in
  let ecmp = fun (_, a) (_, b) -> cmp a b in
  let groups = group ecmp (List.sort ecmp enum) in
  if exit_on_unique && List.length groups = List.length lst
  then raise Exit
  else
    let mapped_groups =
      List.map 
	(fun lst ->
	   let items = List.map snd lst in
	   let items = f items in
	   List.map2 (fun (n, _) x -> n, x) lst items)
	groups in
    List.map snd
      (List.sort
	 (fun (a, _) (b, _) -> a - b)
	 (List.flatten mapped_groups))

let uniq_stable cmp lst =
  Option.filter
    (map_equiv cmp
       (function
	    a :: etc -> Some a :: List.map (fun x -> None) etc
	  | [] -> [])
       lst)

(** Uniquify an already sorted list. *)
let uniq_sorted cmp lst =
  let result = ref [] in
  let rec loop tmp =
    match tmp with
      a :: b :: etc ->
	if cmp a b = 0
	then loop (b :: etc)
	else begin
	  result := a :: !result;
	  loop (b :: etc)
	end
    | [a] ->
	result := a :: !result;
	loop []
    | [] -> List.rev !result
  in loop lst

(** Sort a list according to the comparison function and then uniqify. *)
let uniq cmp lst = uniq_sorted cmp (List.sort cmp lst)

(** Enumerate equal elements, preserving order:
  map_count compare \['a';'b';'a';'c';'b';'a';'a';'c'\] ->
  \[(1,'a'); (1,'b'); (2,'a'); (1,'c'); (2,'b'); (3,'a'); (4,'a'); (2,'c')\] *)
let map_count cmp lst =
  map_equiv cmp (mapi (fun i x -> (i +1, x))) lst

(** Apply modify to the elements of list until they are all distinct
    over compare.  (Short lists only, this is a slow algorithm.) *)
let make_distinct compare modify lst =
  let count pred lst = List.length (List.filter pred lst) in
  let copies x lst = count (fun y -> compare x y = 0) lst in
  (* Return the list of each element's copy count *)
  let counts lst = List.map (fun x -> copies x lst) lst in
  (* We are done when every element has a count of 1 *)
  let finished lst =
    not (List.mem false
	   (List.map ((=) 1) (counts lst))) in
  let rec loop = function
      [] -> []
    | result when finished result -> result
    | result ->
	let sorted =
	  List.sort (fun x y -> copies y result - copies x result) result in 
	match sorted with
	  [] -> []
	| a :: etc -> loop (modify a :: etc) in
  loop lst

let is_uniq cmp lst = List.length (uniq cmp lst) = List.length lst

let rec seq a b =
  match a with
    x when x > b -> []
  | x -> x :: seq (x + 1) b

let list_init f i = List.map f (seq 0 (i - 1))

(** [[1;2]; [3;4]] -> [[1;3]; [2;4]] *)
let transpose (a : 'a list list) : 'a list list =
  List.map
    (fun i -> List.map (fun l -> List.nth l i) a)
    (seq 0 (List.length (List.hd a) - 1))

let extend length init lst =
  (*Log.put ~v:0 ("list_extend from " ^ string_of_int (List.length lst) ^ " to " ^ string_of_int length);*)
  lst @ (List.map (fun n -> init) (seq (List.length lst) (length - 1)))

let rec intersperse elem lst =
  match lst with
    [] -> []
  | [a] -> [a]
  | a :: etc -> a :: elem :: intersperse elem etc

let split3 triples =
  List.fold_left
    (fun (alist, blist, clist) (a, b, c) -> a :: alist, b :: blist, c :: clist)
    ([],[],[]) (List.rev triples)

let sorted_intersection cmp lists =
  let rec discard_head pred lst =
    match lst with
      [] -> []
    | hd :: tl ->
	if pred hd
	then discard_head pred tl
	else hd :: tl in
  let rec loop result tails =
    match tails with
      [] -> []
    | [] :: _ -> List.rev result
    | (head1 :: tail1) :: tails ->
	let cmphead1 = cmp head1 in
	let toosmall = (fun x -> cmphead1 x > 0) in
	(* Discard any elements that are too small *)
	let tails' = List.map (discard_head toosmall) tails in
	if List.mem [] tails'
	then result
	else
	  let justright = (fun x -> cmphead1 x = 0) in
	  let heads = List.map List.hd tails' in
	  let cmps = List.map justright heads in
	  if not (List.mem false cmps)
	    (* All elements matched head1, add it to the result *)
	  then loop (head1 :: result) (tail1 :: (List.map List.tl tails'))
	    (* Discard any elements that matched, since some didn't *)
	  else loop result (tail1 :: (List.map (discard_head justright) tails')) in
  loop [] lists

let intersection cmp lists =
  sorted_intersection cmp (List.map (List.sort cmp) lists)

let filter_out pred lst =
  List.filter (fun x -> not (pred x)) lst

let format ?(prefix="") string_of_elem lst =
  let sep = ";" ^ (if String.contains prefix '\n' then prefix else "") ^ " " in

  let rec format_list_tail string_of_elem =
  function
      [] -> "]"
    | [elem] -> string_of_elem elem ^ "]"
    | elem :: etc -> string_of_elem elem ^ sep ^ format_list_tail string_of_elem etc in

  prefix ^ "[" ^ format_list_tail string_of_elem lst