let pi : float = 3.141592653589793 in
let days_per_year : float = 365.24 in
let solar_mass : float = ((4.0 *. pi) *. pi) in
let dt : float = 0.01 in

let make_body : float -> float -> float -> float -> float -> float -> float -> vec[float] =
	fun (x: float) (y : float) (z : float) (vx : float) (vy : float) (vz : float) (mass : float) ->
	let v : vec[float] = vector 7 0.0 in
	vector_set v 0 x;
	vector_set v 1 y;
	vector_set v 2 z;
	vector_set v 3 vx;
	vector_set v 4 vy;
	vector_set v 5 vz;
	vector_set v 6 mass;
	v
in

let set_body_x : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 0 v;
    0
in

let body_x : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 0
in

let set_body_y : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 1 v;
    0
in

let body_y : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 1
in

let set_body_z : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 2 v;
    0
in

let body_z : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 2
in

let set_body_vx : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 3 v;
    0
in

let body_vx : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 3
in

let set_body_vy : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 4 v;
    0
in

let set_vy : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 4
in

let set_body_vz : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 5 v;
    0
in

let body_vz : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 5
in

let set_body_mass : vec[float] -> float -> int =
	fun (b : vec[float]) (v : float) ->
		vector_set b 6 v;
    0
in

let body_mass : vec[float] -> float =
	fun (b : vec[float]) ->
		vector_get b 6
in

let sun : vec[float] =
  make_body 0.0 0.0 0.0 0.0 0.0 0.0 solar_mass

in

let jupiter : vec[float] =
  make_body 4.84143144246472090
            -1.16032004402742839
            -1.03622044471123109e-1
            (1.66007664274403694e-3 *. days_per_year)
            (7.69901118419740425e-3 *. days_per_year)
            (-6.90460016972063023e-5 *. days_per_year)
            (9.54791938424326609e-4 *. solar_mass)
in

let saturn : vec[float] =
  make_body 8.34336671824457987
            4.12479856412430479
            -4.03523417114321381e-1
            (-2.76742510726862411e-3 *. days_per_year)
            (4.99852801234917238e-3 *. days_per_year)
            (2.30417297573763929e-5 *. days_per_year)
            (2.85885980666130812e-4 *. solar_mass)
in

let uranus : vec[float] =
  make_body 1.28943695621391310e1
            -1.51111514016986312e1
            -2.23307578892655734e-1
            (2.96460137564761618e-03 *. days_per_year)
            (2.37847173959480950e-03 *. days_per_year)
            (-2.96589568540237556e-05 *. days_per_year)
            (4.36624404335156298e-05 *. solar_mass)
in

let neptune : vec[float] =
  make_body 1.53796971148509165e+01
            -2.59193146099879641e+01
            1.79258772950371181e-01
            (2.68067772490389322e-03 *. days_per_year)
            (1.62824170038242295e-03 *. days_per_year)
            (-9.51592254519715870e-05 *. days_per_year)
            (5.15138902046611451e-05 *. solar_mass)
in

let system : vec[vec[float]] =
  let v : vec[vec[float]] = vector 5 (vector 0 0.0) in
  vector_set v 0 sun;
  vector_set v 1 jupiter;
  vector_set v 2 saturn;
  vector_set v 3 uranus;
  vector_set v 4 neptune;
  v
in

let system_size : int = 5 in

let rec offset_momentum_loop : int -> float -> float -> float -> int =
  fun (i1 : int) (px : float) (py : float) (pz : float) ->
    (if (i1 = system_size) then
      vector_set (vector_get system 0) 3 ((0.0 -. px) /. solar_mass);
      vector_set (vector_get system 0) 4 ((0.0 -. py) /. solar_mass);
      vector_set (vector_get system 0) 5 ((0.0 -. pz) /. solar_mass);
      0
    else
      let j : vec[float] = vector_get system i1 in
      offset_momentum_loop (i1 + 1)
                           (px +. ((vector_get j 3) *. (vector_get j 6)))
                           (py +. ((vector_get j 4) *. (vector_get j 6)))
                           (pz +. ((vector_get j 5) *. (vector_get j 6)));
      0);
    0
