module Env = Map.Make (String)

let rec lookup_ty env ty =
  match ty with
  | Type.Int | Bool | Float | Any -> ty
  | Var x -> ( match Env.find_opt x env with Some ty -> ty | None -> ty)
  | Arrow (ts, t) -> Arrow (List.map (lookup_ty env) ts, lookup_ty env t)
  | Product ts -> Product (List.map (lookup_ty env) ts)
  | Vec t -> Vec (lookup_ty env t)
  | Variant cases ->
      Variant
        (List.map
           (fun (Type.Case (name, tys)) ->
             Type.Case (name, List.map (lookup_ty env) tys))
           cases)
  | Rec (name, ty) ->
      let env = Env.remove name env in
      Rec (name, lookup_ty env ty)

let env_of_decls decls =
  List.fold_left
    (fun env decl ->
      match decl with
      | Typed_ast.DType (name, ty) -> Env.add name (lookup_ty env ty) env)
    Env.empty decls

let rec normalize (program : Typed_ast.program) : Typed_ast.program =
  let env = env_of_decls program.decls in
  { program with expr = normalize_e env program.expr }

and normalize_e (type_env : Type.t Env.t) (expr : Typed_ast.e) =
  {
    expr with
    desc =
      (match expr.desc with
      | Const c -> Const c
      | Var x -> Var x
      | Lam (args, e) ->
          let args' =
            List.map (fun (x, ty) -> (x, lookup_ty type_env ty)) args
          in
          Lam (args', normalize_e type_env e)
      | App (e, exprs) ->
          App (normalize_e type_env e, List.map (normalize_e type_env) exprs)
      | Prim1 (op, e) -> Prim1 (op, normalize_e type_env e)
      | Prim2 (op, e1, e2) ->
          Prim2 (op, normalize_e type_env e1, normalize_e type_env e2)
      | If (e1, e2, e3) ->
          If
            ( normalize_e type_env e1,
              normalize_e type_env e2,
              normalize_e type_env e3 )
      | Let (binding, e) ->
          Let (normalize_binding type_env binding, normalize_e type_env e)
      | Letrec (bindings, e) ->
          let bindings' = List.map (normalize_binding type_env) bindings in
          Letrec (bindings', normalize_e type_env e)
      | Tuple es -> Tuple (List.map (normalize_e type_env) es)
      | Get (e, i) -> Get (normalize_e type_env e, i)
      | Vect (e1, e2) -> Vect (normalize_e type_env e1, normalize_e type_env e2)
      | VectGet (e1, e2) ->
          VectGet (normalize_e type_env e1, normalize_e type_env e2)
      | VectSet (e1, e2, e3) ->
          VectSet
            ( normalize_e type_env e1,
              normalize_e type_env e2,
              normalize_e type_env e3 )
      | VectLength e -> VectLength (normalize_e type_env e)
      | Ascribe (e, ty) ->
          Ascribe (normalize_e type_env e, lookup_ty type_env ty)
      | For (x, e1, e2, e3) ->
          For
            ( x,
              normalize_e type_env e1,
              normalize_e type_env e2,
              normalize_e type_env e3 )
      | Variant (name, es, ty) ->
          Variant
            ( name,
              List.map (fun e -> normalize_e type_env e) es,
              lookup_ty type_env ty )
      | Match (e, patterns, ty) ->
          let patterns' =
            List.map
              (fun pattern ->
                Typed_ast.
                  { pattern with body = normalize_e type_env pattern.body })
              patterns
          in
          Match (e, patterns', lookup_ty type_env ty)
      | Seq es -> Seq (List.map (normalize_e type_env) es));
  }

and normalize_binding (type_env : Type.t Env.t) (binding : Typed_ast.binding) =
  {
    binding with
    e = normalize_e type_env binding.e;
    ty = lookup_ty type_env binding.ty;
  }
