(*open Ast*)


type t = 
  Int
| Float
| String
| Void
| Tensor of int

let string_of_styp = function
    Int -> "int"
  | Float -> "float"
  | String -> "string"
  | Void -> "void"
  | Tensor(i) -> (string_of_int i) ^ "-tensor"

type sbind = InferredBind of Ast.line * string * sexpr
| TensorBind of Ast.line * string * squalifier
| TensorBindAndAssign of Ast.line * string * squalifier * sexpr
and
sexpr =
  Id of Ast.line * t * string
| Noexpr of t
| Binop of Ast.line * t * sexpr * Ast.op * sexpr
| Unop of Ast.line * t * Ast.uop * sexpr
| Assign of Ast.line * t * string * sexpr
| Call of Ast.line * t * sfunc_decl * sexpr list
| Fliteral of Ast.line * t * float
| Iliteral of Ast.line * t * int
| Sliteral of Ast.line * t * string
| TensorId of Ast.line * t * string * squalifier
| TensorAssign of Ast.line * t * string * squalifier * sexpr
| Itensor of Ast.line * t * stimmediate
| BindExpr of Ast.line * t * sbind
and squalifier = sexpr list

and stimmediate = 
STenDim of int list * stimmediate list 
| SExprDim of int * sexpr list

 

and sfunc_decl = {
	ftyp : t;
	sfname : string;
	sformals : (t * string) list;
	sbody : sstmt
}

and variable_decl = { (* Sast.Id wrapper*)
	vtyp  : t;
	sname : string;
	vexpr : sexpr;
}

and sstmt = 
  SBlock of sstmt list	
| Expr of sexpr
| If of sexpr * sstmt * sstmt
| For of sexpr * sexpr * sexpr * sstmt
| While of sexpr * sstmt        	
| Return of sexpr
| Callback of (Llvm.llbuilder -> Llvm.llbuilder)

(* program entry point *)
type sprogram = {
       functions : sfunc_decl list;
       globals : variable_decl list;
}

let get_line_sexpr = function
    Id(l, _, _) -> l
  | Binop(l, _, _, _, _) -> l
  | Unop(l, _, _, _) -> l
  | Assign(l, _, _, _) -> l
  | Call(l, _, _, _) -> l
  | Fliteral(l, _, _) -> l
  | Iliteral(l, _, _) -> l
  | Sliteral(l, _, _) -> l
  | TensorId(l, _, _, _) -> l
  | TensorAssign(l, _, _, _, _) -> l
  | Itensor(l, _, _) -> l
  | BindExpr(l, _, _) -> l
  | Noexpr(_) -> raise(Exceptions.NoExprHasNoLineNumberException)

let get_line_sbind = function
    InferredBind(l, _, _) -> l
  | TensorBind(l, _, _) -> l
  | TensorBindAndAssign(l, _, _, _) -> l

let rec string_of_sbind = function
    InferredBind(_, s, e) -> "let " ^ s ^ " = (" ^ string_of_sexpr e ^ ")"
  | TensorBind(_, s, q) -> "let " ^ s ^ "_{" ^ String.concat ", " (List.map string_of_sexpr q) ^ "}"
  | TensorBindAndAssign(_, s, q, e) -> "let " ^ s ^ "_{" ^ String.concat ", " (List.map string_of_sexpr q) ^ "} = " ^
string_of_sexpr e

and string_of_stensor = function
   STenDim(dim, tl) -> "(" ^ String.concat " x "(List.map string_of_int dim)  ^ ") [" ^ String.concat ", " (List.map string_of_stensor tl) ^ "]"
   |SExprDim(i, el) ->  "(" ^ string_of_int i ^ ") [" ^ String.concat ", " (List.map string_of_sexpr el) ^ "]"



and string_of_sexpr = function
    Id(_, t, s) -> (string_of_styp t) ^ " " ^ s
   |Noexpr(t) -> (string_of_styp t) ^ " Noexpr"
   |Binop(_, t, e1, op, e2) -> (string_of_styp t) ^ " (" ^ (string_of_sexpr e1) ^ ") " ^ (Ast.string_of_op op) ^ " (" ^ (string_of_sexpr e2) ^ ")"
   |Unop(_, t, uop, e) -> (string_of_styp t) ^ " " ^ (Ast.string_of_uop uop) ^ " (" ^ (string_of_sexpr e) ^ ")"
   |Assign(_, t, var, e) -> (string_of_styp t) ^ " " ^ var ^ " = (" ^ (string_of_sexpr e) ^ ")"
   |Call(_, t, fn, el) -> (string_of_styp t) ^ " " ^ (string_of_sfunc_call fn) ^ "(" ^ String.concat "), (" (List.map string_of_sexpr el) ^ ")"
   |Fliteral(_, t, f) -> (string_of_styp t) ^ " " ^ string_of_float f
   |Iliteral(_, t, i) -> (string_of_styp t) ^ " " ^ string_of_int i
   |Sliteral(_, t, s) -> (string_of_styp t) ^ " " ^ s
   |TensorId(_, t, s, q) -> string_of_styp t ^ " " ^ s ^ "_{" ^ (String.concat ", " (List.map string_of_sexpr q))  ^ "}"
   |TensorAssign(_, t, s, sl, e) -> string_of_styp t ^ " " ^ s ^ "_{" ^ (String.concat ", " (List.map string_of_sexpr sl)) ^ ") =" ^ string_of_sexpr e
   |Itensor(_, t, i) -> string_of_styp t ^ " " ^ string_of_stensor i
   |BindExpr(_, t, b) -> (string_of_styp t) ^ " " ^ string_of_sbind b

and string_of_sfunc_call sfdecl =
   string_of_styp sfdecl.ftyp ^ " " ^
   sfdecl.sfname ^ "(" ^ String.concat ", " (List.map (fun n -> string_of_styp (fst n) ^ " " ^ (snd n)) sfdecl.sformals) ^ "): " (* ): *)

and string_of_sfunc_decl sfdecl = 
   string_of_styp sfdecl.ftyp ^ " " ^ 
   sfdecl.sfname ^ "(" ^ String.concat ", " (List.map (fun n -> string_of_styp (fst n) ^ " " ^ (snd n)) sfdecl.sformals) ^
  ")\n" ^
  string_of_sstmt sfdecl.sbody ^
  "\n"

and string_of_sstmt = function
    SBlock(stmts) -> "{\n" ^ String.concat "" (List.map string_of_sstmt stmts) ^ "}\n"
  | Expr(e) -> "(" ^ string_of_sexpr e ^ ");\n"
  | Return(e) -> "return (" ^ string_of_sexpr e ^ ");\n"
  | If(e, s, SBlock([])) -> "if (" ^ string_of_sexpr e ^ ")\n" ^ string_of_sstmt s
  | If(e, s1, s2) -> "if (" ^ string_of_sexpr e ^ ")\n" ^ string_of_sstmt s1 ^ " else\n" ^ string_of_sstmt s2
  | For(e1, e2, e3, s) -> "for (" ^ string_of_sexpr e1 ^ " ; " ^ string_of_sexpr e2 ^ " ; " ^ string_of_sexpr e3 ^ ") " ^ string_of_sstmt s
  | While(e, s) -> "while (" ^ string_of_sexpr e ^ ") " ^ string_of_sstmt s
  | Callback(_) -> ""    


let string_of_sglobal global = string_of_styp global.vtyp ^ " " ^ global.sname ^ ";\n"

let string_of_program sast = 
  String.concat "" (List.map string_of_sglobal sast.globals) ^ "\n" ^
  String.concat "\n" (List.map string_of_sfunc_decl sast.functions)
