type op = Add | Sub | Mult | Div | Equal | Neq |
      	Less | Leq | Greater | Geq | And | Or | Concat
type uop = Neg | Not

type line = string

type comment = (string * string) list

type typ = Void | Float | Int | String | TensorType of qualifier_rettype
and
qualifier_rettype = expr list
and
accessor = IntIndex of int | ExprIndex of expr
and
bind = InferredBind of line * string * expr
| TensorBind of line * string * qualifier_rettype
| TensorBindAndAssign of line * string * qualifier_rettype * expr
and
expr =
  Id of line * string                	| Noexpr
| Binop of line * expr * op * expr   	| Unop of line * uop * expr
| Assign of line * string * expr     	| Call of line * string * expr list
| Iliteral of line * int
| Fliteral of line * float
| Sliteral of line * string
| TensorId of line * string * qualifier_rettype
| TensorAssign of line * string * qualifier_rettype * expr
| Tensor of line * expr list
| BindExpr of line * bind


type stmt = Block of comment * stmt list
| Expr of comment * expr
| If of comment * expr * stmt * stmt
| For of comment * expr * expr * expr * stmt
| While of comment * expr * stmt
| Return of comment * expr



type func_decl = {
  comm : (line * string) list;
	typ : typ;
	fname : string;
	formals : (typ * string) list;
	body : stmt
}

type program = bind list * func_decl list

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

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

let string_of_op = function
    Add -> "+"
  | Sub -> "-"
  | Mult -> "*"
  | Div -> "/"
  | Equal -> "=="
  | Neq -> "~="
  | Less -> "<"
  | Leq -> "<="
  | Greater -> ">"
  | Geq -> ">="
  | And -> "&&"
  | Or -> "||"
  | Concat -> "^"

let string_of_uop = function
    Neg -> "-"
  | Not -> "~"

let string_of_qual_list  q = String.concat ", " q 

let rec string_of_bind = function
    InferredBind(_, s, e) -> "let " ^ s ^ " = " ^ string_of_expr e
  | TensorBind(_, s, q) -> "let " ^ s ^ "_{" ^ string_of_qual_list_expr q ^ "}"
  | TensorBindAndAssign(_, s, q, e) -> "let " ^ s ^ "_{" ^ string_of_qual_list_expr q ^ "} = " ^
        string_of_expr e


and string_of_qual_list_expr q = String.concat ", " (List.map string_of_expr q) 

and string_of_expr = function
    Id(_,s) -> s
  | Binop(_, e1, op, e2) -> string_of_expr e1 ^ " " ^ string_of_op op ^ " " ^ string_of_expr e2
  | Unop(_,uop, e) -> string_of_uop uop ^ " " ^ string_of_expr e
  | Assign(_,var, e) -> var ^ " = " ^ string_of_expr e
  | Iliteral(_,i) -> string_of_int i
  | Sliteral(_,s) -> s
  | Fliteral(_,f) -> string_of_float f
  | BindExpr(_,b) -> string_of_bind b
  | Noexpr -> ""
  | Call(_,fn, el) -> fn ^ "(" ^ String.concat ", " (List.map string_of_expr el) ^ ")"
  | TensorId(_, s, q) -> s ^ "_{" ^ string_of_qual_list_expr q  ^ "}"
  | TensorAssign(_, s, sl, e) -> s ^ "_{" ^ string_of_qual_list_expr sl ^
                              "} = " ^ string_of_expr e
  | Tensor(_, t) -> "[" ^ String.concat ", " (List.map string_of_expr t)  ^ "]"

(*and string_of_tensor = function
   |ExprDim(el) -> "[" ^ String.concat ", " (List.map string_of_expr el) ^ "]"*)

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


let string_of_typ = function
    Void -> "void"
  | Float -> "float"
  | Int -> "int"
  | String -> "string"
  | TensorType(q) -> string_of_int (List.length q)

let string_of_fdecl fdecl =
  string_of_typ fdecl.typ ^ " " ^
  fdecl.fname ^ "(" ^ String.concat ", " (List.map snd fdecl.formals) ^
  ")\n" ^
  string_of_stmt fdecl.body

let string_of_program (globals, functions) =
  String.concat "\n" (List.map string_of_bind globals) ^ "\n" ^
  String.concat "\n" (List.map string_of_fdecl functions)





