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

structure Probe : sig

    val expand : MidIL.var * FieldDef.field_def * HighIL.var -> HighIL.assign list

  end = struct

    structure SrcIL = HighIL
    structure SrcOp = HighOps
    structure DstIL = MidIL
    structure DstOp = MidOps
    structure VMap = SrcIL.Var.Map

  (* a tree representation of nested iterations over the image space, where the
   * height of the tree corresponds to the number of dimensions and at each node
   * we have as many children as there are iterations.
   *)
    structure IT =
      struct
	datatype ('nd, 'lf) iter_tree
	  = LF of 'lf
	  | ND of ('nd * ('nd, 'lf) iter_tree list)

	fun create (depth, width, ndAttr, f, lfAttr, init) = let
	      fun mk (d, i, arg) = if (d < depth)
		    then ND(ndAttr arg, List.tabulate(width, fn j => mk(d+1, j, f(j, arg))))
		    else LF(lfAttr arg)
	      in
		mk (0, 0, init)
	      end

	fun map (nd, lf) t = let
	      fun mapf (LF x) = LF(lf x)
		| mapf (ND(i, kids)) = ND(nd i, List.map mapf kids)
	      in
		mapf t
	      end

	fun foldr f init t = let
	      fun fold (LF x, acc) = f(x, acc)
		| fold (ND(_, kids), acc) = List.foldr fold acc kids
	      in
		fold t
	      end

      end
  (* generate a new variable indexed by dimension *)
    local
      val dimNames = Vector.fromList[ "x", "y", "z" ];
    in
    fun newVar_dim (prefix, d) =
	  DstIL.Var.new (prefix ^ Vector.sub(dimNames, d))

    fun assign (x, rator, args) = (x, DstIL.OP(rator, args))
    fun cons (x, args) = (x, DstIL.CONS 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 probing the field (D^k (v * h)) at pos *)
    fun probe (result, (k, v, h), pos) = let
	  val ImageInfo.ImgInfo{dim, ty=([], ty), ...} = v
	  val dimTy = DstOp.VecTy dim
	  val s = Kernel.support h
	  val sTy = DstOp.VecTy(2*s)
	(* generate the transform code *)
	  val x = DstIL.Var.new "x"	(* image-space position *)
	  val f = DstIL.Var.new "f"
	  val nd = DstIL.Var.new "nd"
	  val n = DstIL.Var.new "n"
	  val transformCode = [
		  assign(x, DstOp.Transform v, [pos]),
		  assign(nd, DstOp.Floor dim, [x]),
		  assign(f, DstOp.Sub dimTy, [x, nd]),
		  assign(n, DstOp.TruncToInt dim, [nd])
		]
	(* generate code to load the voxel data *)
	  val voxIter = let
		fun f (i, (offsets, id)) = (i - (s - 1) :: offsets, i::id)
		fun g (offsets, id) = {
			offsets = offsets,
			vox = DstIL.Var.new(String.concat("v" :: List.map Int.toString id))
		      }
		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)
			    val t1 = DstIL.Var.new "t1"
			    val t2 = DstIL.Var.new "t2"
			    val (indices, code) = computeIndices (i+1, offsets)
			    val code =
				  intLit(t1, offset) ::
				  assign(t2, DstOp.Select i, [n]) ::
				  assign(index, DstOp.Add(DstOp.IntTy), [t1, t2]) ::
				  code
			    val indices = index::indices
			    in
			      (indices, code)
			    end
		      val (indices, indicesCode) = computeIndices (0, ~(s-1) :: offsets)
		      val a = DstIL.Var.new "a"
		      in
			indicesCode :: [
			    assign(a, DstOp.VoxelAddress v, indices),
			    assign(vox, DstOp.LoadVoxels(ty, 2*s))
			  ] @ code
		      end
		in
		  IT.foldr genCode [] voxIter
		end
	  val voxVars = IT.foldr (fn ({vox, ...}, vs) => vox::vs) [] voxIter
	(* generate the code for computing the convolution coefficients *)
	  val convCoeffs = Vector.tabulate (dim,
		fn d => Vector.tabulate (k+1,
		  fn 0 => newVar_dim("h", d)
		   | 1 => newVar_dim("dh", d)
		   | i => newVar_dim(concat["d", Int.toString i, "h"], d)))
	  fun coefficient (d, k) = Vector.sub(Vector.sub(convCoeffs, d), k)
	  fun genCoeffCode d = if (d < dim)
		then let
		  val code = genCoeffCode (d+1)
		  val fd = newVar_dim "f"
		  val a = DstIL.Var.new "a"
		  val tmps = List.tabulate(2*s,
			fn i => (DstIL.Var.new("t"^Int.toString i), s - (i+1)))
		  fun mkArg ((t, n), code) = let
			val t' = DstIL.Var.new "r"
			in
			  realLit (t', n) ::
			  assign (t, DstOp.Add DstOp.realTy, [fd, t']) ::
			  code
			end
		  val coeffCode =
			cons(a, List.map #1 tmps) ::
			List.tabulate (k+1, fn i =>
			  assign(coefficient(d, i), DstOp.EvalKernel(2*s, h, i), [a]))
		  in
		    assign(fd, DstOp.Select d, f) ::
		    (List.foldr mkArg (coeffCode @ code) tmps)
		  end
		else []
	  val coeffCode = genCoeffCode
	(* generate the reduction code *)
	  fun genReduce (d, IT.ND(kids), code) =
		if (d < dim)
		  then List.foldr (fn (nd, code) => genReduce(d+1, nd, code)) code kids
		  else let (* the kids will all be leaves *)
		    val vv = DstIL.Var.new "vv"
		    fun getVox (IT.LF{vox, offsets}) = vox
		    val hh = coefficient (d, 0)	(* FIXME: what is the right value for k? *)
		    in
		      cons (vv, List.map getVox kids) ::
		      assign (t, DstOp.Dot, [hh, vv]) :: code
		    end
	  val reduceCode = genReduce (1, voxIter, [])
	  in
	    transformCode @ loadCode @ coeffCode @ reduceCode
	  end

    end

    fun expand (result, FieldDef.CONV(0, img, h), pos) = let
	  fun expand' (result, FieldDef.CONV(k, v, h)) = let
		val x = DstIL.Var.new "x"
		val xformStm = (x, DstIL.OP(DstOp.Transform img, [pos']))
		in
		  probe (result, (k, v, h), x) @ [xformStm]
		end
	    | expand' (result, FieldDef.NEG fld) = let
		val r = DstIL.Var.new "value"
		val stms = expand' (r, fld)
		in
		  expand' (r, fld) @ [assign(r, DstOp.Neg ty, [r])]
		end
	    | expand' (result, FieldDef.SUM(fld1, dlf2)) = raise Fail "expandInside: SUM"
	  in
	    List.rev (expand (result, fld))
	  end

  end