(* gen-output.sml
 *
 * COPYRIGHT (c) 2012 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Generate strand output functions.  The output formats always have a single axis for the
 * data elements followed by one, or more, axes for the output structure.  There are four
 * cases that we handle:
 *
 *	grid, fixed-size elements:
 *		nrrd has object axis followed by grid axes
 *
 *	collection, fixed-size elements
 *		nrrd has object axis followed by a single axis
 *
 *	grid, dynamic-size elements
 *		nLengths nrrd has size 2 for objects (offset, length) followed by grid axes
 *		nData nrrd has object axis followed by a single axis
 *
 *	collection, dynamic-size elements
 *		nLengths nrrd has size 2 for objects (offset, length) followed by a single axis
 *		nData nrrd has object axis followed by a single axis
 *
 * The object axis kind depends on the output type, but it will either be one of the tensor types
 * that Teem knows about or else nrrdKindList.  In any case, the data elements are written as a
 * flat vector following the in-memory layout.  The other axes in the file will have nrrdKindSpace
 * as their kind.
 *
 * TODO: some of this code will be common across all targets (e.g., writing outputs to files), so
 * we will want to refactor it.
 *
 * TODO: for sequences of tensors (e.g., tensor[3]{2}), we should use a separate axis for the
 * sequence dimension with kind nrrdKindList.
 *)

