(**************************************************************************)
(*                                                                        *)
(*  Copyright (C) Johannes Kanig, Stephane Lescuyer                       *)
(*  Jean-Christophe Filliatre, Romain Bardou and Francois Bobot           *)
(*                                                                        *)
(*  This software is free software; you can redistribute it and/or        *)
(*  modify it under the terms of the GNU Library General Public           *)
(*  License version 2.1, with the special exception on linking            *)
(*  described in file LICENSE.                                            *)
(*                                                                        *)
(*  This software 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.                  *)
(*                                                                        *)
(**************************************************************************)

(* This is code which has been taken from
(*  Ocamlgraph: a generic graph library for OCaml                         *)
(*  Copyright (C) 2004-2008                                               *)
(*  Sylvain Conchon, Jean-Christophe Filliatre and Julien Signoles        *)
 and has been modified since then by the Mlpost authors *)

module M = struct
  type t = int * int
  let equal = Pervasives.(=)
  let compare = Pervasives.compare
  let hash = Hashtbl.hash
end

type elt = float * float
type inputelt = M.t

module H = Hashtbl.Make(M)

type cell = {
  mutable c : int;
  mutable data : elt;
  mutable father : cell
}

type t = cell H.t (* a forest *)

let init l = 
  let h = H.create 997 in
  List.iter 
    (fun ((a,b) as x) ->
      let t = float_of_int a, float_of_int b in
       let rec cell = { c = 1; data = t; father = cell } in 
       H.add h x cell) 
    l;
  h

let rec find_aux cell = 
  if cell.father == cell then 
    cell
  else 
    let r = find_aux cell.father in 
    cell.father <- r; 
    r

let find x h = (find_aux (H.find h x)).data

let avg ra rb = 
  let ax,ay = ra.data and bx,by = rb.data in
  let ac = float_of_int ra.c and bc = float_of_int rb.c in
  let z = ac +. bc in
  (ac *. ax +. bc *. bx) /. z, (ac *. ay +. bc *. by) /. z

let union x y h = 
  let rx = find_aux (H.find h x) in
  let ry = find_aux (H.find h y) in
  if rx != ry then begin
    if rx.c > ry.c then begin
      ry.father <- rx;
      rx.data <- avg rx ry;
      rx.c <- rx.c + ry.c
    end else if rx.c < ry.c then begin
      rx.father <- ry;
      ry.data <- avg rx ry;
      ry.c <- rx.c + ry.c
    end else begin
      ry.father <- rx;
      rx.data <- avg rx ry;
      rx.c <- rx.c + ry.c
    end
  end

let fold_classes f acc h = 
  let seen = Hashtbl.create 127 in
  H.fold
    (fun _ v acc ->
      let r = find_aux v in
      let d = r.data in
      if Hashtbl.mem seen r then acc
      else (Hashtbl.add seen r () ;f d acc)
    ) h acc
