(*
File: CODEGEN.ML
Description: Generates LLVM IR
*)

module L = Llvm
module I = Int64
open Ast

module StringMap = Map.Make(String)


(* SECTION 1: ast.program -> Llvm.module *)
let translate functions =

  let image_row_size = 50
  and image_col_size = 50 in

(* SECTION 2: TYPES *)
(* SECTION 2a: Defining LLVM Types *)
  let context      = L.global_context ()
  in let the_module   = L.create_module context "VSCOde"
  and double_t     = L.double_type  context
  and i32_t        = L.i32_type     context
  and i8_t         = L.i8_type      context
  and i1_t         = L.i1_type      context
  and void_t       = L.void_type    context
  and array_t      = L.array_type
  in let str_t        = L.pointer_type i8_t
  and print_t      = L.var_arg_function_type i32_t [| L.pointer_type i8_t |]
  in let image_t      = L.named_struct_type context "image_t"
    in L.struct_set_body image_t [| (array_t (array_t double_t image_col_size) image_row_size );
                                    (array_t (array_t double_t image_col_size) image_row_size );
                                    (array_t (array_t double_t image_col_size) image_row_size ) |] false;


(* SECTION 2b: AST typ to LLVM typ conversion *)
  let ltype_of_typ = function
      Int               -> i32_t
    | Bool              -> i1_t
    | Void              -> void_t
    | Double            -> double_t
    | String            -> str_t
    | Matrix            -> void_t (* this should never be called because we calculate matrix size and match DimMatrix. Is void_t ok?*)
    | DimMatrix (r, c)  -> (array_t (array_t double_t c) r )
    | Image   ->           image_t
  in


(* SECTION 3: FUNCTION/BLOCKS *)
(* SECTION 3a: Making map of all of the user functions
key: function name
val: tuple: 1. ll.value which can be a function, global var, constant, etc
            2. func decl *)

    (* Built our built in functions and add then to build_in_function_decls,
    which is used as the starting point for function_decls *)
    let built_in_function_decls : (L.llvalue * func_decl) StringMap.t =

      (* Gets the row size (in pixels) of a image *)
      let row_size_cpp_t = L.function_type (i32_t) [| str_t |]
      in let row_size_cpp_func = L.declare_function "row_size_cpp" row_size_cpp_t the_module
      in let row_size_func_decl =
        { typ = Int;
          fname = "row_size";
          formals = [(String, "filename")];
          body = [] }
      in let row_size_func = L.define_function "row_size" (L.function_type i32_t [| str_t |]) the_module

      in let row_size_func_body function_ptr =
        let builder = ref (L.builder_at_end context (L.entry_block function_ptr))
        in let path = List.hd (Array.to_list (L.params function_ptr))
        in let row_size_return = L.build_call row_size_cpp_func [| path |] "row_size_ret" !builder
        in ignore(L.build_ret row_size_return !builder)
      in ignore (row_size_func_body row_size_func);

      StringMap.add "row_size" (row_size_func, row_size_func_decl) StringMap.empty
      in


 (* Now add all the user defined functions to function_decls *)
  let function_decls : (L.llvalue * func_decl) StringMap.t =
    let function_decl m fdecl =
      let name = fdecl.fname
      and formal_types = Array.of_list (List.map (fun (t,_) -> ltype_of_typ t) fdecl.formals)
      in let ftype = L.function_type (ltype_of_typ fdecl.typ) formal_types
      in StringMap.add name (L.define_function name ftype the_module, fdecl) m
    in List.fold_left function_decl built_in_function_decls functions
  in


(* SECTION 3b: Adds a return statement to the end of every block *)
  let add_terminal builder_and_3maps_tuple instr =
    match L.block_terminator (L.insertion_block (fst builder_and_3maps_tuple)) with
        Some _ -> ()
      | None -> ignore (instr (fst builder_and_3maps_tuple))
  in

