(** Module for handling distributed parallelism *)

let size = 1

let rank = 0

let master = true


let barrier () = ()

let broadcast x = 
    Lazy.force x

let broadcast_int x = x

let broadcast_int_array x = x

let broadcast_float x = x

let broadcast_float_array x = x

let broadcast_vec x = x


module Vec = struct

  type t = 
  {
    global_first : int ; (* Lower  index in the global array *)
    global_last  : int ; (* Higher index in the global array *)
    local_first  : int ; (* Lower  index in the local array *)
    local_last   : int ; (* Higher index in the local array *)
    data         : Lacaml.D.vec ; (* Lacaml vector containing the data *)
  }

  let dim vec =
    vec.global_last - vec.global_first + 1

  let local_first  vec = vec.local_first
  let local_last   vec = vec.local_last 
  let global_first vec = vec.global_first
  let global_last  vec = vec.global_last 
  let data         vec = vec.data

  let pp ppf v =
    Format.fprintf ppf "@[<2>";
    Format.fprintf ppf "@[ gf : %d@]@;" v.global_first;
    Format.fprintf ppf "@[ gl : %d@]@;" v.global_last;
    Format.fprintf ppf "@[ lf : %d@]@;" v.local_first;
    Format.fprintf ppf "@[ ll : %d@]@;" v.local_last;
    Format.fprintf ppf "@[ data : %a@]@;" (Lacaml.Io.pp_lfvec ()) v.data;
    Format.fprintf ppf "@]@."; 
    ()

  let create n = 
    {
      global_first = 1 ;
      global_last  = n ;
      local_first  = 1 ;
      local_last   = n ;
      data         = Lacaml.D.Vec.create n
    }


  let make n x = 
    let result = create n in
    { result with data =
      Lacaml.D.Vec.make
        (Lacaml.D.Vec.dim result.data)
        x
    }


  let make0 n = 
    make n 0.


  let init n f =
    let result = create n in
    { result with data =
        Lacaml.D.Vec.init
          (Lacaml.D.Vec.dim result.data)
          (fun i -> f (i+result.local_first-1))
    }


  let of_array a =
    let length_a = Array.length a in
    let a = 
      let n = length_a mod size in
      if n > 0 then
        Array.concat [ a ; Array.make (size-n) 0. ]
      else
        a
    in
    let result = create length_a in
    let a_local = Array.make (Array.length a) 0. in
    { result with data = Lacaml.D.Vec.of_array a_local }


  let to_array vec = 
    Lacaml.D.Vec.to_array vec.data 
    |> Array.copy

      
  let of_vec a =
    Lacaml.D.Vec.to_array a
    |> of_array


  let to_vec v =
    to_array v
    |> Lacaml.D.Vec.of_array
    

end



let dot v1 v2 =
  if Vec.dim v1 <> Vec.dim v2 then
    invalid_arg "Incompatible dimensions";
  Lacaml.D.dot (Vec.data v1) (Vec.data v2)