(* Pretty Printing for LaTeX file generation*)
let rec pretty_print_binop (e1, op, e2) = match op with 
    Add -> "(" ^ pretty_print_expr e1 ^ "+" ^ pretty_print_expr e2 ^ ")"
  | Sub -> "(" ^ pretty_print_expr e1 ^ "-" ^ pretty_print_expr e2 ^ ")"
  | Mult -> pretty_print_expr e1 ^ "\\cdot " ^ pretty_print_expr e2
  | Div -> "\\frac {" ^ pretty_print_expr e1 ^ "} {" ^ pretty_print_expr e2 ^ "}"
  | Equal -> pretty_print_expr e1 ^ "==" ^ pretty_print_expr e2 
  | Neq -> pretty_print_expr e1 ^ "\\neq" ^ pretty_print_expr e2
  | Less -> pretty_print_expr e1 ^ "<" ^ pretty_print_expr e2
  | Leq -> pretty_print_expr e1 ^ "\\leq" ^ pretty_print_expr e2
  | Greater -> pretty_print_expr e1 ^ ">" ^ pretty_print_expr e2
  | Geq -> pretty_print_expr e1 ^ "\\geq" ^ pretty_print_expr e2
  | And -> pretty_print_expr e1 ^ "\\& \\&" ^ pretty_print_expr e2
  | Or -> pretty_print_expr e1 ^ "||" ^ pretty_print_expr e2
  | Concat -> pretty_print_expr e1 ^ "\\hbox{ \\string^ }" ^ pretty_print_expr e2

and pretty_print_uop = function
    Neg -> "-"
  | Not -> "\\textasciitilde"

and pretty_print_indices q = String.concat ", " (List.map pretty_print_expr q)

and pretty_print_bind = function
    InferredBind(_, s, e) -> "\\hbox{\\textcolor{violet}{\\bf let}  \\textcolor{BurntOrange}{" ^ s ^ "} = }" ^ pretty_print_expr e
  | TensorBind(_, s, q) -> "\\hbox{\\textcolor{violet}{\\bf let} " ^ s ^ "}_{" ^ pretty_print_indices q ^ "}"
  | TensorBindAndAssign(_, s, q, e) -> 
        "\\hbox{\\textcolor{violet}{\\bf let} " ^ s ^ "}_{" ^ pretty_print_indices q ^ "} = " ^ pretty_print_expr e

and pretty_print_comment = function
    [] -> ""
  | hd :: tl when tl != [] -> "\\hbox{\\textcolor{Purple}{" ^ snd hd ^ "}}" ^ "\\\\\n" ^ pretty_print_comment tl
  | hd :: tl -> "\\hbox{\\textcolor{Purple}{" ^ snd hd ^ "}}" ^ "\n" ^ pretty_print_comment tl

and pretty_print_expr = function
    Id(_, s) -> "\\hbox{" ^ s ^ "} "
  | Binop(_, e1, op, e2) -> pretty_print_binop (e1, op, e2)
  | Unop(_, uop, e) -> "(" ^ pretty_print_uop uop ^ pretty_print_expr e ^ ")"
  | Assign(_, var, e) -> var ^ " = " ^ pretty_print_expr e
  | Iliteral(_, i) -> string_of_int i
  | Sliteral(_, s) -> "\\hbox{\\textcolor{LimeGreen}{``" ^ s ^ "\"}}"
  | Fliteral(_, f) -> string_of_float f
  | BindExpr(_, b) -> pretty_print_bind b
  | Noexpr -> ""
  | Call(_,fn, el) -> "\\hbox{" ^ fn ^ "}(" ^ 
                    String.concat ", " (List.map pretty_print_expr el) ^ ")"
  | TensorId(_, s, q) ->  s ^ "_{" ^ pretty_print_indices q ^ "}"
  | TensorAssign(_, s, sl, e) -> s ^ "}_{" ^ pretty_print_indices sl ^ "} = " ^ pretty_print_expr e
  | Tensor(_, t) -> pretty_print_tensor t

and pretty_print_line l = String.concat " & " (List.map pretty_print_expr l)