structure GenOutput : sig

    val gen : TargetUtil.target_desc * int option -> (TreeIL.Ty.ty * string) list -> CLang.decl list

  end = struct

    structure IL = TreeIL
    structure V = IL.Var
    structure Ty = IL.Ty
    structure CL = CLang
    structure N = CNames
    structure Nrrd = NrrdEnums

    fun mapi f l = let
	  fun mapf (i, [], l) = List.rev l
	    | mapf (i, x::xs, l) = mapf (i+1, xs, f(i, x)::l)
	  in
	    mapf (0, l, [])
	  end

    val nrrdPtrTy = CL.T_Ptr(CL.T_Named "Nrrd")
    val filePtrTy = CL.T_Ptr(CL.T_Named "FILE")
    val sizeTy = CL.T_Named "size_t"
    fun wrldPtr tgt = CL.T_Ptr(CL.T_Named(N.worldTy tgt))
    fun mkInt i = CL.mkInt(IntInf.fromInt i)

  (* return information about the output type.  This is a tuple
   *
   *    (c-type, nrrd-type, nrrd-axis-kind, # elements)
   *)
    fun infoOf (tgt : TargetUtil.target_desc, ty) = (case ty
           of Ty.IntTy => if #longint tgt
		then (CL.int64, Nrrd.TypeLLong, Nrrd.KindScalar, 1)
		else (CL.int32, Nrrd.TypeInt, Nrrd.KindScalar, 1)
            | Ty.TensorTy [] => if #double tgt
		then (CL.double, Nrrd.TypeDouble, Nrrd.KindScalar, 1)
		else (CL.float, Nrrd.TypeFloat, Nrrd.KindScalar, 1)
            | Ty.TensorTy dims => let
                val (axisKind, nElems) = (case dims
                       of [2] => (Nrrd.Kind2Vector, 2)
                        | [3] => (Nrrd.Kind3Vector, 3)
                        | [4] => (Nrrd.Kind4Vector, 4)
                        | [2,2] => (Nrrd.Kind2DMatrix, 4)
                        | [3,3] => (Nrrd.Kind3DMatrix, 9)
                        | _ => (Nrrd.KindList, List.foldl Int.* 1 dims)
                      (* end case *))
                in
                  if #double tgt
                    then (CL.double, Nrrd.TypeDouble, axisKind, nElems)
                    else (CL.float, Nrrd.TypeFloat, axisKind, nElems)
                end
            | Ty.SeqTy(ty, n) => raise Fail "FIXME" (*let
		val (elemTy, nrrdTy, dims) = infoOf (tgt, ty)
		in
		  (elemTy, nrrdTy, n::dims)
		end*)
            | _ => raise Fail(concat["GetOutput.infoOf(", Ty.toString ty, ")"])
	  (* end case *))

  (* variables in the generated code *)
    val wrldV = CL.mkVar "wrld"
    val sizesV = CL.mkVar "sizes"
    val iV = CL.mkVar "i"
    val nV = CL.mkVar "n"
    val cpV = CL.mkVar "cp"
    val ipV = CL.mkVar "ip"
    val msgV = CL.mkVar "msg"
    val offsetV = CL.mkVar "offset"
    val nDataV = CL.mkVar "nData"
    val nLengthsV = CL.mkVar "nLengths"
    val numStableV = CL.mkVar "numStable"
    val numElemsV = CL.mkVar "numElems"
    val outSV = CL.mkVar "outS"
    val DIDEROT_STABLE = CL.mkVar "DIDEROT_STABLE"
    val NRRD = CL.mkVar "NRRD"

  (* dymanic sequence operations *)
    fun seqLength arg = CL.mkApply("Diderot_DynSeqLength", [arg])
    fun seqCopy (elemSz, dst, seq) = CL.mkApply("Diderot_DynSeqCopy", [elemSz, dst, seq])

  (* utility functions for initializing the sizes array *)
    fun sizes i = CL.mkSubscript(sizesV, mkInt i)
    fun setSizes (i, v) = CL.mkAssign(sizes i, v)

  (* code to access state variable
	wrld->outState[i].name
   *)
    fun stateVar name = CL.mkIndirect(CL.mkSubscript(CL.mkIndirect(wrldV, "outState"), iV), name)

  (* code fragment to loop over strands
	for (unsigned int i = 0;  i < wrld->numStrands;  i++) ...
   *)
    fun forStrands stm = CL.mkFor(
	  [(CL.uint32, "i", mkInt 0)],
	  CL.mkBinOp(iV, CL.#<, CL.mkIndirect(wrldV, "numStrands")),
	  [CL.mkPostOp(iV, CL.^++)],
	  stm)

  (* code fragment to test for stable strands in a loop
	if (wrld->status[i] == DIDEROT_STABLE)
	    ...
   *)
    fun ifStable stm = CL.mkIfThen(
	  CL.mkBinOp(CL.mkSubscript(CL.mkIndirect(wrldV, "status"), iV), CL.#==, DIDEROT_STABLE),
	  stm)

  (* code fragment to allocate nrrd data and check for errors
	if (nrrdMaybeAlloc_nva(<nrrdVar>, <nrrdType>, <nDims>, sizes) != 0) {
	    char *msg = biffGetDone(NRRD);
	    biffMsgAdd (wrld->errors, msg);
	    FREE (msg);
	    return true;
	}
   *)
    fun maybeAlloc (nrrdVar, nrrdType, nDims) = 
	  CL.mkIfThen(
	    CL.mkBinOp(
	      CL.mkApply("nrrdMaybeAlloc_nva", [
		  nrrdVar, CL.mkVar nrrdType, mkInt nDims, sizesV
		]),
	      CL.#!=,
	      CL.mkInt 0),
	  (* then *)
	    CL.mkBlock[
		CL.mkDeclInit(CL.charPtr, "msg", CL.mkApply("biffGetDone", [NRRD])),
		CL.mkCall("biffMsgAdd", [CL.mkIndirect(wrldV, "errors"), msgV]),
		CL.mkCall("FREE", [msgV]),
		CL.mkReturn(SOME(CL.mkVar "true"))
	      ]
	  (* endif*))

  (* code fragment to initialize the axes kinds; the data axis (axis[0]) is given, but we skip it
   * (by convention) if it is scalar. The other axes are always KindSpace.
   *)
    fun initAxisKinds (nrrd, dataAxisKind, nAxes) = let
        (* nData->axis[0].kind *)
          fun axisKind i = CL.mkSelect(CL.mkSubscript(CL.mkIndirect(nrrd, "axis"), mkInt i), "kind")
          fun init (i, k) = CL.mkAssign (axisKind i, CL.mkVar(Nrrd.kindToEnum  k))
          val (firstSpace, dataAxis) = (case dataAxisKind
                 of Nrrd.KindScalar => (0, [])
                  | _ => (1, [init(0, dataAxisKind)])
                (* end case *))
          in
            dataAxis @ List.tabulate(nAxes, fn i => init(i+firstSpace, Nrrd.KindSpace))
          end

  (* create the body of an output function for dynamic-size outputs.  The structure of the
   * function body is:
   *
   *	declarations
   *    compute sizes array for nLengths
   *	allocate nrrd for nLengths
   *    compute sizes array for nData
   *	allocate nrrd for nData
   *	copy data from strands to nrrd
   *)
    fun genDynOutput (tgt, nAxes, ty, name) = let
	  val (elemCTy, nrrdType, axisKind, nElems) = infoOf (tgt, ty)
	  val (isArray, nAxes) = (case nAxes of NONE => (false, 1) | SOME n => (true, n))
	(* declarations *)
	  val sizesDecl = CL.mkDecl(CL.T_Array(sizeTy, SOME(nAxes+1)), "sizes", NONE)
	(* count number of elements (and stable strands) *)
	  val countElems = let
		val nElemsInit = CL.mkDeclInit(CL.uint32, "numElems", CL.mkInt 0)
		val cntElems = CL.S_Exp(CL.mkAssignOp(numElemsV, CL.+=, seqLength(stateVar name)))
		in
		  if isArray
		    then [
			CL.mkComment["count number of elements"],
			nElemsInit, forStrands cntElems
		      ]
		    else [
			CL.mkComment["count number of output elements and stable strands"],
			CL.mkDeclInit(CL.uint32, "numStable", CL.mkInt 0),
			nElemsInit,
			forStrands(ifStable(CL.mkBlock[cntElems, CL.S_Exp(CL.mkPostOp(numStableV, CL.^++))]))
		      ]
		end
	(* generate code to allocate the nLengths nrrd *)
	  val lengthsNrrd = let
		val dimSizes = setSizes(0, CL.mkInt 2)  (* nLengths is 2-element vector *)
		in
		  CL.mkComment["allocate nLengths nrrd"] ::
		  (if isArray
		    then dimSizes ::
		      List.tabulate (nAxes, fn i =>
			setSizes(i+1, CL.mkSubscript(CL.mkIndirect(wrldV, "size"), mkInt(nAxes-i-1)))) @
		      [maybeAlloc (nLengthsV, Nrrd.tyToEnum Nrrd.TypeInt, nAxes+1)]
		    else [
			dimSizes, setSizes(1, numStableV),
			maybeAlloc (nLengthsV, Nrrd.tyToEnum Nrrd.TypeInt, 2)
		      ])
		end
	(* generate code to allocate the data nrrd *)
	  val dataNrrd = if (axisKind = Nrrd.KindScalar)
                then [ (* drop data axis for scalar data by convention *)
                    CL.mkComment["allocate nData nrrd"],
                    setSizes(0, numElemsV),
                    maybeAlloc (nDataV, Nrrd.tyToEnum nrrdType, 1)
                  ]
                else [
                    CL.mkComment["allocate nData nrrd"],
                    setSizes(0, mkInt nElems),
                    setSizes(1, numElemsV),
                    maybeAlloc (nDataV, Nrrd.tyToEnum nrrdType, 2)
                  ]
	(* generate the nLengths copy code *)
	  val copyLengths = let
		val pInit = CL.mkDeclInit(CL.T_Ptr CL.uint32, "ip",
		      CL.mkCast(CL.T_Ptr(CL.uint32), CL.mkIndirect(nLengthsV, "data")))
		val offsetDecl = CL.mkDeclInit(CL.uint32, "offset", CL.mkInt 0)
		val copyBlk = CL.mkBlock[
			CL.mkDeclInit(CL.uint32, "n", seqLength(stateVar name)),
			CL.mkAssign(CL.mkUnOp(CL.%*, CL.mkPostOp(ipV, CL.^++)), offsetV),
			CL.mkAssign(CL.mkUnOp(CL.%*, CL.mkPostOp(ipV, CL.^++)), nV),
			CL.S_Exp(CL.mkAssignOp(offsetV, CL.+=, nV))
		      ]
		val copyStm = if isArray
		      then copyBlk
		      else ifStable copyBlk
		in
                  CL.mkComment["initialize nLengths nrrd"] ::
                  pInit ::
                  offsetDecl ::
                  forStrands copyStm ::
                  initAxisKinds (nLengthsV, Nrrd.Kind2Vector, nAxes)
		end
	(* generate the nData copy code *)
	  val copyData = let
		val pInit = CL.mkDeclInit(CL.charPtr, "cp",
		      CL.mkCast(CL.charPtr, CL.mkIndirect(nDataV, "data")))
		val copyStm = CL.mkAssign(cpV, seqCopy(
		      CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof(elemCTy)), cpV, stateVar name))
		val copyStm = if isArray
		      then copyStm
		      else ifStable copyStm
		in
		  CL.mkComment["initialize nLengths nrrd"] ::
                  pInit ::
                  forStrands copyStm ::
                  initAxisKinds (nDataV, axisKind, 1)
		end
	(* the function body *)
	  val stms =
		sizesDecl ::
		countElems @
		lengthsNrrd @
		dataNrrd @
		copyLengths @
		copyData @
		[CL.mkReturn(SOME(CL.mkVar "false"))]	  
	  in
	    ([CL.PARAM([], nrrdPtrTy, "nLengths"), CL.PARAM([], nrrdPtrTy, "nData")], CL.mkBlock stms)
	  end

  (* create the body of an output function for fixed-size outputs.  The structure of the
   * function body is:
   *
   *	declare and compute sizes array
   *	allocate nrrd nData
   *	copy data from strands to nrrd
   *)
    fun genFixedOutput (tgt, nAxes, ty, name) = let
	  val (elemCTy, nrrdType, axisKind, nElems) = infoOf (tgt, ty)
	  val (isArray, nAxes) = (case nAxes of NONE => (false, 1) | SOME n => (true, n))
          val nDataAxes = if (axisKind = Nrrd.KindScalar) then 0 else 1
	(* generate the sizes initialization code *)
	  val initSizes = let
		val dimSizes = let
                      val dcl = CL.mkDecl(CL.T_Array(sizeTy, SOME(nAxes+nDataAxes)), "sizes", NONE)
                      in
                        if (axisKind = Nrrd.KindScalar)
                          then [dcl]
                          else [dcl, setSizes(0, mkInt nElems)]
                      end
		in
		  if isArray
		    then dimSizes @
		      List.tabulate (nAxes, fn i =>
			setSizes(i+nDataAxes, CL.mkSubscript(CL.mkIndirect(wrldV, "size"), mkInt(nAxes-i-1))))
		    else
		      CL.mkDeclInit(sizeTy, "numStable", mkInt 0) ::
		      forStrands (ifStable(CL.S_Exp(CL.mkPostOp(numStableV, CL.^++)))) ::
		      dimSizes @ [setSizes(nDataAxes, numStableV)]
		end
	(* generate the copy code *)
	  val copyCode = let
		val pDecl = CL.mkDeclInit(CL.charPtr, "cp",
		      CL.mkCast(CL.charPtr, CL.mkIndirect(nDataV, "data")))
		val copyBlk = CL.mkBlock[
			CL.mkCall("memcpy", [
			    cpV,
			    CL.mkUnOp(CL.%&, stateVar name),
			    CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof elemCTy)
			  ]),
			CL.mkExpStm(CL.mkAssignOp(cpV, CL.+=,
			  CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof elemCTy)))
		      ]
		val copyStm = if isArray
		      then copyBlk
		      else ifStable copyBlk
		in
		  pDecl :: forStrands copyStm :: initAxisKinds (nDataV, axisKind, nAxes)
		end
	(* the function body *)
	  val stms =
		CL.mkComment["Compute sizes of nrrd file"] ::
		initSizes @
		CL.mkComment["Allocate nData nrrd"] ::
		maybeAlloc (nDataV, Nrrd.tyToEnum  nrrdType, nAxes+nDataAxes) ::
		CL.mkComment["copy data to output nrrd"] ::
		copyCode @
		[CL.mkReturn(SOME(CL.mkVar "false"))]
	  in
	    ([CL.PARAM([], nrrdPtrTy, "nData")], CL.mkBlock stms)
	  end

  (* generate the nrrd-file output and print functions used by standalone executables *)
    fun genOutput (tgt : TargetUtil.target_desc, outputs) = let
	  fun isDyn ty = (case ty of Ty.DynSeqTy _ => true | _ => false)
