(* eval.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Evaluation of "static" expressions.
 *)

structure Eval : sig

  (* raised if there is an error due to faulty code or input values (e.g., loading an
   * image of the wrong shape.
   *)
    exception Error of string list

    datatype value
      = BV of bool
      | SV of string
      | IV of IntInf.int
      | RV of real
      | TV of value list        (* tensors: values will either be RV or TV *)
      | ImgV of ImageInfo.info * Var.var

    val evalStatics : Var.Set.set * Simple.block -> value Var.Map.map

  end = struct

    structure Ty = Types
    structure BV = BasisVars
    structure S = Simple
    structure VMap = Var.Map
    structure VSet = Var.Set
    structure VTbl = Var.Tbl

    exception Error of string list

    datatype value
      = BV of bool
      | SV of string
      | IV of IntInf.int
(* FIXME: we probably should use FloatLit.float values instead of reals! *)
      | RV of real
      | TV of value list        (* tensors: values will either be RV or TV *)
      | ImgV of ImageInfo.info * Var.var

    fun toString (BV b) = Bool.toString b
      | toString (IV i) = IntInf.toString i
      | toString (SV s) = concat["\"", String.toString s, "\""]
      | toString (RV r) = Real.toString r
      | toString (TV _) = "tensor"
(* FIXME: should include x in output *)
      | toString (ImgV(info, x)) = ImageInfo.toString info

    val tbl : (value list -> value) VTbl.hash_table = let
          val tbl = VTbl.mkTable (128, Fail "Eval table")
          fun intBinOp rator [IV a, IV b] = IV(rator(a, b))
          fun tensorBinOp rator [v1, v2] = let
                fun f (TV v1, TV v2) = TV(ListPair.mapEq f (v1, v2))
                  | f (RV r1, RV r2) = RV(rator(r1, r2))
                in
                  f (v1, v2)
                end
          fun realBinOp rator [RV a, RV b] = RV(rator(a, b))
          fun realUnOp rator [RV a] = RV(rator a)
          fun intCmp rator [IV a, IV b] = BV(rator(a, b))
          fun realCmp rator [RV a, RV b] = BV(rator(a, b))
          fun boolCmp rator [BV a, BV b] = BV(rator(a, b))
          fun stringCmp rator [SV a, SV b] = BV(rator(a, b))
          in
            List.app (VTbl.insert tbl) [
                (BV.add_ii,             intBinOp (op +)),
                (BV.add_tt,             tensorBinOp (op +)),
                (BV.sub_ii,             intBinOp (op -)),
                (BV.sub_tt,             tensorBinOp (op -)),
                (BV.mul_ii,             intBinOp (op * )),
                (BV.mul_rr,             realBinOp (op * )),
(*
                (BV.mul_rt,             tensorOp Op.Scale),
                (BV.mul_tr,             fn (y, sv, [t, r]) => tensorOp Op.Scale (y, sv, [r, t])),
*)
                (BV.div_ii,             intBinOp IntInf.quot),
                (BV.div_rr,             realBinOp (op /)),
(*
                (BV.div_tr,             tensorOp Op.InvScale),
*)
                (BV.lt_ii,              intCmp (op <)),
                (BV.lt_rr,              realCmp (op <)),
                (BV.lte_ii,             intCmp (op <=)),
                (BV.lte_rr,             realCmp (op <=)),
                (BV.gte_ii,             intCmp (op >=)),
                (BV.gte_rr,             realCmp (op >=)),
                (BV.gt_ii,              intCmp (op >)),
                (BV.gt_rr,              realCmp (op >)),
                (BV.equ_bb,             boolCmp (op =)),
                (BV.equ_ii,             intCmp (op =)),
                (BV.equ_ss,             stringCmp (op =)),
                (BV.equ_rr,             realCmp Real.==),
                (BV.neq_bb,             boolCmp (op <>)),
                (BV.neq_ii,             intCmp (op <>)),
                (BV.neq_ss,             stringCmp (op <>)),
                (BV.neq_rr,             realCmp Real.!=),
                (BV.neg_i,              fn [IV i] => IV(~i)),
(*
                (BV.neg_t,              tensorOp Op.Neg),
                (BV.neg_f,              fn [FV fld] => FV(FieldDef.neg fld)),
                (BV.op_at,              fn (y, _, xs) => assign(y, Op.Probe, xs)),
                (BV.op_D,               fn [FV fld] => FV(FieldDef.diff fld)),
                (BV.op_norm,            tensorOp Op.Norm),
*)
                (BV.op_not,             fn [BV b] => BV(not b)),
(*
                (BV.fn_CL,              fn (y, _, xs) => assign(y, Op.CL, xs)),
                (BV.op_convolve,        fn [Img info, KV h] => FV(FieldDef.CONV(0, info, h))),
                (BV.fn_inside,          fn (y, _, xs) => assign(y, Op.Inside, xs)),
*)
                (BV.fn_max,             realBinOp Real.min),
                (BV.fn_min,             realBinOp Real.max),
                (BV.fn_modulate,        tensorBinOp (op * )),
(*
                (BV.fn_principleEvec,   vectorOp Op.PrincipleEvec),
                (BV.kn_bspln3,          kernel Kernel.bspln3),
                (BV.kn_bspln5,          kernel Kernel.bspln5),
                (BV.kn_ctmr,            kernel Kernel.ctmr),
                (BV.kn_tent,            kernel Kernel.tent),
                (BV.kn_c1tent,          kernel Kernel.c1tent),
                (BV.kn_c2ctmr,          kernel Kernel.c2ctmr),
*)
                (BV.i2r,                fn [IV i] => RV(real(IntInf.toInt i)))
              ];
            tbl
          end

    fun loadImage ([Ty.DIM dim, Ty.SHAPE shp], SV filename) = let
          val Ty.DimConst dim = TypeUtil.resolveDim dim
          val dd = let
                val Ty.Shape dd = TypeUtil.resolveShape shp
                fun doDim (Ty.DimConst d) = d
                  | doDim (Ty.DimVar d) = let val Ty.DimConst d = TypeUtil.resolveDim d in d end
                in
                  List.map doDim dd
                end
          val nrrd = NrrdInfo.getInfo filename
          in
            case ImageInfo.mkInfo (nrrd, dim, dd)
             of SOME info => info
              | NONE => let
                  fun rngToS [] = "real"
                    | rngToS dd = concat["tensor[", String.concatWith "," (List.map Int.toString dd), "]"]
                  fun error msg = raise Error("image file \"" :: filename :: "\" " :: msg)
                  in
                  (* figure out what the mismatch is *)
                    if (NrrdInfo.dim nrrd <> dim)
                      then error [
                          "has dimension ", Int.toString(NrrdInfo.dim nrrd),
                          ", expected ", Int.toString dim
                        ]
                      else error ["has unexpected range ", rngToS dd]
                  end
            (* end case *)
          end

    fun evalVar env x = (case VMap.find (env, x)
           of SOME v => v
            | NONE => raise Fail("undefined variable " ^ Var.uniqueNameOf x)
          (* end case *))

    fun apply (env, f, mvs, xs) =
          if List.all (fn x => VMap.inDomain(env, x)) xs
            then (* try *)(
(* FIXME: what about fn_load? *)
                if Var.same(f, BV.fn_image)
                  then let
                    val [imgName] = xs
                    in
                      SOME(ImgV(loadImage(mvs, evalVar env imgName), imgName))
                    end
                  else (case VTbl.find tbl f
                     of SOME evalFn => SOME(evalFn (List.map (evalVar env) xs))
                      | NONE => NONE
                    (* end case *))
                ) handle ex as Error msg => raise ex
                       | ex => (
                           TextIO.output (TextIO.stdErr, concat [
                               Var.uniqueNameOf f, "(",
                               String.concatWith "," (List.map Var.uniqueNameOf xs),
                               ") fails with exception ", exnName ex, "\n"
                             ]);
              raise ex)
            else NONE

    fun evalExp (env, e) = (case e
           of S.E_Var x => VMap.find (env, x)
            | S.E_Lit(Literal.Int i) => SOME(IV i)
            | S.E_Lit(Literal.Float f) => SOME(RV(FloatLit.toReal f))
            | S.E_Lit(Literal.String s) => SOME(SV s)
            | S.E_Lit(Literal.Bool b) => SOME(BV b)
            | S.E_Tuple _ => raise Fail "E_Tuple"
            | S.E_Apply(f, mvs, xs, _) => apply(env, f, mvs, xs)
            | S.E_Cons xs => (case evalArgs(env, xs)
                 of NONE => NONE
                  | SOME vs => SOME(TV vs)
                (* end case *))
            | S.E_Slice(x, indices, _) => (case VMap.find (env, x)
                 of SOME v => let
                      fun slice (TV vs, SOME ix :: ixs) = (case VMap.find (env, ix)
                             of SOME(IV i) => slice (List.nth(vs, IntInf.toInt i), ixs)
                              | NONE => raise Subscript
                            (* end case *))
                        | slice (TV vs, NONE :: ixs) =
                            TV(List.map (fn v => slice(v, ixs)) vs)
                        | slice (v, []) = v
                      in
                        SOME(slice(v, indices)) handle Subscript => NONE
                      end
                  | _ => NONE
                (* end case *))
            | S.E_Input(ty, name, desc, optDefault) => raise Fail "impossible"
            | S.E_LoadImage info => SOME(ImgV info)
            | S.E_Coerce{srcTy, dstTy, x} => NONE
          (* end case *))

    and evalArgs (env, args) = let
          fun eval ([], vs) = SOME(List.rev vs)
            | eval (x::xs, vs) = (case VMap.find(env, x)
                 of SOME v => eval(xs, v::vs)
                  | NONE => NONE
                (* end case *))
          in
            eval (args, [])
          end

    fun getInput (ty, name, optDefault) = (case ty
           of Ty.T_Bool =>
                Inputs.getInput(name, (Option.map BV) o Bool.fromString, optDefault)
            | Ty.T_Int =>
                Inputs.getInput(name, (Option.map IV) o IntInf.fromString, optDefault)
            | Ty.T_String => Inputs.getInput(name, fn s => SOME(SV s), optDefault)
            | Ty.T_Tensor(Ty.Shape[]) =>
                Inputs.getInput(name, (Option.map RV) o Real.fromString, optDefault)
            | Ty.T_Tensor(Ty.Shape[Ty.DimConst d]) => let
                fun fromString s = let
                    (* first split into fields by "," *)
                      val flds = String.fields (fn #"," => true | _ => false) s
                    (* then tokenize by white space and flatten *)
                      val toks = List.concat(List.map (String.tokens Char.isSpace) flds)
                    (* then convert to reals *)
                      val vals = List.map (RV o valOf o Real.fromString) toks
                      in
                        if (List.length vals = d)
                          then SOME(TV(vals))
                          else NONE
                      end
                      handle _ => NONE
                in
                  Inputs.getInput(name, fromString, optDefault)
                end
            | Ty.T_Tensor shp => raise Fail "TODO: general tensor inputs"
            | _ => raise Fail(concat[
                  "input ", name, " has invalid type ", TypeUtil.toString ty
                ])
          (* end case *))

    fun evalStatics (statics, blk) = let
          fun evalBlock (env, S.Block stms) = let
                exception Done of value VMap.map
                fun evalStm (stm, env) = (case stm
                       of S.S_Var _ => raise Fail "unexpected variable decl"
                        | S.S_Assign(x, S.E_Input(ty, name, desc, optDefault)) =>
                            if VSet.member(statics, x)
                              then let
                                val optDefault = Option.map (evalVar env) optDefault
                                val input = getInput (ty, name, optDefault)
                                in
                                  case input
                                   of SOME v => VMap.insert(env, x, v)
                                    | NONE => raise Fail("error getting required input " ^ name)
                                  (* end case *)
                                end
                              else env
                        | S.S_Assign(x, e) => (case evalExp(env, e)
                             of SOME v =>
(Log.msg(concat["eval assignment: ", Var.uniqueNameOf x, " = ", toString v, "\n"]);
                                VMap.insert(env, x, v)
)
                              | NONE => env
                            (* end case *))
                        | S.S_IfThenElse(x, b1, b2) => (case VMap.find(env, x)
                             of SOME(BV true) => evalBlock(env, b1)
                              | SOME(BV false) => evalBlock(env, b2)
                              | SOME _ => raise Fail "type error"
                              | NONE => raise (Done env)
                            (* end case *))
                        | S.S_New _ => raise Fail "unexpected new strand"
                        | S.S_Die => raise Fail "unexpected die"
                        | S.S_Stabilize => raise Fail "unexpected stabilize"
                        | S.S_Print _ => raise Fail "unexpected print"
                      (* end case *))
                in
                  (List.foldl evalStm env stms) handle Done env => env
                end
          in
            evalBlock (VMap.empty, blk)
          end

  end