and pretty_print_matrix m = String.concat " \\\\ " (List.map (fun n -> (match n with Tensor(_, t) -> pretty_print_line t | _ -> "")) m) 

and pretty_print_tensor t = match List.hd t with 
    Tensor(_, st) -> (match List.hd st with 
         Tensor(_, _) -> "\\left[" ^ String.concat ", " (List.map pretty_print_btensor t) ^ "\\right]"
        |_ -> "\\begin{pmatrix} " ^ pretty_print_matrix t ^ "\\end{pmatrix}")
   | _ -> "\\begin{pmatrix} " ^ String.concat " & " (List.map string_of_expr t) ^ "\\end{pmatrix}"

and pretty_print_btensor t = match t with
    Tensor(_, st) -> "\\left[" ^ String.concat ", " (List.map pretty_print_btensor st)
 ^ "\\right]"
   |_ -> pretty_print_expr t


let rec pretty_print_stmt = function
    Block(comm, stmts) -> pretty_print_comment comm ^ "\\\\" ^
                    "\\{ \n \\begin{adjustwidth}{2.5em}{0pt} \n" ^
                   String.concat "" (List.map pretty_print_stmt stmts) ^ 
                   "\\end{adjustwidth} \\} \\\\ \n"
  | Expr(comm, e) -> pretty_print_comment comm ^  "\\\\" ^
                    "$" ^ pretty_print_expr e ^ "$;\\\\ \n"
  | Return(comm, e) -> pretty_print_comment comm ^ "\\\\" ^
                  "\\textcolor{Aquamarine}{\\bf return} $" ^ 
                  pretty_print_expr e ^ "$;\n"
  | If(comm, e, s, Block(_, [])) -> pretty_print_comment comm ^ "\\\\" ^
                          "if ($" ^ pretty_print_expr e ^ "$)" ^ 
                           pretty_print_stmt s
  | If(comm, e, s1, s2) -> pretty_print_comment comm ^ "\\\\" ^
                     "if ($" ^ pretty_print_expr e ^ "$) " ^ 
                     pretty_print_stmt s1 ^ " else " ^ pretty_print_stmt s2
  | For(comm, e1, e2, e3, s) -> pretty_print_comment comm ^ "\\\\" ^
                          "for ($" ^ pretty_print_expr e1 ^ "$ ; $" ^ 
                          pretty_print_expr e2 ^ "$ ; $" ^ pretty_print_expr e3 ^ 
                          "$) " ^ pretty_print_stmt s
  | While(comm, e, s) -> pretty_print_comment comm ^  "\\\\" ^
                        "while ($" ^ pretty_print_expr e ^ "$) " ^ pretty_print_stmt s

let pretty_print_typ = function
    Void -> "void"
  | Float -> "float"
  | Int -> "int"
  | String -> "string"
  | TensorType(q) -> string_of_int (List.length q) ^ "-tensor"

let pretty_print_fdecl fdecl =
  pretty_print_comment fdecl.comm ^ "\n" ^
  "\\textcolor{blue}{\\bf " ^ pretty_print_typ fdecl.typ ^ "} " ^
  "\\textcolor{BlueViolet}{ \\bf " ^ fdecl.fname ^ "}(" ^ String.concat ", " (List.map snd fdecl.formals) ^
  ") \n" ^
  pretty_print_stmt fdecl.body

let pretty_print_glob g = "$" ^ pretty_print_bind g ^ "$;\\\\ \n"

let pretty_program (globals, functions) =
     "\\documentclass{article} \n" ^
     "\\usepackage{fullpage} \n" ^
     "\\usepackage{listings} \n" ^
     "\\usepackage{color} \n" ^
     "\\usepackage{amsmath} \n" ^
     "\\usepackage[dvipsnames]{xcolor} \n" ^
     "\\usepackage{changepage} \n" ^
     "\\parindent 0pt \n" ^
     "\\renewcommand{\\baselinestretch}{1.0} \n \n " ^
     "\\setlength\\delimitershortfall{0pt} \n" ^
     "\\setlength\\delimiterfactor{1200} \n \n" ^
     "\\begin{document} \n" ^
     String.concat "" (List.map pretty_print_glob globals) ^
     String.concat "\\\\ \n" (List.map pretty_print_fdecl functions) ^
     "\\end{document}"