(* FIXME: use biffMsgAddF and return error status *)
	  fun error (fmt, msg) = CL.mkBlock[
		  CL.mkCall("fprintf", [CL.mkVar "stderr", CL.mkStr fmt, msg]),
		  CL.mkCall("exit", [CL.mkInt 1])
		]
	  val outDecls = if List.exists (isDyn o #1) outputs
		then [CL.mkDecl(nrrdPtrTy, "nLengths", NONE), CL.mkDecl(nrrdPtrTy, "nData", NONE)]
		else [CL.mkDecl(nrrdPtrTy, "nData", NONE)]
	  val prDecls = outDecls @ [CL.mkDecl(filePtrTy, "outS", NONE)]
	  fun nrrdNew v = CL.mkAssign(v, CL.mkApply("nrrdNew", []))
	  fun nrrdNuke v = CL.mkCall("nrrdNuke", [v])
	  fun writeNrrd (ty, name) =
		if isDyn ty
		  then [
		      nrrdNew (nLengthsV),
		      nrrdNew (nDataV),
		      CL.mkIfThenElse(
			CL.mkApply(N.outputGet(tgt, name), [wrldV, nLengthsV, nDataV]),
		      (* then *)
			  error ("Error getting nrrd data: %s\n",
			    CL.mkApply("biffMsgStrGet", [CL.mkIndirect(wrldV, "errors")])),
		      (* else *)
			CL.mkIfThen(
			  CL.mkBinOp(
			    CL.mkApply("nrrdSave", [
				CL.mkStr(OS.Path.joinBaseExt{base=name^"-len", ext=SOME "nrrd"}),
				nLengthsV, CL.mkVar "NULL"
			      ]),
			    CL.#||,
			    CL.mkApply("nrrdSave", [
				CL.mkStr(OS.Path.joinBaseExt{base=name^"-data", ext=SOME "nrrd"}),
				nDataV, CL.mkVar "NULL"
			      ])),
			(* then *)
			  error ("Error saving nrrd: %s\n", CL.mkApply("biffGetDone", [NRRD]))
			(* endif *))
		      (* endif *)),
		      nrrdNuke nLengthsV,
		      nrrdNuke nDataV
		    ]
		  else [
		      nrrdNew (nDataV),
		      CL.mkIfThenElse(
			CL.mkApply(N.outputGet(tgt, name), [wrldV, nDataV]),
		      (* then *)
			  error ("Error getting nrrd data: %s\n",
			    CL.mkApply("biffMsgStrGet", [CL.mkIndirect(wrldV, "errors")])),
		      (* else *)
			CL.mkIfThen(
			  CL.mkApply("nrrdSave", [
			      CL.mkStr(OS.Path.joinBaseExt{base=name, ext=SOME "nrrd"}),
			      nDataV, CL.mkVar "NULL"
			    ]),
			(* then *)
			  error ("Error saving nrrd: %s\n", CL.mkApply("biffGetDone", [NRRD]))
			(* endif *))
		      (* endif *)),
		      nrrdNuke nDataV
		    ]
	  fun printNrrd (ty, name) = [] (* FIXME *)
	  in [
	    CL.D_Func(["static"], CL.voidTy, "WriteOutput", [CL.PARAM([], wrldPtr tgt, "wrld")],
	      CL.mkBlock(outDecls @ List.foldr (fn (output, l) => writeNrrd output @ l) [] outputs)),
	    CL.D_Func(["static"], CL.voidTy, "PrintOutput", [CL.PARAM([], wrldPtr tgt, "wrld")],
	      CL.mkBlock(prDecls @ List.foldr (fn (output, l) => printNrrd output @ l) [] outputs))
	  ] end
		      
    fun gen (tgt : TargetUtil.target_desc, nAxes) = let
	  fun getFn (ty, name) = let
		val funcName = N.outputGet(tgt, name)
		fun mkFunc (params, body) =
		      CL.D_Func([], CL.boolTy, funcName, CL.PARAM([], wrldPtr tgt, "wrld")::params, body)
		in
		  case ty
		   of Ty.DynSeqTy ty' => mkFunc (genDynOutput(tgt, nAxes, ty', name))
		    | _ => mkFunc (genFixedOutput(tgt, nAxes, ty, name))
		  (* end case *)
		end
	  fun gen' outputs = let
		val getFns = List.map getFn outputs
		in
		  if (#exec tgt)
		    then getFns @ genOutput(tgt, outputs)
		    else getFns
		end
	  in
	    gen'
	  end

  end