(* 
%use "standard_prelude.fsh" ;;
*)

(* some basic functions *)

(* composition and identity *)

let comp g f x = g (f x) ;;
let identity x = x ;;


(* curried projections *)

let fst x y = x ;;
let snd x y = y ;;


(* incrementing and decrementing *)

let incr x = x := x+1 ;;
let decr x = x := x-1 ;;


(* Shape operations *)

(* 
check x y 
checks that the fact x is true 
and then evaluates y 
*)

(* well_shaped x checks that x is not error 

let well_shaped (x: exp a) = check (#x #= #x) true
;;

*)

(* 
folddim : (#a -> x) 
	  -> (size -> x -> x) 
	  -> #[a] -> x

The shape of an array is essentially 
a list of sizes, paired with a shape 
for the entries. Corresponding to
this is a foldright operation, called 
folddim. Perhaps it should be 
the primitive, instead of numdim.
*)

let folddim f0 f1 xsh = 
  let h x = f0 (undim x)
  and g n f x = 
	f1 (lendim x) (f (preddim x))
  in primrec g h (numdim xsh) xsh
;;

(* 
zeroShape : #[a] -> #a
extracts the shape of the entry 
of a zero dimensional array 
*)

let zeroShape = 
  folddim identity snd
;;

(*
extendShape : #[a] -> #b -> #[b]
uses the "outer" shape of its first argument
to augment the shape of its second argument.
For example, 
extendShape {2,3:int_shape} {4:bool_shape} =
{ 2,3 : 4 : bool_shape }
*)

let extendShape =
  folddim f0 f1
  where 
      f0 x y = zerodim y
  and f1 n h y = succdim n (h y) 
;;

(*
extendShape2 : #[a] -> #[b] -> #[b]
is like extendShape, 
but does not introduce a zerodim. For example,
extendShape2 {2,3:int_shape} {4:bool_shape} =
{ 2,3,4 : bool_shape }
*)

let extendShape2 = 
  folddim f0 f1
  where
      f0 x y = y
  and f1 n h y = succdim n (h y) 
;;


(* number_of_entries: #[a] -> size *)

let number_of_entries = 
  folddim (fun (x: exp #a) -> 1) (fun (x:size) -> fun (y:size) -> x* y) 
;;

(* shape_zipop : (#a -> #[b] -> #[c]) -> 
		 (size -> size -> size) -> 
		 #[a] -> #[b] -> #[c]

allows one to manipulate shapes as if they were lists  
*)


let shape_zipop f g = 
  folddim f f1 
where
  f1 n h b = succdim (g n (lendim b)) (h (preddim b)) 
;;

(*
snocdim : size -> #[a] -> #[a]
adds the new size as an innermost dimension.
For example,
snocdim 4 {2,3:int_shape} = { 2,3,4 : int_shape }
*)

let snocdim n sh = 
  extendShape2 sh (succdim n (zerodim (zeroShape sh)))
;;


(* 
shape2array : #[a] -> [int]
converts a shape to an array of integers 
representing sizes. For example,
shape2array {2,3:bool_shape} = 
  fill { 2 : int_shape } with [2,3]
*)

let shape2array sh = 
  new #x = succdim (numdim sh) (zerodim int_shape)
  and  i = 0 in 
    folddim f0 f1 sh
    where
      f0 sh0 = skip 
    and 
      f1 (m:size) C = 
	  (x[i] := m ;
	  incr i ;
	  C) 
  return x 
;;


(* Procedures *)



(* doall, doall2 and doall3 

doall allow procedures to be applied to 
all the entries in an array, 
(without using the index as a parameter)
like a multi-dimensional forall. 
doallk does the same for a procedure
taking k arguments. 
*)


let doall pr (x: var [a]) = 
  primrec g (do_nonce pr) (numdim #x) x
  where 
      g n = do_once
  and do_once pr1 (x1: var [a]) = 	// 1dim 
    for i< lendim #x1 do
      pr1 sub(x1,i) 
    done 
  and do_nonce pr0 (x0: var [a]) = pr0 (get x0)  	// 0dim 
;;

(*
idoall : (var [int] -> var a -> comm) -> var [a] -> comm

is an indexed version of doall, i.e. allows the procedure to access
the array index as first argument.  
*)

let idoall pr (x:var [a]) = 
let m = numdim #x in 
new #ndx = {m:int_shape} in 
let g n pr1 (x1:var[a]) = 		
      for (i< (lendim #x1)) do
        ndx[(m-n)-1] := i ;
        pr1 sub(x1,i) 
      done 
in
  primrec g (fun (x:var [a]) -> pr ndx (get x)) (numdim #x) x
end
;;

(* doall2 is a two array version of doall - indexing is done using the
second array argument *)

let doall2 pr (x:var[a]) (y:var[b]) = 
  primrec g (do_nonce2 pr) (numdim #y) x y
  where 
      g n = do_once2 
  and do_once2 pr1 (x1:var[a]) (y1:var[b]) = 
    for (i< (lendim #y1)) do
      pr1 sub(x1,i) sub(y1,i)
    done 
  and do_nonce2 pr0 (x0:var[a]) (y0:var[b]) = 
    pr0 (get x0) (get y0)
;;

let idoall2 pr (x:var[a]) (y:var[b]) = 
let m = numdim #x in 
new #ndx = {m:int_shape} in 
let do_nonce2 pr0 ndx (x0:var[a]) (y0:var[b]) =     
      pr0 ndx (get x0) (get y0) 
and g n pr1 (x1:var[a]) (y1:var[b]) = 		
      for (i< (lendim #x1)) do
        ndx[(m-n)-1] := i ;
        pr1 sub(x1,i) sub(y1,i)
      done 
in
  primrec g (do_nonce2 pr ndx) (numdim #x) x y
end
;;


let doall2_offset left pr (x:var[a]) (y:var[b]) = 
  primrec g (do_nonce2 pr) (numdim #y) left x y
  where 
      g n pr1 left1 (x1:var[a]) (y1:var[b])  = 
        for (i< (lendim #y1)) do
          pr1 (preddim left1) sub(x1,lendim left1 + i) sub(y1,i)
        done 
  and do_nonce2 pr0 left0 (x0:var[a]) (y0:var[b]) = 
    pr0 (get x0) (get y0)
;;


let doall3 pr (x:var[a]) (y:var[b]) (z: var[c]) = 
  primrec g (do_nonce3 pr) (numdim #x) x y z
  where 
      g n = do_once3 
  and do_once3 pr1 (x1:var[a]) (y1:var[b]) (z1: var[c]) =
    for (i< (lendim #z1)) do
      pr1 sub(x1,i) sub(y1,i) sub(z1,i) 
    done 
  and do_nonce3 pr0 (x0:var[a]) (y0:var[b]) (z0: var[c]) =
    pr0 (get x0) (get y0) (get z0)
;;


let doall3_offset (left1: var[int]) (left2:var[int]) 
		  pr 
		  (x:var[a]) (y:var[b]) (z: var[c]) =
  primrec g (do_nonce3 pr) (numdim #x) x y z
  where 
      g n = do_once3_offset 
		left1[(lendim #left1 - n - 1 )] 
		left2[(lendim #left2 - n - 1 )] 
  and do_once3_offset (offset1: int) (offset2: int) 
			pr1 (x1:var[a]) (y1:var[b]) (z1: var[c]) =
    for (i< (lendim #z1)) do
      pr1 sub(x1,offset1 + i) 
	  sub(y1,offset2 + i) 
   	  sub(z1,i) 
    done 
  and do_nonce3 pr0 (x0:var[a]) (y0:var[b]) (z0: var[c]) =
    pr0 (get x0) (get y0) (get z0)
;;


(* functions from procedures and shapes *)

let ap2exp f x = new y = x in f y end 
;;


(*
proc2fun : (var a -> var b -> comm) 
		-> (#b -> #a) 
		-> b -> a

Note that in f = proc2fun f_pr f_sh 
the first argument to f_pr is a variable 
whose data type is that output of f.
Also, f_pr should *not* write to 
its second argument.
*)

let proc2fun f_pr f_sh x = 
  new #y = f_sh #x
  and z = x in 
    f_pr y z
  return y
;;

(* 
proc2fun2 : (var a -> var b -> var c -> comm) 
		-> (#b -> #c -> #a) 
		-> b -> c -> a

proc2fun2 is like proc2fun 
but produces a function of two arguments. 
*)

let proc2fun2 f_pr f_sh x y =
  new #z = f_sh #x #y 
  and x0 = x
  and y0 = y in  
  f_pr z x0 y0
  return z
;;

(* 
ap2er : (var a -> comm) -> a -> a
e2ap : (a -> a) -> var a -> comm

Variations on ap2exp 
*)

let ap2er proc arg = 
	new aux = arg 
	in proc aux 
	return aux ;;


let e2ap f (x:var a) = x:= f x ;;



(* Second-Order Functions *)


(* Mapping *)

(* the procedure *)
let map_pr f  = 
  let pr (y:var b) (x:var a) = y := f !x in 
  doall2 pr
;;

(* in-place update *)
let update f x = map_pr f x x ;;


(* the shape *)
let map_sh fsh =
  folddim (zerodim . fsh) succdim
;;

(* the function 
map : (a -> b) -> [a] -> [b]
*)
let map f = proc2fun (map_pr f) (map_sh #f) 
;;


(* selfmap : (a -> a) -> [a] -> [a]
will use update instead of map_pr if the shapes allow it. *)

let selfmap f x = 
  let sh = zeroShape #x in 
  if 	#f sh #= sh 
  then 	new y = x in 
	  update f y 
	return y
  else map f x
;;


(*
imap : ([int] -> a -> b) -> [a] -> [b]

is an indexed map (see idoall above)
*)

let imap_pr f  = 
   let pr (ndx:var [int]) y x = 
	y := f !ndx !x in 
  idoall2 pr
;;

let imap f x = proc2fun (imap_pr f) (map_sh (#f {numdim #x : int_shape})) x
;;

(*
map_pr_offset : var [int] -> (a -> b) -> var [b] -> var [a] -> comm
*)

let map_pr_offset left f =
  let pr y (x:var a) = y := f !x in 
  doall2_offset left pr
;;

(* 
map_offset : var [int] -> (a -> b) -> var [b] -> [a] -> comm
  *)

let map_offset left f y = 
  ap2exp (map_pr_offset left f y) 
;;

(* 
zipop : (a -> b -> c) 
	-> [a] -> [b] -> [c]
takes a function of two arguments, 
and two arrays of the same outer shape, 
and produces a third array 
of the same outer shape, 
whose entries are obtained by applying 
the function to corresponding entries.

zip is not defined as the language 
does not currently support pairing.
*)

let zipop_pr f =

  let pr z (x:var a) (y:var b) = z := f !x !y
  in
  doall3 pr
;;

let zipop_sh fsh  =
  let f0 x y = 
    zerodim (fsh x (undim y)) 
  and f1 (n:size) h y = 
    check (n = (lendim y))
    (succdim n (h (preddim y)))
  in 
    folddim f0 f1 
;;

let zipop f = 
  proc2fun2 (zipop_pr f) (zipop_sh #f) 
;;


(* reduction 

Reduction of a function f is like folding 
of f except that f is required to have 
the property that #(f x y) = #x.
This is true for datum operations, 
but also for many array operations,
such as pointwise multiplication 
occuring in inner product.
*)

(* accum : (a -> b -> a) 
	   -> var a -> var [b] -> comm
reduces in place, 
c.f. update above 
*)
let accum f z = 
  let pr y = z:= f !z !y 
  in 
    doall pr
;;

let reduce_pr f x z y = 
  (z := x ;
  accum f z y)
;;

let reduce_sh fsh (xsh: exp #a) ysh =
  check (fsh xsh (zeroShape ysh) #= xsh )
		xsh 
;;		

let reduce f x = 
  proc2fun (reduce_pr f x) (reduce_sh #f #x) 	
;;


(* 
fold : (x -> a -> x) 
	-> x -> [a] -> x
can be applied to operations like append
which produce values having new shapes. 
These are computed by unwinding 
a primitive recursion. 
*)

let fold f x y  =
  if #f #x (zeroShape #y) #= #x 
  then reduce f x y 
  else 
  primrec g (fold_nonce f) (numdim #y) x y
  where 
      g n = fold_once 
  and fold_once h x y = 
    primrec (foldsucc h y) x (lendim #y) 
  and foldsucc h (y:exp b) i x = h x sub(y,i) 
  and fold_nonce f x (y:exp b) = f x (get y)
;;


(* Basic Datum Functions *)

(* integers *)

let zero_int = (0:int) ;;
let plus_int (x:int) y = x + y ;;	
let subtract_int (x:int) y = x-y;;
let times_int (x:int) y = x * y ;;
let divide_int (x:int) y = x div y ;;
let equal_int (x:int) y = x = y ;;
let less_than_int (x:int) y = x< y ;;
let less_than_or_equal_int (x:int) y = x<= y ;;
let sum_int = reduce plus_int 0 ;;
let inner_product_int x y = sum_int (zipop times_int x y) ;;

let modulus (x:int) y = x - (x div y * y) ;;

(* booleans *)

let conjunction (x:bool) y = x && y ;;
let disjunction (x:bool) y = x || y ;;
let eq_bool (x:bool) y = x=y ;;    

let bool2int (x:bool) = if x then 1 else 0 ;;
let all = reduce conjunction true ;;


(* floats *)

let zero_float = 0.0 ;;
let plus_float x y = x +. y ;;
let subtract_float x y = x -. y;;
let times_float x y = x *. y ;;
let less_than_float x y = x<. y ;;
let less_than_or_equal_float x y = x<=. y ;;

let sum_float = reduce plus_float zero_float ;;
let inner_product_float x y = sum_float (zipop times_float x y) ;; 



(* Initialising arrays *)

(* copy : #[a] -> b -> [b] *)

let copy_pr (y:var [a]) (x: var a) = 
  let pr z = z := x in 
   doall pr y
;;

let copy sh = proc2fun copy_pr (extendShape sh)
;;



(* Diagonal indexing *)

let diagvector k = 
	new #x = { k: int_shape} in 
	for (i < k) do
	  x[i] := i 
        done
	return x
;;

(* reversing a vector *)

let reverse_pr (y:var b) (x:var a) = 
  for (i < n ) do
      sub(y,i) := sub(x,n-1-i)
    done
  where n = (lendim #!x)
;;

let reverse = proc2fun reverse_pr identity
;;


(* 
indexFill : #[a] -> #[int] -> [[int]]
fills an array shape with arrays 
containing the index information at
each entry.

No longer used. Better to use imap. 
 

let indexFill_pr sh = 
  primrec indexFill_step indexFill_pr0 m
  where
    indexFill_step n pr indices iA = 
	for (i<(lendim #!iA)) do
	indices[m - n - 1] := i ;	(* set index entry to be i *)
	pr indices sub(iA,i) 		(* set the other indices *)
      done 
  and indexFill_pr0 indices iA = 
      get iA := indices 		(* set iA entry to indices *)
  and m = numdim sh ;;

let indexFill sh sh1 =
  new #iA = extendShape sh sh1  
  and #indices = sh1 in 
  indexFill_pr sh indices iA
  return iA
;;

*)


(* No longer used. Try imap 

let indexArray sh = 	
	(* puts index array at each entry, 
	   e.g. the (i,j) position in a matrix has array [i,j] *)
  indexFill sh {(numdim sh) : int_shape}
;;

*)

(* Array entry techniques *)

let multi_sub (indices: exp [int]) array =
  let g n (x:var [a]) = sub(x,indices[n]) in 
  primrec g array (lendim #indices)
;;

let entry indices array = 
  get (multi_sub indices array)
;;

let entry2 array indices = 
  entry indices array 
;;


let multi_subexp_pr y indices x = 
    y := multi_sub indices x 
;;

let multi_subexp_sh ndx_sh array_sh = 
    primrec (fun x -> preddim) 
	array_sh (lendim ndx_sh) 
;;

let multi_subexp = proc2fun2 multi_subexp_pr multi_subexp_sh  
;;

(* 
entryexp : [int] -> [a] -> a
*)

let entryexp_pr y indices array = 
  y:= entry indices array 
;;

let entryexp_sh x = zeroShape 
;;

let entryexp = 
  proc2fun2 entryexp_pr  entryexp_sh
;;


let entryexp2 array indices = 
  entryexp indices array 
;;


(* 
transpose : [a] -> [a]
moves entries to the position 
obtained by reversing their indices.
*)

let transpose_pr y x = 
  let pr xs ndx y1 = 
    new xdn = reverse !ndx in 
	y1 := entry xdn xs 
    end
  in 
    idoall (pr x) y
;;


let transpose_sh = 
    folddim zerodim snocdim 
;;

let transpose = proc2fun transpose_pr transpose_sh 
;;