(* ocamlc -I +ocamlmpi mpi.cma str.cma hopfield_measure2.ml -o hopfield_measure2ml
   mpirun -np 2 hopfield_measure2ml N M filename *)
open Mpi
open Printf

let acoef = 1.0
let bcoef = 1.0
let dcoef = 2.0

let set_pos num pos file = 
  let input = open_in file in
    for i=1 to num do
      let strlist = Str.split_delim (Str.regexp ",") (input_line input) in
      let a = (float_of_string (List.nth strlist 0)) in
      let b = (float_of_string (List.nth strlist 1)) in
	pos.(i-1) <- (a,b)
    done;
    close_in input

let set_distance num pos distance = 
  let distance_total = ref 0.0 in
    for i=1 to num do
      for j=1 to num do
	if i=j then
	  distance.(i-1).(j-1) <- 0.0
	else(
	  let a_x = fst pos.(i-1) and a_y = snd pos.(i-1) and b_x = fst pos.(j-1) and b_y = snd pos.(j-1) in
	  let c = (a_x -. b_x) *. (a_x -. b_x) and d = (a_y -. b_y) *. (a_y -. b_y) in
	  let e = sqrt (c +. d) in
	    distance.(i-1).(j-1) <- e;
	    distance_total := !distance_total +. distance.(i-1).(j-1));
      done;
    done;
    for i=0 to (num-1) do
      for j=0 to (num-1) do
	distance.(i).(j) <- 10.0 *. distance.(i).(j) /. !distance_total
      done;
    done

      

let size = comm_size comm_world
let myrank = comm_rank comm_world

let set_fst_index num =
  let a = ref 0 in
    for j=0 to (myrank-1) do
      a := !a + ((num+j)/size)
    done;
    !a

let set_snd_index num =
  let a = ref 0 in
    for j=0 to myrank do
      a := !a + ((num+j)/size)
    done;
    a := !a - 1;
    !a


let energy num unit distance =
  let term1 = ref 0.0 and term2 = ref 0.0 and term3 = ref 0.0 in
  let jm = ref 0 and jp = ref 0 in
    for i=0 to (num-1) do 
      for j=0 to (num-1) do
	for m=0 to (num-1) do
	  if m!=j then term1 := !term1 +. unit.(i).(j) *. unit.(i).(m)
	done;
      done;
    done;
    term1 := 0.5 *. acoef *. !term1;

    for j=0 to (num-1) do
      for i=0 to (num-1) do
	for m=0 to (num-1) do
	  if m!=i then term2 := !term2 +. unit.(i).(j) *. unit.(m).(j)
	done;
      done;
    done;
    term2 := 0.5 *. bcoef *. !term2;

    for i=0 to (num-1) do
      for j=0 to (num-1) do
	if (j-1) = -1 then jm := num-1 else jm := j-1;
	if (j+1) = num then jp := 0 else jp := j+1;
	for m=0 to (num-1) do
	  term3 := !term3 +. distance.(i).(m) *. unit.(i).(j) *. (unit.(m).(!jp) +. unit.(m).(!jm));
	done;
      done;
    done;
    term3 := 0.5 *. dcoef *. !term3;
    
    !term1 +. !term2 +. !term3

let energy2 fst_index snd_index num unit distance =
  let term1 = ref 0.0 and term2 = ref 0.0 and term3 = ref 0.0 in
  let term4 = ref 0.0 and term5 = ref 0.0 and term6 = ref 0.0 in
  let jm = ref 0 and jp = ref 0 in
    for i=fst_index to snd_index do 
      for j=0 to (num-1) do
	for m=0 to (num-1) do
	  if m!=j then term1 := !term1 +. unit.(i).(j) *. unit.(i).(m)
	done;
      done;
    done;
    term4 := allreduce_float !term1 Float_sum comm_world;
    term4 := 0.5 *. acoef *. !term4;

    for j=fst_index to snd_index do
      for i=0 to (num-1) do
	for m=0 to (num-1) do
	  if m!=i then term2 := !term2 +. unit.(i).(j) *. unit.(m).(j)
	done;
      done;
    done;
    term5 := allreduce_float !term2 Float_sum comm_world;
    term5 := 0.5 *. bcoef *. !term5;

    for i=fst_index to snd_index do
      for j=0 to (num-1) do
	if (j-1) = -1 then jm := num-1 else jm := j-1;
	if (j+1) = num then jp := 0 else jp := j+1;
	for m=0 to (num-1) do
	  term3 := !term3 +. distance.(i).(m) *. unit.(i).(j) *. (unit.(m).(!jp) +. unit.(m).(!jm));
	done;
      done;
    done;
    term6 := allreduce_float !term3 Float_sum comm_world;
    term6 := 0.5 *. dcoef *. !term6;
    
    !term4 +. !term5 +. !term6




let set_unit num unit =
  if myrank = 0 then
    (for i=0 to (num-1) do
       for j=0 to (num-1) do
	 unit.(i).(j) <- Random.float 1.0
       done;
     done);
  let a = broadcast unit 0 comm_world in
    if myrank!=0 then 
      for i=0 to (num-1) do
	for j=0 to (num-1) do
	  unit.(i).(j) <- a.(i).(j)
	done;
      done

let _ =
  let num_tmp = if myrank=0 then (int_of_string Sys.argv.(1)) else 1 in
  let fornum_tmp = if myrank=0 then (int_of_string Sys.argv.(2)) else 1 in
  let num = broadcast num_tmp 0 comm_world in
  let fornum = broadcast fornum_tmp 0 comm_world in
  let file = broadcast Sys.argv.(3) 0 comm_world in
  let elapsed_time = ref 0.0 in
  let pos = Array.make num (0.0,0.0) in
  let unit = Array.make_matrix num num 0.0 in
  let distance = Array.make_matrix num num 0.0 in
  let fst_index = set_fst_index num and snd_index = set_snd_index num in
    set_pos num pos file;
    set_distance num pos distance;
    set_unit num unit;

    barrier comm_world;
    elapsed_time := Sys.time();
    for i=1 to fornum do
      (* if myrank=0 then printf "mpi energy:%f, single energy%f\n" (energy2 fst_index snd_index num unit distance) 
	 (energy num unit distance)  *)
      ignore(energy2 fst_index snd_index num unit distance);
    done;
    barrier comm_world;
    if myrank=0 then printf "elapsed time:%f second\n" (Sys.time() -. !elapsed_time)