(* probe.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Expansion of probe operations in the HighIL to MidIL translation.
 *)

structure Probe : sig

    val expand : {
	    result : MidIL.var,		(* result variable for probe *)
	    img : MidIL.var,		(* probe image argument *)
	    v : ImageInfo.info,		(* summary info about image *)
	    h : Kernel.kernel,		(* reconstruction kernel *)
	    k : int,			(* number of levels of differentiation *)
	    pos : MidIL.var		(* probe position argument *)
	  } -> MidIL.assign list

  end = struct

    structure SrcIL = HighIL
    structure SrcOp = HighOps
    structure DstIL = MidIL
    structure DstTy = MidILTypes
    structure DstOp = MidOps
    structure DstV = DstIL.Var
    structure VMap = SrcIL.Var.Map
    structure IT = Shape

  (* generate a new variable indexed by dimension *)
    fun newVar_dim (prefix, d, ty) =
	  DstV.new (prefix ^ Partials.axisToString(Partials.axis d), ty)

    fun assign (x, rator, args) = (x, DstIL.OP(rator, args))
    fun cons (x, args) = (x, DstIL.CONS(DstV.ty x, args))
    fun realLit (x, i) = (x, DstIL.LIT(Literal.Float(FloatLit.fromInt i)))
    fun intLit (x, i) = (x, DstIL.LIT(Literal.Int(IntInf.fromInt i)))

  (* generate code for a evaluating a single element of a probe operation *)
    fun probeElem {
	    dim,	(* dimension of space *)
	    h, s,	(* kernel h with support s *)
	    n, f,	(* Dst vars for integer and fractional components of position *)
	    voxIter	(* iterator over voxels *)
	  } (result, pdOp) = let
	  val vecsTy = DstTy.vecTy(2*s) (* vectors of coefficients cover support of kernel *)
	  val vecDimTy = DstTy.vecTy dim
	(* generate the variables that hold the convolution coefficients.  The
	 * resulting list is in slowest-to-fastest axes order.
	 *)
	  val convCoeffs = let
		val Partials.D l = pdOp
		fun mkVar (_, [], coeffs) = coeffs
		  | mkVar (i, d::dd, coeffs) = (case d
		       of 0 => mkVar(i+1, dd, newVar_dim("h", i, vecsTy) :: coeffs)
			| 1 => mkVar(i+1, dd, newVar_dim("dh", i, vecsTy) :: coeffs)
			| _ => mkVar(i+1, dd, newVar_dim(concat["d", Int.toString d, "h"], i, vecsTy) :: coeffs)
		      (* end case *))
		in
		  mkVar (0, l, [])
		end
	(* for each dimension in space, we evaluate the kernel at the coordinates for that axis.
	 * the coefficients are
	 *	h_{s-i} (f - i) for 1-s <= i <= s
	 *)
	  val coeffCode = let
		fun gen (x, k, (d, code)) = let
		    (* note that for 1D images, the f vector is a scalar *)
		      val fd = if (dim > 1)
			    then newVar_dim ("f", d, DstTy.realTy)
			    else f
		      val a = DstV.new ("a", vecsTy)
		    (* note that we reverse the order of the list since the convolution
		     * space is flipped from the image space and we want the voxel vector
		     * to be in increasing address order.
		     *)
		      val tmps = List.rev(List.tabulate(2*s,
			    fn i => (DstV.new("t"^Int.toString i, DstTy.realTy), i - s)))
		      fun mkArg ((t, 0), code) = (t, DstIL.VAR fd) :: code
			| mkArg ((t, n), code) = let
			    val (rator, n) = if (n < 0) then (DstOp.Sub, ~n) else (DstOp.Add, n)
			    val t' = DstV.new ("r", DstTy.realTy)
			    in
			      realLit (t', n) ::
			      assign (t, rator DstTy.realTy, [fd, t']) ::
			      code
			    end
		      val code =
			    cons(a, List.map #1 tmps) ::		    
			    assign(x, DstOp.EvalKernel(2*s, h, k), [a]) ::
			      code
		      val code = List.foldr mkArg code tmps
		      val code = if (dim > 1)
			    then assign(fd, DstOp.Select(DstTy.vecTy dim, d), [f]) :: code
			    else code
		      in
			(d+1, code)
		      end
		val Partials.D l = pdOp
		in
		(* we iterate from fastest to slowest axis *)
		  #2 (ListPair.foldr gen (0, []) (convCoeffs, List.rev l))
		end
	(* generate the reduction code in reverse order *)
	  fun genReduce (result, [hh], IT.LF{vox, offsets}, code) =
		assign (result, DstOp.Dot(2*s), [vox, hh]) :: code
	    | genReduce (result, hh::r, IT.ND(_, kids), code) = let
		val tv = DstV.new ("tv", vecsTy)
		val tmps = List.tabulate(2*s, fn i => DstV.new("t"^Int.toString i, DstTy.realTy))
		fun lp ([], [], code) = code
		  | lp (t::ts, kid::kids, code) = genReduce(t, r, kid, lp(ts, kids, code))
		val code = cons(tv, tmps) :: assign(result, DstOp.Dot(2*s), [hh, tv]) :: code
		in
		  lp (tmps, kids, code)
		end
	    | genReduce _ = raise Fail "genReduce"
	  val reduceCode = genReduce (result, convCoeffs, voxIter, [])
	  in
	    coeffCode @ reduceCode
	  end

    fun doVoxelSample (result, v, k, s, diffIter, {h, n, f, img}, offset) = let
	  val stride = ImageInfo.stride
	  val dim = ImageInfo.dim v
	  val vecsTy = DstTy.vecTy(2*s) (* vectors of coefficients cover support of kernel *)
	(* generate code to load the voxel data; since we use a vector load operation to load the
	 * fastest dimension, the height of the tree is one less than the dimension of space.
	 *)
	  val voxIter = let
		fun f (i, (offsets, id)) = (i - (s - 1) :: offsets, i::id)
		fun g (offsets, id) = {
			offsets = ~(s-1) :: offsets,
			vox = DstV.new(String.concat("v" :: List.map Int.toString id), vecsTy)
		      }
		in
		  IT.create (dim-1, 2*s, fn _ => (), f, g, ([], []))
		end
	  val loadCode = let
		fun genCode ({offsets, vox}, code) = let
		      fun computeIndices (_, []) = ([], [])
			| computeIndices (i, offset::offsets) = let
			    val index = newVar_dim("i", i, DstTy.intTy)
			    val t1 = DstV.new ("t1", DstTy.intTy)
			    val t2 = DstV.new ("t2", DstTy.intTy)
			    val (indices, code) = computeIndices (i+1, offsets)
			    val code = if (dim > 1)
				  then
				    intLit(t1, offset) ::
				    assign(t2, DstOp.Select(DstTy.IVecTy dim, i), [n]) ::
				    assign(index, DstOp.Add(DstTy.intTy), [t1, t2]) ::
				    code
				  else
				    intLit(t1, offset) ::
				    assign(index, DstOp.Add(DstTy.intTy), [t1, n]) ::
				    code
			    val indices = index::indices
			    in
			      (indices, code)
			    end
		      val (indices, indicesCode) = computeIndices (0, offsets)
		      val a = DstV.new ("a", DstTy.AddrTy v)
		      in
			indicesCode @ [
			    assign(a, DstOp.VoxelAddress(v, offset), img::indices),
			    assign(vox, DstOp.LoadVoxels(v, 2*s), [a])
			  ] @ code
		      end
		in
		  IT.foldr genCode [] voxIter
		end
	(* generate code to evaluate and construct the result tensor *)
	  val probeElem = probeElem {dim = dim, h = h, s = s, n = n, f = f, voxIter = voxIter}
	  fun genProbe (result, IT.ND(_, kids as (IT.LF _)::_), code) = let
	      (* the kids will all be leaves *)
		fun genProbeCode (IT.LF arg, code) = probeElem arg @ code
		fun getProbeVar (IT.LF(t, _)) = t
		in
		  List.foldr genProbeCode (cons (result, List.map getProbeVar kids) :: code) kids
		end
	    | genProbe (result, IT.ND(ty, kids), code) = let
(* FIXME: the type of the tmps depends on the types of the kids *)
		val tmps = List.tabulate(dim, fn i => DstV.new("t"^Int.toString i, ty))
		val code = cons(result, tmps) :: code
		fun lp ([], [], code) = code
		  | lp (t::ts, kid::kids, code) = genProbe(t, kid, lp(ts, kids, code))
		in
		  lp (tmps, kids, code)
		end
	    | genProbe (result, IT.LF(t, pdOp), code) = (* for scalar fields *)
		probeElem (result, pdOp) @ code
	  val probeCode = if (k > 0)
		then let
		(* for gradients, etc. we have to transform back to world space *)
		  val ty = DstV.ty result
		  val tensor = DstV.new("tensor", ty)
		  val xform = assign(result, DstOp.TensorToWorldSpace(v, ty), [img, tensor])
		  in
		    genProbe (tensor, diffIter, [xform])
		  end
		else genProbe (result, diffIter, [])
	  in
(* FIXME: for dim > 1 and k > 1, we need to transform the result back into world space *)
	    loadCode @ probeCode
	  end

  (* generate code for probing the field (D^k (v * h)) at pos *)
    fun expand {result, img, v, h, k, pos} = let
	  val dim = ImageInfo.dim v
	  val s = Kernel.support h
	  val vecsTy = DstTy.vecTy(2*s) (* vectors of coefficients to cover support of kernel *)
	  val vecDimTy = DstTy.vecTy dim
	(* generate the transform code *)
	  val x = DstV.new ("x", vecDimTy)	(* image-space position *)
	  val f = DstV.new ("f", vecDimTy)
	  val nd = DstV.new ("nd", vecDimTy)
	  val n = DstV.new ("n", DstTy.IVecTy dim)
	  val toImgSpaceCode = [
		  assign(x, DstOp.PosToImgSpace v, [img, pos]),
		  assign(nd, DstOp.Floor dim, [x]),
		  assign(f, DstOp.Sub vecDimTy, [x, nd]),
		  assign(n, DstOp.RealToInt dim, [nd])
		]
	(* generate the shape of the differentiation tensor with variables representing
	 * the elements
	 *)
	  fun diffIter () = let
		val partial = Partials.partial dim
		fun f (i, (_::dd, axes)) = (dd, Partials.axis i :: axes)
		fun labelNd (_::dd, _) = DstTy.tensorTy dd
		fun labelLf (_, axes) = let
			val r = DstV.new(
				String.concat("r" :: List.map Partials.axisToString axes),
				DstTy.realTy)
			in
			  (r, partial axes)
			end
 		in
		  IT.create (k, dim, labelNd, f, labelLf, (List.tabulate(k, fn _ => dim), []))
		end
	  val vars = {h=h, n=n, f=f, img=img}
	  in
	    case ImageInfo.voxelShape v
	     of [] => toImgSpaceCode @ doVoxelSample (result, v, k, s, diffIter(), vars, 0)
	      | [d] => let
                  val sampleTy = DstTy.TensorTy(List.tabulate(k, fn _ => d))
		  fun doSamples (offset, xs, code) = if (offset < 0)
			then code @ [cons(result, xs)]
			else let
			  val res = DstV.new ("probe" ^ Int.toString offset, sampleTy)
			  val code = doVoxelSample (res, v, k, s, diffIter(), vars, offset) @ code
			  in
			    doSamples (offset-1, res::xs, code)
			  end
		  in
		    toImgSpaceCode @ doSamples (d-1, [], [])
		  end
	      | _ => raise Fail "image data with order > 1 not supported yet"
	    (* end case *)
	  end

  end