(* SECTION 4: PRINT *)
(* SECTION 4a: Formating strings for printing *)
  let str_format_str builder = L.build_global_stringptr "%s\n" "fmt" builder
  in let int_format_str builder = L.build_global_stringptr "%d\n" "fmt" builder
  in let double_format_str builder = L.build_global_stringptr "%.3f\n" "fmt" builder
  in let mat_format_str builder row column =
    let make_str r c =
      let rec row str orig_c r' c' =
        if (r' = 0) then str
        else
            if (c' = 0) then let new_str  = str ^ "\n"
              in row new_str orig_c (r'-1) orig_c
            else let new_str = str ^ "%.3f "
              in row new_str orig_c r' (c'-1)
      in row "" c r c
    in let mat_str = make_str row column
  in L.build_global_stringptr mat_str "fmt" builder
  in

(* SECTION 4b: Calling C's print function *)
let print_func : L.llvalue =
  L.declare_function "printf" print_t the_module
  in let print_by_type ltyp e' builder mat_dim_map img_dim_map m_row m_col =
    (match ltyp with
        "i32" ->
          ((L.build_call print_func [| int_format_str builder ; e' |]
          "printf" builder), (mat_dim_map, img_dim_map))
      | "i1" -> (* prints 1 for true, 0 for false *)
          ((L.build_call print_func [| int_format_str builder ; e' |]
          "printf" builder), (mat_dim_map, img_dim_map))
      | "double" ->
          ((L.build_call print_func [| double_format_str builder ; e' |]
          "printf" builder), (mat_dim_map, img_dim_map))
      | "i8*" ->  (* string *)
          ((L.build_call print_func [| str_format_str builder ; e' |]
          "printf" builder), (mat_dim_map, img_dim_map))
      | "void" -> raise (Failure ("Can't print void type"))
      | _ ->  (* must be matrix *)
          ((L.build_call print_func [| (mat_format_str builder m_row m_col) ; e' |]
          "printf" builder), (mat_dim_map, img_dim_map)))
  in

(* SECTION 5: MATRICES *)
(* SECTION 5a: General Matrix Helper Functions*)
  (* SECTION 5a i: Checks if an id is a matrix *)
    let check_if_matrix e_typ =
      (match e_typ with
          "i32" -> false
        | "i1" -> false
        | "void" -> false
        | "double" -> false
        | "i8*" -> false (* string *)
        | _ -> true (* now it should be a matrix... *)
      )

  (* SECTION 5a ii: Gets Matrix Dimensions *)
    and get_matrix_dim m =
      let rec len_check = function
          []                  -> true (* empty list of lists *)
        | _ :: []             -> true (* only one list *)
        | fst :: snd :: []    -> List.length fst = List.length snd
        | fst :: snd :: tail  -> len_check (fst::[snd]) && len_check (snd::tail)
              in if len_check m
                  then let rows = (List.length m)
                  and cols = (List.length (List.hd m))
                    in if rows = 1 && cols = 0 then (0,0) (* An "empty" matrix has 1 row (empty) *)
                    else (rows, cols)
              else raise (Failure ("Not all rows in matrix are the same length"))

  (* SECTION 5a iii: Accesses element at given i, j position *)
    and build_matrix_access m row_index column_index builder locals_map mat_dim_map =
     (try let value = StringMap.find m locals_map
       in (let (r, c) = StringMap.find m mat_dim_map
        in if L.int64_of_const row_index < Some (I.of_int r)
          && L.int64_of_const column_index < Some (I.of_int c)
          && L.int64_of_const row_index >= Some (I.of_int 0)
          && L.int64_of_const column_index >= Some (I.of_int 0)
        then L.build_load (L.build_gep (value) [| L.const_int i32_t 0; row_index;
            column_index |] m builder) m builder
        else raise (Failure ("Index out of matrix bounds")))
      with Not_found -> raise (Failure ("Variable not found 7: " ^ m)))
    in


(* SECTION 5c: Matrix Binop (called within expr) *)
  let build_binop_op op = (match op with
      Add -> L.build_fadd
    | Sub -> L.build_fsub
    | Mult -> L.build_fmul
    | _ -> raise (Failure "Invalid matrix binop"))
  in

  (* SECTION 5c i: Matrix * Matrix add/subtract *)
    let binop_mat_sum op builder mat_dim_map img_dim_map locals_map m1 m2 m1_row m1_col =
      let new_mat_dim_map = StringMap.add "binop_result" (m1_row, m1_col) mat_dim_map
      and new_mat = L.build_alloca (ltype_of_typ (DimMatrix(m1_row, m1_col))) "binop_result" builder
      in for i=0 to (m1_row - 1) do
        for j=0 to (m1_col - 1) do
          let elem1' = build_matrix_access m1 (L.const_int i32_t i) (L.const_int i32_t j)
              builder locals_map new_mat_dim_map
          and elem2' = build_matrix_access m2 (L.const_int i32_t i) (L.const_int i32_t j)
              builder locals_map new_mat_dim_map
          in let result_v = (build_binop_op op) elem1' elem2' "tmp" builder
          in let result_p = L.build_gep new_mat [| L.const_int i32_t 0;
              L.const_int i32_t i; L.const_int i32_t j |] "" builder
          in ignore(L.build_store result_v result_p builder);
        done
      done;
      ((L.build_load (L.build_gep new_mat [| L.const_int i32_t 0 |] "binop_result" builder)
          "binop_result" builder), (new_mat_dim_map, img_dim_map))

  (* SECTION 5c ii: Matrix * Matrix mult *)
    and binop_mat_mult builder mat_dim_map img_dim_map locals_map m1 m2 m1_row m1_col m2_col =
      let new_mat_dim_map = StringMap.add "binop_result" (m1_row, m2_col) mat_dim_map
      and new_mat = L.build_alloca (ltype_of_typ (DimMatrix(m1_row, m2_col))) "binop_result" builder
      and tmp_product = L.build_alloca double_t "tmpproduct" builder
      in ignore(L.build_store (L.const_float double_t 0.0) tmp_product builder);
      for i=0 to (m1_row-1) do
        for j=0 to (m2_col-1) do
          ignore(L.build_store (L.const_float double_t 0.0) tmp_product builder);
          for k=0 to (m1_col-1) do
            let m1_float_val = build_matrix_access m1 (L.const_int i32_t i)
                (L.const_int i32_t k) builder locals_map new_mat_dim_map
            and m2_float_val = build_matrix_access m2 (L.const_int i32_t k)
                (L.const_int i32_t j) builder locals_map new_mat_dim_map
            in let product_m1_m2 = L.build_fmul m1_float_val m2_float_val "tmp" builder
            in ignore(L.build_store ( L.build_fadd product_m1_m2
                (L.build_load tmp_product "addtmp" builder) "tmp" builder) tmp_product builder);
          done;
          let new_mat_element = L.build_gep new_mat [| L.const_int i32_t 0;
              L.const_int i32_t i; L.const_int i32_t j |] "tmpmat" builder
          in let tmp_product_val = L.build_load tmp_product "resulttmp" builder
          in ignore(L.build_store tmp_product_val new_mat_element builder);
        done
      done;
      ((L.build_load (L.build_gep new_mat [| L.const_int i32_t 0 |] "binop_result" builder)
          "binop_result" builder), (new_mat_dim_map, img_dim_map))
      in


  (* SECTION 5c iii: Matrix * Scalar mult *)
    let conduct_scalar_mult op builder m1 num mat_dim_map img_dim_map locals_map =
      let (m1_row, m1_col) = StringMap.find m1 mat_dim_map
          in (match op with
               Mult ->
                  let new_mat_dim_map = StringMap.add "binop_result" (m1_row, m1_col) mat_dim_map
                  in let new_mat = L.build_alloca (ltype_of_typ (DimMatrix(m1_row, m1_col)))
                      "binop_result" builder
                  in for i=0 to (m1_row - 1) do
                    for j=0 to (m1_col - 1) do
                      let elem1' = build_matrix_access m1 (L.const_int i32_t i)
                          (L.const_int i32_t j) builder locals_map new_mat_dim_map
                      in let result_v = (build_binop_op op) elem1' num "tmp" builder
                      in let result_p = L.build_gep new_mat [| L.const_int i32_t 0;
                          L.const_int i32_t i; L.const_int i32_t j |] "" builder
                      in ignore(L.build_store result_v result_p builder);
                    done
                  done;
                  ((L.build_load (L.build_gep new_mat [| L.const_int i32_t 0 |]
                        "binop_result" builder) "binop_result" builder), (new_mat_dim_map, img_dim_map))
              | _ -> raise (Failure "Cannot perform scalar/matrix operation other than multiplication")
          ) in

    let binop_dbl_mat_mult op builder mat_dim_map img_dim_map locals_map m1 double_k =
      let num = L.const_float double_t double_k
      in conduct_scalar_mult op builder m1 num mat_dim_map img_dim_map locals_map
    and binop_dblid_mat_mult op builder mat_dim_map img_dim_map locals_map m1 double_id =
      let find_lvalue_of_id name =
        (try let value = StringMap.find name locals_map
          in L.build_load value name builder
        with Not_found -> raise (Failure ("Variable not found 2: " ^ name)))
      in let num = find_lvalue_of_id double_id
      in conduct_scalar_mult op builder m1 num mat_dim_map img_dim_map locals_map

  (* SECTION 5d iv: Loading an Image *)
    in let load_image name load_return mat_dim_map img_dim_map locals_map builder =
      let img = L.build_alloca (ltype_of_typ (Image)) "img" builder
      in let r_ptr = L.build_struct_gep img 0 "r_ptr" builder
      and g_ptr = L.build_struct_gep img 1 "g_ptr" builder
      and b_ptr = L.build_struct_gep img 2 "b_ptr" builder

      in let rec turn_llval_to_list as_llval as_list iter =
        let size = 3 * image_row_size * image_col_size
        in (match iter with
          | n when n = size -> as_list (* if it reaches image size then return *)
          | _ -> let next_element = (L.build_load (L.build_gep load_return
                [|L.const_int i32_t iter|] "element_ptr" builder) "element" builder)
                in
                (* FILL IN MATRICES *)
                  let mat_list_ix = (iter / 3)
                  in let curr_row = (mat_list_ix / image_col_size)
                  in let curr_col = (mat_list_ix mod image_col_size)
                    in if iter mod 3 = 2
                      then (* FILL IN RED MATRIX *) let red_elem_ptr =
                          L.build_gep r_ptr [| L.const_int i32_t 0;
                          L.const_int i32_t curr_row; L.const_int i32_t curr_col |] "red_mat_ptr" builder
                          in ignore(L.build_store next_element red_elem_ptr builder);
                    else if iter mod 3 = 1
                      then (* FILL IN GREEN MATRIX *) let green_elem_ptr =
                          L.build_gep g_ptr [| L.const_int i32_t 0;
                          L.const_int i32_t curr_row; L.const_int i32_t curr_col |] "green_mat_ptr" builder
                      in ignore(L.build_store next_element green_elem_ptr builder);
                    else if iter mod 3 = 0
                      then (* FILL IN BLUE MATRIX *) let blue_elem_ptr =
                          L.build_gep b_ptr [| L.const_int i32_t 0;
                          L.const_int i32_t curr_row; L.const_int i32_t curr_col |] "green_mat_ptr" builder
                      in ignore(L.build_store next_element blue_elem_ptr builder);
                    else raise (Failure "Internal error"); (* END ENTIRE IF BLOCK *)
                  let new_as_list = as_list@[next_element]
                  in (turn_llval_to_list as_llval new_as_list (iter + 1))
        )
        in ignore(turn_llval_to_list load_return [] 0);

        let new_img_dim_map = StringMap.add name (image_row_size, image_col_size) img_dim_map
        in let new_locals_map = StringMap.add name img locals_map

        in (builder, (new_locals_map, (mat_dim_map, new_img_dim_map)))
    in


(* SECTION 6: BINOP EXPR *)
  let binop_scalar typs op e1' e2' builder mat_dim_map img_dim_map =
    (match typs with
        ("double", "double") -> (* 1. DOUBLE *)
          (((match op with
              Add         -> L.build_fadd
            | Sub         -> L.build_fsub
            | Mult        -> L.build_fmul
            | Div         -> L.build_fdiv
            | Equal       -> L.build_fcmp L.Fcmp.Oeq
            | Neq         -> L.build_fcmp L.Fcmp.One
            | Less        -> L.build_fcmp L.Fcmp.Olt
            | Leq         -> L.build_fcmp L.Fcmp.Ole
            | Greater     -> L.build_fcmp L.Fcmp.Ogt
            | Geq         -> L.build_fcmp L.Fcmp.Oge
            | _           -> raise (Failure "internal error: semant should have rejected invalid op on double")
          ) e1' e2' "tmp" builder), (mat_dim_map, img_dim_map))
      | ("i32", "i32") ->     (* 2. INT *)
        (((match op with
              Add         -> L.build_add
            | Sub         -> L.build_sub
            | Mult        -> L.build_mul
            | Div         -> L.build_sdiv
            | Mod         -> L.build_srem
            | Equal       -> L.build_icmp L.Icmp.Eq
            | Neq         -> L.build_icmp L.Icmp.Ne
            | Less        -> L.build_icmp L.Icmp.Slt
            | Leq         -> L.build_icmp L.Icmp.Sle
            | Greater     -> L.build_icmp L.Icmp.Sgt
            | Geq         -> L.build_icmp L.Icmp.Sge
            | _           -> raise (Failure "internal error: semant should have rejected invalid op on int")
          ) e1' e2' "tmp" builder), (mat_dim_map, img_dim_map))
      | ("i1", "i1") ->     (* 3. BOOL *)
        (((match op with
              And         -> L.build_and
            | Or          -> L.build_or
            | Equal       -> L.build_icmp L.Icmp.Eq
            | Neq         -> L.build_icmp L.Icmp.Ne
            | _           -> raise (Failure "internal error: semant should have rejected invalid op on bool")
          ) e1' e2' "tmp" builder), (mat_dim_map, img_dim_map))
      | (_, _) -> raise (Failure ("cannot do binop on different types"))
    )
    in


(* SECTION 7: EXPR *)
(* returns tuple of (LL.value, mat_dim_map) *)
  let rec expr locals_map mat_dim_map img_dim_map builder (e : expr) = match e with
     IntLit i                -> (L.const_int i32_t i, (mat_dim_map, img_dim_map))
    | StrLit i               -> (* this should be a string pointer *)
                                ((L.build_global_stringptr i "str" builder), (mat_dim_map, img_dim_map))
    | DblLit d               -> ((L.const_float double_t d), (mat_dim_map, img_dim_map))
    | BoolLit b              ->  if b then (L.const_int i1_t 1, (mat_dim_map, img_dim_map))
                                 else (L.const_int i1_t 0, (mat_dim_map, img_dim_map))
    | Id name                -> (try let name' = StringMap.find name locals_map
                                  in ((L.build_load name' name builder), (mat_dim_map, img_dim_map))
                                with Not_found -> raise (Failure ("Variable not found 3: " ^ name)))
    | Binop (e1, op, e2)     ->
          (* e1' is the ll value ll const array array*)
          let e1' = fst (expr locals_map mat_dim_map img_dim_map builder e1)
          and e2' = fst (expr locals_map mat_dim_map img_dim_map builder e2)
          in

          let typ1 = L.string_of_lltype (L.type_of e1')
          and typ2 = L.string_of_lltype (L.type_of e2')
          in

          (* BINOP MATRIX CASE 1: ID op ID *)
          if ((check_if_matrix typ1 = true) && (check_if_matrix typ2 = true))
            then (match (e1, e2) with
              (Id m1, Id m2) ->
                let (m1_row, m1_col) = StringMap.find m1 mat_dim_map
                and (m2_row, m2_col) = StringMap.find m2 mat_dim_map
                  in (match op with
                      Add | Sub -> if (m1_row, m1_col) = (m2_row, m2_col)
                        then (binop_mat_sum op builder mat_dim_map img_dim_map locals_map m1 m2 m1_row m1_col)
                        else raise (Failure "Matrices must have same dimensions for binop addition")
                    | Mult -> if m1_col = m2_row
                        then (binop_mat_mult builder mat_dim_map img_dim_map locals_map m1 m2 m1_row m1_col m2_col)
                        else raise (Failure "Matrices do not have valid dimensions for binop multiplication")
                    | _ -> raise (Failure "Can't divide two matrices")
                  )
              | _ -> raise (Failure "only id matrix op id matrix")
             )

          (* BINOP MATRIX CASE 2: Matrix mult scalar *)
          else if (check_if_matrix typ1 = true) then
            (match (e1, e2) with
                (Id m1, DblLit k) -> binop_dbl_mat_mult op builder mat_dim_map img_dim_map locals_map m1 k
              | (Id m1, Id id) -> binop_dblid_mat_mult op builder mat_dim_map img_dim_map locals_map m1 id
              | _ -> raise (Failure "No other types allowed for binop scalar multiplication"))
          else if (check_if_matrix typ2 = true) then
            (match (e1, e2) with
                (DblLit k, Id m1) -> binop_dbl_mat_mult op builder mat_dim_map img_dim_map locals_map m1 k
              | (Id id, Id m1) -> binop_dblid_mat_mult op builder mat_dim_map img_dim_map locals_map m1 id
              | _ -> raise (Failure "No other types allowed for binop scalar multiplication"))

          (* NON-MATRIX SCALAR BINOP *)
          else
            let typs = (typ1, typ2)
            in binop_scalar typs op e1' e2' builder mat_dim_map img_dim_map
    | Unop (op, e)            ->
          let e' = fst (expr locals_map mat_dim_map img_dim_map builder e)
          in let typ = L.string_of_lltype (L.type_of e')
          in (((match op with
            Neg ->
              (match typ with
                  "i32" -> L.build_neg
                | "double" -> L.build_fneg
                | _ -> raise (Failure "internal error: semant should have rejected invalid unop type")
              )
          | Not when typ = "i1" -> L.build_not
          | _ -> raise (Failure "internal error: semant should have rejected illegal unop")
        ) e' "tmp" builder), (mat_dim_map, img_dim_map))
    | Assign (name, e)       ->
          ( match e with
              MatLit m ->   (* x = [1]; Allocate space according to MatLit and
                update the locals_map for name to have the right amount of space*)
                    let (row, col) = get_matrix_dim m
                    in let local_var = L.build_alloca (ltype_of_typ (DimMatrix(row, col))) name builder
                    in let new_locals_map = StringMap.add name local_var locals_map
                    in let e' = fst (expr new_locals_map mat_dim_map img_dim_map builder e)
                    in ignore (L.build_store e' local_var builder);
                    (e', (mat_dim_map, img_dim_map))
              | ImageLit (_, _, _) ->
                    let local_var = L.build_alloca (ltype_of_typ (Image)) name builder
                    in let new_locals_map = StringMap.add name local_var locals_map
                    in let e' = fst (expr new_locals_map mat_dim_map img_dim_map builder e)
                    in ignore(L.build_store e' local_var builder);
                    (local_var, (mat_dim_map, img_dim_map))
              | Id idName ->
                    let name_value = StringMap.find name locals_map (* value is the llvalue *)
                    in let name_typ = L.type_of name_value
                    in if name_typ = (ltype_of_typ (DimMatrix (0, 0)))
                      then (try let _ = StringMap.find idName locals_map
                        and (row, col) = StringMap.find idName mat_dim_map
                        in let local_var = L.build_alloca (ltype_of_typ (DimMatrix(row, col))) name builder
                        in let new_locals_map = StringMap.add name local_var locals_map
                        and new_mat_dim_map = StringMap.add name (row, col) mat_dim_map
                        in let matrix_llarray = fst (expr new_locals_map new_mat_dim_map img_dim_map builder e)
                        in ignore (L.build_store (fst (expr new_locals_map new_mat_dim_map img_dim_map builder e))
                            local_var builder);
                        (matrix_llarray, (mat_dim_map, img_dim_map))
                      with Not_found -> raise (Failure ("Variable not found 4: " ^ idName)))
                    else
                      let value = StringMap.find name locals_map
                      in let e' = fst (expr locals_map mat_dim_map img_dim_map builder e)
                      in ignore(L.build_store e' value builder);
                      (e', (mat_dim_map, img_dim_map))
              | _ -> (* Non matrix type assign *)
                    let value = StringMap.find name locals_map (* value is the llvalue *)
                    in let e' = fst (expr locals_map mat_dim_map img_dim_map builder e)
                    in ignore(L.build_store e' value builder);
                    (e', (mat_dim_map, img_dim_map)) (* store v inside e' *)
          )
    | Noexpr                 -> (L.const_int i32_t 0, (mat_dim_map, img_dim_map))
    | Noassign (_)           -> (L.const_int i32_t 0, (mat_dim_map, img_dim_map))
                                (* Placeholder for when something gets assigned to it *)
    | MatLit m               ->
          let flipped     = List.map List.rev m
          in let lists       = List.map (List.map (fun e -> fst (expr locals_map
                                mat_dim_map img_dim_map builder e))) flipped
          in let listArray   = List.map Array.of_list lists
          in let listArray2  = List.rev (List.map (L.const_array double_t) listArray)
          in let arrayOfArray  = Array.of_list listArray2
          in ((L.const_array (array_t double_t (List.length (List.hd m))) arrayOfArray), (mat_dim_map, img_dim_map))
    | MatAccess (m, r, c)    ->
          let r' = fst (expr locals_map mat_dim_map img_dim_map builder r)
          and c' = fst (expr locals_map mat_dim_map img_dim_map builder c)
          in ((build_matrix_access m r' c' builder locals_map mat_dim_map), (mat_dim_map, img_dim_map))
    | ImageLit (m1, m2, m3)  -> (* m1, m2, and m3 are all IDs of matrices with dimensions (4032,3024) *)
                                (* Allocate space for the struct *)
                                let img = L.build_alloca (ltype_of_typ (Image)) "img" builder
                                (* Get pointers to every element in struct *)
                                in let m1_pointer = L.build_struct_gep img 0 "m1_pointer" builder
                                (* m1 is element ptr to img at idx 0*)
                                and m2_pointer = L.build_struct_gep img 1 "m2_pointer" builder
                                and m3_pointer = L.build_struct_gep img 2 "m3_pointer" builder
                                (* Get the pointer to the DimMat m1, m2, and m3 *)
                                in let m1' = StringMap.find m1 locals_map
                                (* m1' is the pointer to the LLValue of the mat *)
                                and m2' = StringMap.find m2 locals_map
                                and m3' = StringMap.find m3 locals_map
                                (* Get the pointer to the DimMat m1, m2, and m3 *)
                                in let m1_val = L.build_load m1' "m1_val" builder
                                and m2_val = L.build_load m2' "m2_val" builder
                                and m3_val = L.build_load m3' "m3_val" builder
                                (* Store appropriate values into the pointers *)
                                in ignore(L.build_store m1_val m1_pointer builder);
                                ignore(L.build_store m2_val m2_pointer builder) ;
                                ignore(L.build_store m3_val m3_pointer builder) ;
                                ((L.build_load (L.build_gep img [| L.const_int i32_t 0|]
                                    "img" builder ) "img" builder), (mat_dim_map, img_dim_map))
    | ImageRedAccess img_id  -> let img_val = StringMap.find img_id locals_map
                                in let pointer_to_red = L.build_struct_gep img_val 0 "i_red" builder
                                in ((L.build_load pointer_to_red "actual_red" builder), (mat_dim_map, img_dim_map))
    | ImageGreenAccess img_id-> let img_val = StringMap.find img_id locals_map
                                in let pointer_to_green = L.build_struct_gep img_val 1 "i_green" builder
                                in ((L.build_load pointer_to_green "actual_green" builder), (mat_dim_map, img_dim_map))
    | ImageBlueAccess img_id -> let img_val = StringMap.find img_id locals_map
                                in let pointer_to_blue = L.build_struct_gep img_val 2 "i_blue" builder
                                in ((L.build_load pointer_to_blue "actual_blue" builder), (mat_dim_map, img_dim_map))
    | MatrixRowSize m        -> let size = StringMap.find m mat_dim_map
                                in (L.const_int i32_t (fst (size)), (mat_dim_map, img_dim_map))
    | MatrixColSize m        -> let size = StringMap.find m mat_dim_map
                                in (L.const_int i32_t (snd (size)), (mat_dim_map, img_dim_map))
    | Cast (_, _)            -> (L.const_int i32_t 10, (mat_dim_map, img_dim_map)) (* TODO jk  TO REMOVE *)
    | Call ("save", [e])     ->
          let img_struct_alloc = L.build_alloca image_t "tmp_img_alloc" builder
          in let img_struct_val = fst (expr locals_map mat_dim_map img_dim_map builder e)
          in ignore(L.build_store img_struct_val img_struct_alloc builder);

          let pointer_to_red = L.build_struct_gep img_struct_alloc 0 "i_red" builder
          in let _  = L.build_load pointer_to_red "actual_red" builder
          in let pointer_to_blue = L.build_struct_gep img_struct_alloc 2 "i_blue" builder
          in let _ = L.build_load pointer_to_blue "actual_blue" builder
          in let pointer_to_green = L.build_struct_gep img_struct_alloc 1 "i_green" builder
          in let _ = L.build_load pointer_to_green "actual_green" builder

          in let red_mat_ptr = L.build_gep pointer_to_red [| L.const_int i32_t 0; L.const_int i32_t 0 |] "ptr_red" builder
          in let blue_mat_ptr = L.build_gep pointer_to_blue [| L.const_int i32_t 0; L.const_int i32_t 0 |] "ptr_blue" builder
          in let green_mat_ptr = L.build_gep pointer_to_green [| L.const_int i32_t 0; L.const_int i32_t 0 |] "ptr_green" builder

          in let ptr_typ = L.pointer_type (array_t double_t image_row_size)
          in let save_cpp_t = L.function_type void_t [| ptr_typ; ptr_typ; ptr_typ |]
          in let save_cpp_func = L.declare_function "save_cpp" save_cpp_t the_module

          in ignore(L.build_call save_cpp_func [| red_mat_ptr; green_mat_ptr; blue_mat_ptr |] "" builder);
          (L.const_int i32_t 0, (mat_dim_map, img_dim_map)) (* return void *)
    | Call ("print", [e])    ->
          let e' = fst (expr locals_map mat_dim_map img_dim_map builder e)
          in let ltyp = L.string_of_lltype (L.type_of e')
          in if check_if_matrix ltyp then (match e with
              Id m ->
                let (m_row, m_col) = StringMap.find m mat_dim_map
                in  print_by_type ltyp e' builder mat_dim_map img_dim_map m_row m_col
            | _ -> raise (Failure ("Cannot print anonymous matrix")))
          else print_by_type ltyp e' builder mat_dim_map img_dim_map (-1) (-1)
    | Call (fname, act) ->
          (try let (fdef, fdecl) = StringMap.find fname function_decls
            in let actuals = List.rev (List.map fst (List.map
              (expr locals_map mat_dim_map img_dim_map builder) (List.rev act))  )
            in let result = (match fdecl.typ with
                Void -> ""
              | _ -> fname ^ "_result")
              in (try ((L.build_call fdef (Array.of_list actuals) result builder), (mat_dim_map, img_dim_map))
              with _ -> raise (Failure ("Call function failed on " ^ fname)))
          with _ -> raise(Failure ("cannot find function " ^ fname)))
    in


(* SECTION 8: STMT *)
(* Returns (builder, locals_map) tuple *)
  let rec stmt (the_function, fdecl) (builder, (locals_map, (mat_dim_map, img_dim_map))) = function
     Block sl -> List.fold_left (stmt (the_function, fdecl)) (builder, (locals_map, (mat_dim_map, img_dim_map))) sl
    | Expr e -> ignore(fst (expr locals_map mat_dim_map img_dim_map builder e)); (builder, (locals_map, (mat_dim_map, img_dim_map)))
    | Local (typ, name, e) ->
        (* If local assigning a matrix, we need to calculate the matrix dimensions to allocate memory. *)
        (match typ with
            Matrix  ->
              (match e with
                  Noassign (_) ->
                    (* Allocate a 0x0 dimension matrix for a matrix declaration (no size) e.g. matrix m; *)
                      let local_var = L.build_alloca (ltype_of_typ (DimMatrix(0,0))) name builder
                      in let new_locals_map = StringMap.add name local_var locals_map
                      in (builder, (new_locals_map, (mat_dim_map, img_dim_map)))
                | MatLit m ->
                      let (row, col) = get_matrix_dim m
                      in let new_mat_dim_map = StringMap.add name (row, col) mat_dim_map
                      in let mat_alloc = L.build_alloca (ltype_of_typ (DimMatrix(row, col))) name builder
                      in let mat_val = (fst (expr locals_map new_mat_dim_map img_dim_map builder e))
                      in let new_locals_map = StringMap.add name mat_alloc locals_map
                      in ignore(L.build_store mat_val mat_alloc builder);
                      (builder, (new_locals_map, (new_mat_dim_map, img_dim_map)))
                | Id idName ->
                     (try let _ = StringMap.find idName locals_map
                        and (row, col) = StringMap.find idName mat_dim_map
                        in let mat_alloc = L.build_alloca (ltype_of_typ (DimMatrix(row, col))) name builder
                        in let new_locals_map = StringMap.add name mat_alloc locals_map
                        and new_mat_dim_map = StringMap.add name (row, col) mat_dim_map
                        in let mat_val = (fst (expr new_locals_map new_mat_dim_map img_dim_map builder e))
                        in ignore (L.build_store mat_val mat_alloc builder);
                        (builder, (new_locals_map, (new_mat_dim_map, img_dim_map)))
                      with Not_found -> raise (Failure ("Variable not found 5: " ^ idName)))
                | Binop(m1, _, m2) as e ->
                      let assign_mat_to_binop =
                      let (new_matrix, (new_mat_dim_map, _)) = expr locals_map mat_dim_map img_dim_map builder e
                     in (try let (row, col) = StringMap.find "binop_result" new_mat_dim_map
                        in let mat_alloc = L.build_alloca (L.type_of new_matrix) name builder
                        in let new_locals_map = StringMap.add name mat_alloc locals_map
                        and new_mat_dim_map = StringMap.add name (row, col) mat_dim_map
                        in let mat_val = fst (expr new_locals_map new_mat_dim_map img_dim_map builder e)
                        in ignore (L.build_store mat_val mat_alloc builder);
                        (builder, (new_locals_map, (new_mat_dim_map, img_dim_map)))
                        with Not_found -> raise (Failure ("Variable not found 6: " ^ name)))
                      in (match (m1, m2) with
                          (Id _, Id _) -> assign_mat_to_binop
                        | (Id _, DblLit _) -> assign_mat_to_binop
                        | (DblLit _, Id _) -> assign_mat_to_binop
                        | _ -> raise (Failure "Can't perform binop on unnamed matrices")
                    )
                | ImageRedAccess _ | ImageGreenAccess _ | ImageBlueAccess _ ->
                      let (row, col) = (image_row_size,image_col_size)
                      in let new_mat_dim_map = StringMap.add name (row, col) mat_dim_map
                      in let mat_alloc = L.build_alloca (ltype_of_typ (DimMatrix(row, col))) name builder
                      in let mat_val = (fst (expr locals_map new_mat_dim_map img_dim_map builder e))
                      in let new_locals_map = StringMap.add name mat_alloc locals_map
                      in ignore(L.build_store mat_val mat_alloc builder);
                      (builder, (new_locals_map, (new_mat_dim_map, img_dim_map)))
                | Call (fname, act) ->
                     (match fname with
                        "dim" ->
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act)))
                            in let dim_cpp_t = L.function_type (L.pointer_type double_t) [| str_t |]
                            in let dim_cpp_func = L.declare_function "dim_cpp" dim_cpp_t the_module
                            in let path = List.hd actuals
                            in let dim_return = L.build_call dim_cpp_func [| path |] "dim_ret" builder
                            in (* double* *) let row = L.build_load (L.build_gep dim_return
                                [| L.const_int i32_t 0 |] "row_ptr" builder ) "row_val" builder
                            in let col = L.build_load (L.build_gep dim_return [| L.const_int i32_t 1 |]
                                "col_ptr" builder ) "col_val" builder
                            in let mat_alloc = L.build_alloca (ltype_of_typ (DimMatrix(1, 2))) "tmp" builder
                            in let mat_row = L.build_gep mat_alloc [| L.const_int i32_t 0; L.const_int i32_t 0;
                                L.const_int i32_t 0 |] "" builder
                            in let mat_col = L.build_gep mat_alloc [| L.const_int i32_t 0; L.const_int i32_t 0;
                                L.const_int i32_t 1 |] "" builder
                            in let new_locals_map = StringMap.add name mat_alloc locals_map
                            in let new_mat_dim_map = StringMap.add name (1, 2) mat_dim_map
                            in ignore(L.build_store row mat_row builder);
                            ignore(L.build_store col mat_col builder);
                            (builder, (new_locals_map, (new_mat_dim_map, img_dim_map)))
                        | _ ->
                            (try let (fdef, fdecl) = StringMap.find fname function_decls
                                in let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                    img_dim_map builder) (List.rev act))  )
                                in let result = (match fdecl.typ with
                                      Void -> ""
                                    | _ -> fname ^ "_result")
                                in (try let mat_val = ((L.build_call fdef (Array.of_list actuals) result builder))
                                  in let new_mat_dim_map = StringMap.add name (1, 2) mat_dim_map
                                  in let mat_alloc = L.build_alloca (ltype_of_typ (DimMatrix(1, 2))) name builder
                                  in let new_locals_map = StringMap.add name mat_alloc locals_map
                                  in ignore (L.build_store mat_val mat_alloc builder);
                                  (builder, (new_locals_map, (new_mat_dim_map, img_dim_map)))
                                with _ -> raise (Failure ("Image call function failed on " ^ fname)))
                            with _ -> raise(Failure ("Image Cannot find function " ^ fname))))
              | _ -> raise (Failure ("Can't assign non matrix to a matrix (either noassign, matlit, id)")))
          | Image  ->
              (match e with
                  Noassign (_) ->
                      let local_var = L.build_alloca (ltype_of_typ (Image)) name builder
                      in let new_locals_map = StringMap.add name local_var locals_map
                      in (builder, (new_locals_map, (mat_dim_map, img_dim_map)))
                | ImageLit (_, _, _) ->
                      let (row, col) = (image_row_size,image_col_size)
                      in let new_img_dim_map = StringMap.add name (row, col) img_dim_map
                      in let img_alloc = L.build_alloca (ltype_of_typ (Image)) name builder
                      and img_val = (fst (expr locals_map mat_dim_map new_img_dim_map builder e))
                      in let new_locals_map = StringMap.add name img_alloc locals_map
                      in ignore(L.build_store img_val img_alloc builder);
                      (builder, (new_locals_map, (mat_dim_map, new_img_dim_map)))
                | Id idName ->
                      (try let _ = StringMap.find idName locals_map
                        and (row, col) = (image_row_size,image_col_size)
                        in let img_alloc = L.build_alloca (ltype_of_typ (Image)) name builder
                        in let new_locals_map = StringMap.add name img_alloc locals_map
                        and new_img_dim_map = StringMap.add name (row, col) img_dim_map
                        in let image_val = (fst (expr new_locals_map mat_dim_map new_img_dim_map builder e))
                        in ignore (L.build_store image_val img_alloc builder);
                        (builder, (new_locals_map, (mat_dim_map, new_img_dim_map)))
                      with Not_found -> raise (Failure ("Variable not found 1: " ^ idName)))
                | Call (fname, act) ->
                      (match fname with
                        | "edgedetect" ->
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act))  )
                            in let edgedetect_cpp_t = L.function_type (L.pointer_type double_t) [| str_t |]
                            in let edgedetect_cpp_func = L.declare_function "edgedetect_cpp" edgedetect_cpp_t
                                the_module
                            in let path = List.hd actuals
                            in let edgedetect_return = L.build_call edgedetect_cpp_func [| path |]
                                "edgedetect_ret" builder
                            in load_image name edgedetect_return mat_dim_map img_dim_map locals_map builder
                        | "grayscale" ->
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act))  )
                            in let grayscale_cpp_t = L.function_type (L.pointer_type double_t) [| str_t |]
                            in let grayscale_cpp_func = L.declare_function "grayscale_cpp" grayscale_cpp_t the_module
                            in let path = List.hd actuals
                            in let grayscale_return = L.build_call grayscale_cpp_func [| path |]
                                "grayscale_ret" builder
                            in load_image name grayscale_return mat_dim_map img_dim_map locals_map builder
                        | "brighten" ->
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act))  )
                            in let brighten_cpp_t = L.function_type (L.pointer_type double_t) [| str_t |]
                            in let brighten_cpp_func = L.declare_function "brighten_cpp" brighten_cpp_t the_module
                            in let path = List.hd actuals
                            in let brighten_return = L.build_call brighten_cpp_func [| path |] "brighten_ret" builder
                            in load_image name brighten_return mat_dim_map img_dim_map locals_map builder
                        | "blur" ->
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act))  )
                            in let blur_cpp_t = L.function_type (L.pointer_type double_t) [| str_t |]
                            in let blur_cpp_func = L.declare_function "blur_cpp" blur_cpp_t the_module
                            in let path = List.hd actuals
                            in let blur_return = L.build_call blur_cpp_func [| path |] "blur_ret" builder
                            in load_image name blur_return mat_dim_map img_dim_map locals_map builder
                        | "load" ->
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act))  )
                            in let load_cpp_t = L.function_type (L.pointer_type double_t) [| str_t |]
                            in let load_cpp_func = L.declare_function "load_cpp" load_cpp_t the_module
                            in let path = List.hd actuals
                            in let load_return = L.build_call load_cpp_func [| path |] "load_ret" builder
                            in (* double* *) load_image name load_return mat_dim_map img_dim_map locals_map builder
                        | _ -> (try let (fdef, fdecl) = StringMap.find fname function_decls in
                            let actuals = List.rev (List.map fst (List.map (expr locals_map mat_dim_map
                                img_dim_map builder) (List.rev act))  )
                            in let result =
                                (match fdecl.typ with
                                    Void -> ""
                                  | _ -> fname ^ "_result")
                            in (try let image_val = ((L.build_call fdef (Array.of_list actuals) result
                                builder)) in let img_alloc = L.build_alloca (ltype_of_typ (Image)) name builder
                            in let new_locals_map = StringMap.add name img_alloc locals_map
                            in ignore (L.build_store image_val img_alloc builder);
                            (builder, (new_locals_map, (mat_dim_map, img_dim_map)))
                            with _ -> raise (Failure ("Image call function failed on " ^ fname)))
                        with _ -> raise(Failure ("Image Cannot find function " ^ fname))))
                      | _ -> raise (Failure ("Can't assign non image to a image (either noassign, imagelit, id)"))
                  )
            (* Non Matrix/Image*)
          | _       ->
              let local_var = L.build_alloca (ltype_of_typ typ) name builder
              in let new_locals_map = StringMap.add name local_var locals_map
              in ignore (L.build_store (fst (expr new_locals_map mat_dim_map img_dim_map builder e)) local_var builder);
              (builder, (new_locals_map, (mat_dim_map, img_dim_map)))
        )
    | Return e -> ignore(match fdecl.typ with
                      (* Special "return nothing" instr *)
                      Void -> L.build_ret_void builder
                      (* Build return statement *)
                    | _ -> L.build_ret (fst (expr locals_map mat_dim_map img_dim_map builder e)) builder );
                  (builder, (locals_map, (mat_dim_map, img_dim_map)))
    | If (predicate, then_stmt, else_stmt) ->
           let bool_val = fst (expr locals_map mat_dim_map img_dim_map builder predicate) in
           let merge_bb = L.append_block context "merge" the_function in
                 let build_br_merge = L.build_br merge_bb in (* partial function *)

           let then_bb = L.append_block context "then" the_function in
           add_terminal (stmt (the_function, fdecl) ((L.builder_at_end context then_bb), (locals_map,
                (mat_dim_map, img_dim_map))) then_stmt)
             build_br_merge;

           let else_bb = L.append_block context "else" the_function in
           add_terminal (stmt (the_function, fdecl) ((L.builder_at_end context else_bb), (locals_map,
                (mat_dim_map, img_dim_map))) else_stmt)
             build_br_merge;

           ignore(L.build_cond_br bool_val then_bb else_bb builder);
           (L.builder_at_end context merge_bb, (locals_map, (mat_dim_map, img_dim_map)))
    | While (predicate, body) ->
          let pred_bb = L.append_block context "while" the_function in
          ignore(L.build_br pred_bb builder);

          let body_bb = L.append_block context "while_body" the_function in
          add_terminal (stmt (the_function, fdecl) ((L.builder_at_end context body_bb), (locals_map,
                (mat_dim_map, img_dim_map))) body)
            (L.build_br pred_bb);

          let pred_builder = L.builder_at_end context pred_bb in
          let bool_val = fst (expr locals_map mat_dim_map img_dim_map pred_builder predicate) in

          let merge_bb = L.append_block context "merge" the_function in
          ignore(L.build_cond_br bool_val body_bb merge_bb pred_builder);
          (L.builder_at_end context merge_bb, (locals_map, (mat_dim_map, img_dim_map)))
    | For (e1, e2, e3, body) ->
          stmt (the_function, fdecl) (builder, (locals_map, (mat_dim_map, img_dim_map))) ( Block [Expr e1 ;
              While (e2, Block [body ; Expr e3]) ] )

  in


(* SECTION 9: Fill in the body of the given function *)
  let build_function_body fdecl =
    let (the_function, _) = StringMap.find fdecl.fname function_decls
    in let builder_for_stmt = L.builder_at_end context (L.entry_block the_function)


    (* Set formal arguments *)
    in let add_formal m (t, n) p =
      L.set_value_name n p;
    let local = L.build_alloca (ltype_of_typ t) n builder_for_stmt
        in ignore (L.build_store p local builder_for_stmt);
    StringMap.add n local m

    in let formals = List.fold_left2 add_formal StringMap.empty fdecl.formals (* list of tuples (typ, name) *)
        (Array.to_list (L.params the_function)) (* list of the actual parameter values *)

    (* Build the code for each statement in the function. The StringMap.empty as the second value in the second tuple is a StringMap to keep track of matrix sizes *)
    in let builder_and_3maps_tuple = stmt (the_function, fdecl) (builder_for_stmt, (formals, (StringMap.empty,
        StringMap.empty))) (Block fdecl.body)

    (* Add a return if the last block falls off the end *)
    in add_terminal builder_and_3maps_tuple (match fdecl.typ with
        Void -> L.build_ret_void
      | Double -> L.build_ret (L.const_float double_t 0.0)
      | t -> L.build_ret (L.const_int (ltype_of_typ t) 0))
  in


List.iter build_function_body functions; the_module
