(**************************************************************************************)
(* Copyright (C) 2008 Stefano Zacchiroli <zack@debian.org> and                        *)
(* Jaap Boender <boender@pps.jussieu.fr> and                                          *)
(* Pietro Abate <pietro.abate@pps.jussieu.fr>                                         *)
(*                                                                                    *)
(*  This library is free software: you can redistribute it and/or modify              *)
(*  it under the terms of the GNU Lesser General Public License as                    *)
(*  published by the Free Software Foundation, either version 3 of the                *)
(*  License, or (at your option) any later version.  A special linking                *)
(*  exception to the GNU Lesser General Public License applies to this                *)
(*  library, see the COPYING file for more information.                               *)
(**************************************************************************************)

(*
 * More info about small world networks
 * http://en.wikipedia.org/wiki/Small-world_network
 *)

(** Small World analisys *)

open Graph
open ExtLib
open Dose_common

include Util.Logging (struct
  let label = "dose_algo.statistics"
end)

module Make (G : Sig.I) = struct
  module VS = Set.Make (G.V)
  module UndG = Imperative.Graph.Concrete (G.V)

  let undirect g =
    let g2 = UndG.create () in
    G.iter_vertex
      (fun v ->
        UndG.add_vertex g2 v ;
        G.iter_succ (fun v' -> UndG.add_edge g2 v v') g v ;
        G.iter_pred (fun v' -> UndG.add_edge g2 v' v) g v)
      g ;
    g2

  (* www.cs.cornell.edu/home/kleinber/...book/networks-book-ch20.pdf *)
  let clustering_coefficient graph vertex =
    let neighbours = G.succ graph vertex in
    let n = List.length neighbours in
    if n = 0 then 0.0
    else if n = 1 then 1.0
    else
      let n_edges =
        List.fold_left
          (fun old_sum v ->
            old_sum
            + List.fold_left
                (fun old_sum' v' ->
                  if G.mem_edge graph v v' && v <> v' then old_sum' + 1
                  else old_sum')
                0
                neighbours)
          0
          neighbours
      and max_edges = if G.is_directed then n * (n - 1) else n * (n - 1) / 2 in
      float_of_int n_edges /. float_of_int max_edges

  let average_distance graph vertex =
    let rec add_successors distance visited vertices =
      (* Add successors breadth-first, we want the shortest path we can find *)
      let succs = ref [] in
      let (n, sum) =
        List.fold_left
          (fun (old_n, old_sum) v ->
            if not (VS.mem v !visited) then (
              visited := VS.add v !visited ;
              succs := G.succ graph v :: !succs ;
              (old_n + 1, old_sum + distance))
            else (old_n, old_sum))
          (0, 0)
          vertices
      in
      let sf = List.flatten !succs in
      if sf <> [] then
        let (n', sum') = add_successors (distance + 1) visited sf in
        (n + n', sum + sum')
      else (n, sum)
    in
    let visited = ref (VS.singleton vertex) in
    let (n, sum) = add_successors 1 visited (G.succ graph vertex) in
    if sum = 0 then 0.0 else float_of_int sum /. float_of_int n

  module MSin = Map.Make (struct
    type t = G.V.t * G.t ref

    let compare (v1, _) (v2, g) = G.in_degree !g v1 - G.in_degree !g v2
  end)

  module MSout = Map.Make (struct
    type t = G.V.t * G.t ref

    let compare (v1, _) (v2, g) = G.out_degree !g v1 - G.out_degree !g v2
  end)

  let _avgdegree = ref None

  let _outdata = ref None

  let _indata = ref None

  let _outdatadegree = ref MSout.empty

  let _indatadegree = ref MSin.empty

  let degree graph =
    let nv = G.nb_vertex graph in
    let outmax = ref 0 in
    let inmax = ref 0 in
    let outh = Hashtbl.create 1031 in
    let inh = Hashtbl.create 1031 in
    (* how many vertices have degree X *)
    let add h d =
      if d = 0 then ()
      else
        try Hashtbl.replace h d (Hashtbl.find h d + 1)
        with Not_found -> Hashtbl.add h d 1
    in
    let total =
      G.fold_vertex
        (fun v sum ->
          let outdeg = G.out_degree graph v in
          let indeg = G.in_degree graph v in
          add outh outdeg ;
          add inh indeg ;
          _indatadegree := MSin.add (v, ref graph) indeg !_indatadegree ;
          _outdatadegree := MSout.add (v, ref graph) outdeg !_outdatadegree ;
          sum + indeg)
        graph
        0
    in
    (float_of_int total /. float_of_int nv, !outmax, !inmax, outh, inh)

  let computeDegree graph =
    if Option.is_some !_avgdegree then ()
    else
      let (avdeg, _maxout, _maxin, outdata, indata) = degree graph in
      _avgdegree := Some avdeg ;
      _outdata := Some outdata ;
      _indata := Some indata

  let maxOutDegree graph =
    computeDegree graph ;
    snd (MSout.max_binding !_outdatadegree)

  let maxInDegree graph =
    computeDegree graph ;
    snd (MSin.max_binding !_indatadegree)

  let averageDegree graph =
    computeDegree graph ;
    Option.get !_avgdegree

  let zdp graph =
    G.fold_vertex
      (fun v sum ->
        if G.in_degree graph v = 0 && G.out_degree graph v = 0 then sum + 1
        else sum)
      graph
      0

  let scatteredPlotIn graph =
    computeDegree graph ;
    Option.get !_indata

  let scatteredPlotOut graph =
    computeDegree graph ;
    Option.get !_outdata

  let scatteredPlotBoth graph =
    let add h i o =
      try Hashtbl.replace h (i, o) (Hashtbl.find h (i, o) + 1)
      with Not_found -> Hashtbl.add h (i, o) 1
    in
    let h = Hashtbl.create 1031 in
    G.iter_vertex
      (fun v -> add h (G.in_degree graph v) (G.out_degree graph v))
      graph ;
    h

  (* http://en.wikipedia.org/wiki/Centrality *)
  let centralityDegree graph fd =
    let n = float_of_int (G.nb_vertex graph) in
    let cd v = float_of_int v /. (n -. 1.0) in
    let m =
      G.fold_vertex
        (fun v max ->
          let s = List.length (fd graph v) in
          let m = cd s in
          if m > max then m else max)
        graph
        0.0
    in
    let c =
      G.fold_vertex
        (fun v sum ->
          let s = List.length (fd graph v) in
          sum +. (m -. cd s))
        graph
        0.0
    in
    c /. (n -. 2.0)

  let centralityOutDegree graph = centralityDegree graph G.succ

  let centralityInDegree graph = centralityDegree graph G.pred

  let clustering graph =
    let n = float_of_int (G.nb_vertex graph) in
    let c =
      G.fold_vertex
        (fun v acc -> acc +. clustering_coefficient graph v)
        graph
        0.0
    in
    c /. n

  let averageShortestPathLength graph =
    let n = float_of_int (G.nb_vertex graph) in
    let c =
      G.fold_vertex (fun v acc -> acc +. average_distance graph v) graph 0.0
    in
    c /. n

  (* strongly directed components *)
  (* weakly directed compoenents = strongly directed compoenents if the graph
   * is not direct !!! *)
  let components graph =
    let module C = Components.Make (G) in
    C.scc_array graph

  let weaklycomponents graph =
    let module C = Components.Make (UndG) in
    C.scc_array (undirect graph)

  let numberComponents c = Array.length c

  let averageComponents c =
    let sum = Array.fold_left (fun acc i -> acc + List.length i) 0 c in
    float_of_int sum /. float_of_int (Array.length c)

  let largestComponent c =
    Array.sort (fun x y -> compare (List.length x) (List.length y)) c ;
    List.length c.(Array.length c - 1)

  let density graph =
    let n = float_of_int (G.nb_vertex graph) in
    let ps_edg = n *. (n -. 1.0) in
    float_of_int (G.nb_edges graph) /. ps_edg

  let averageTwoStepReach graph =
    let module S = Set.Make (struct
      type t = G.vertex

      let compare = compare
    end) in
    let n = float_of_int (G.nb_vertex graph) in
    let t =
      G.fold_vertex
        (fun i0 total ->
          let s =
            G.fold_succ
              (fun i1 set1 ->
                G.fold_succ (fun i2 set2 -> S.add i2 set2) graph i1 set1)
              graph
              i0
              S.empty
          in
          total +. float_of_int (S.cardinal s))
        graph
        0.0
    in
    t /. n

  let removezdp graph =
    let g = G.copy graph in
    G.iter_vertex
      (fun v ->
        if G.in_degree graph v = 0 && G.out_degree graph v = 0 then
          G.remove_vertex g v)
      graph ;
    g

  (*
  (* bullshit *)
  let brokerage graph =
    let n = float_of_int (G.nb_vertex graph) in
    let ps_edg = n *. (n -. 1.0) in
    let total = ref 0 in
    G.iter_vertex (fun i0 ->
      G.iter_vertex (fun i1 ->
        if not (G.mem_edge graph i0 i1) then
          incr total
      ) graph
    ) graph;
    ((float_of_int !total) /. ps_edg)

  let shorter_path_length gr v =
    let module Bfs = Graph.Traverse.Bfs(G) in
    let seen = Hashtbl.create 1031 in
    let level = ref 0 in
    Bfs.prefix_component (fun v ->
      incr level ;
      Hashtbl.add see v !level
    ) gr v
    seen
  ;;

  let eccentricity gr =
    let vv = Hashtbl.create 1031 in
    G.iter_vertex (fun v ->
      let h = shorted_path_length gr v in
      Hashtbl.add vv v (Hashtbl.fold (fun k v acc -> max v acc) h 0)

    ) gr
  *)
end
