(* basis.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * Type definitions for Basis functions.
 *)

structure Basis =
  struct
    local
      structure N = BasisNames
      structure Ty = Types
      structure TV = TypeVar

      fun --> (tys1, ty) = Ty.T_Fun(tys1, [ty])
      infix -->

      val N2 = Ty.NatConst 2
      val N3 = Ty.NatConst 3

    (* short names for kinds *)
      val NK = Ty.TK_NAT
      val SK = Ty.TK_SHAPE
      val TK = Ty.TK_TYPE

      fun ty t = ([], t)
      fun all (kinds, mkTy) = let
	    val tvs = List.map (fn k => TV.new k) kinds
	    in
	      (tvs, mkTy tvs)
	    end
      fun allNK mkTy = let
	    val tv = TV.new NK
	    in
	      ([tv], mkTy tv)
	    end

    fun field (k, d, dd) = Ty.T_Field{diff=k, dim=d, shape=dd}
    fun tensor ds = Ty.T_Tensor(Ty.Shape ds)

    in

    val basis = [
	(* operators *)
	  (N.op_at, all([NK, NK, SK],
	    fn [k, d, dd] => let
		val k = Ty.NatVar k
		val d = Ty.NatVar d
		val dd = Ty.ShapeVar dd
		in
		  [field(k, d, dd), tensor[d]] --> Ty.T_Tensor dd
		end)),
(*
    val op_add = Atom.atom "+"
    val op_sub = Atom.atom "-"
    val op_mul = Atom.atom "*"
    val op_div = Atom.atom "/"
    val op_lt = Atom.atom "<"
    val op_lte = Atom.atom "<="
    val op_eql = Atom.atom "=="
    val op_neq = Atom.atom "!="
    val op_gte = Atom.atom ">="
    val op_gt = Atom.atom ">"
*)
	  (N.op_at, all([NK, NK, SK],
	    fn [k, d, dd] => let
		val k = Ty.NatVar k
		val d = Ty.NatVar d
		val dd = Ty.ShapeVar dd
		in
		  [field(k, d, dd), tensor[d]]
		    --> field(Ty.NatExp(k, ~1), d, Ty.ShapeExt(dd, d))
		end)),
(*
    val op_orelse = Atom.atom "||"
    val op_andalso = Atom.atom "&&"
*)
	  (N.op_norm, all([SK],
	    fn [dd] => [Ty.T_Tensor(Ty.ShapeVar dd)] --> Ty.realTy)),
	(* functions *)
	  (N.fn_CL,	ty([tensor[N3, N3]] --> Ty.vec3Ty)),
	  (N.fn_convolve, all([NK, NK, SK],
	    fn [k, d, dd] => let
		val k = Ty.NatVar k
		val d = Ty.NatVar d
		val dd = Ty.ShapeVar dd
		in
		  [Ty.T_Kernel k, Ty.T_Image{dim=d, shape=dd}]
		    --> field(k, d, dd)
		end)),
	  (N.fn_cos,	ty([Ty.realTy] --> Ty.realTy)),
	  (N.fn_dot,	allNK(fn tv => [tensor[Ty.NatVar tv]]
			  --> tensor[Ty.NatVar tv])),
(*
    val fn_inside = Atom.atom "inside"
*)
	  (N.fn_load,	all([NK, SK],
	    fn [d, dd] => let
		val d = Ty.NatVar d
		val dd = Ty.ShapeVar dd
		in
		  [Ty.T_String] --> Ty.T_Image{dim=d, shape=dd}
		end)),
	  (N.fn_pow,	ty([Ty.realTy, Ty.realTy] --> Ty.realTy)),
(*
    val fn_principleEvec = Atom.atom "principleEvec"
*)
	  (N.fn_sin,	ty([Ty.realTy] --> Ty.realTy)),
	(* kernels *)
	  (N.kn_bspln3,	ty(Ty.T_Kernel(Ty.NatConst 2))),
	  (N.kn_tent,	ty(Ty.T_Kernel(Ty.NatConst 0)))
	]

    end (* local *)
  end