in

let offset_momentum : int -> int = 
  fun (_ : int) -> 
    offset_momentum_loop 0 0.0 0.0 0.0;
    0
in

let rec energy_loop_o : int -> float -> float =
  fun (o : int) (e : float) ->
    if (o = system_size) then
      e
    else
      let o1 : vec[float] = vector_get system o in
      let sqs : float = ((vector_get o1 3) *. (vector_get o1 3)) +.
                        ((vector_get o1 4) *. (vector_get o1 4)) +.
                        ((vector_get o1 5) *. (vector_get o1 5)) in
      let e : float = e +. ((0.5 *. (vector_get o1 6)) *. sqs) in
        energy_loop_i o o1 (o + 1) e
  and
  energy_loop_i : int -> vec[float] -> int -> float -> float =
    fun (o : int) (o1 : vec[float]) (i : int) (e : float) ->
      if (i = system_size) then
        energy_loop_o (o + 1) e
      else
        let i1 : vec[float] = vector_get system i in
        let dx : float = (vector_get o1 0) -. (vector_get i1 0) in
        let dy : float = (vector_get o1 1) -. (vector_get i1 1) in
        let dz : float = (vector_get o1 2) -. (vector_get i1 2) in
        let dist : float = sqrt ((dx *. dx) +. (dy *. dy) +. (dz *. dz)) in
        let e : float = e -. (((vector_get o1 6) *. (vector_get i1 6)) /. dist) in
        energy_loop_i o o1 (i + 1) e

in

let energy : int -> float = fun (_ : int) -> energy_loop_o 0 0.0 in

let rec advance_loop_o : int -> int =
  fun (o : int) ->
    (if (o = system_size) then
      0
    else
      let o1 : vec[float] = vector_get system o in
      advance_loop_i (o + 1)
                     (vector_get o1 3)
                     (vector_get o1 4)
                     (vector_get o1 5)
                     o1;
      advance_loop_o (o + 1));
      0
and
  advance_loop_i : int -> float -> float -> float -> vec[float] -> int =
    fun (i3 : int) (vx : float) (vy : float) (vz : float) (o1 : vec[float]) ->
      (if (i3 < system_size) then
        let i1 : vec[float] = vector_get system i3 in
        let dx : float = (vector_get o1 0) -. (vector_get i1 0) in
        let dy : float = (vector_get o1 1) -. (vector_get i1 1) in
        let dz : float = (vector_get o1 2) -. (vector_get i1 2) in
        let dist2 : float = (dx *. dx) +. (dy *. dy) +. (dz *. dz) in
        let mag : float = dt /. (dist2 *. sqrt (dist2)) in
        let dxmag : float = dx *. mag in
        let dymag : float = dy *. mag in
        let dzmag : float = dz *. mag in
        let om : float = vector_get o1 6 in
        let im : float = vector_get i1 6 in
        vector_set i1 3 ((vector_get i1 3) +. (dxmag *. om));
        vector_set i1 4 ((vector_get i1 4) +. (dymag *. om));
        vector_set i1 5 ((vector_get i1 5) +. (dzmag *. om));
        advance_loop_i (i3 + 1)
                       (vx -. (dxmag *. im))
                       (vy -. (dymag *. im))
                       (vz -. (dzmag *. im))
                       o1;
        0
      else
        vector_set o1 3 vx;
        vector_set o1 4 vy;
        vector_set o1 5 vz;
        vector_set o1 0 ((vector_get o1 0) +. (dt *. vx));
        vector_set o1 1 ((vector_get o1 1) +. (dt *. vy));
        vector_set o1 2 ((vector_get o1 2) +. (dt *. vz));
        0);
        0
in

let advance : int -> int =
  fun (_ : int) -> 
    advance_loop_o (0);
    0
in

let run_benchmark : int -> int =
  fun (_ : int) ->
    offset_momentum (0);
    print_float (energy (0));
    (let final_i = read_int () in
    for i = 0 to final_i do
      advance 0
    done);
    print_float (energy (0));
    0
in

time run_benchmark
