structure Elab =
struct
  structure StringMap = Map(type k = string val cmp = String.compare)
  structure IntMap = Map(type k = int val cmp = Int.compare)

  fun primop (s : string) : Syntax.primop =
    case s of
      "exit" => Syntax.PExit
    | "add" => Syntax.PAdd
    | "sub" => Syntax.PSub
    | "mul" => Syntax.PMul
    | "div" => Syntax.PDiv
    | "read" => Syntax.PRead
    | "write" => Syntax.PWrite
    | "writeErr" => Syntax.PWriteErr
    | _ => raise Fail ("invalid op: " ^ s)

  fun enumerate (l : 'a list) : (int * 'a) list =
  let
    fun go _ [] = []
      | go i (x :: xs) = (i, x) :: go (i + 1) xs
  in go 0 l
  end

  fun hdstls ([] : 'a list list) : ('a list * 'a list list) option = SOME ([], [])
    | hdstls ([] :: _) = NONE
    | hdstls ((x :: xs) :: ls) =
      case hdstls ls of
        NONE => NONE
      | SOME (heads, tails) => SOME (x :: heads, xs :: tails)

  fun transpose ([] : 'a list list) : 'a list list = []
    | transpose l =
      case hdstls l of
        NONE => []
      | SOME (heads, tails) => heads :: transpose tails

  datatype env = Env of { vars: int StringMap.map, types: (int * int) StringMap.map, structTypes: env StringMap.map }

  val emptyEnv = Env { vars = StringMap.empty, types = StringMap.empty, structTypes = StringMap.empty }

  fun bindVar (name : string) (sym : int) (Env env) : env = Env { vars = StringMap.insert name sym (#vars env), types = #types env, structTypes = #structTypes env }

  fun lookupVar (name : string) (Env env) : int =
    case StringMap.lookup name (#vars env) of
      SOME x => x
    | NONE => raise Fail ("unbound identifier " ^ name)

  fun bindDataCons (cons : (string * Syntax.etype option) list) (Env env) : env =
  let
    val nCons = length cons
    val (_, vars, types) =
      foldl
        (fn ((name, _), (i, vars, types)) =>
          (i + 1, StringMap.insert name (Gensym.new ()) vars, StringMap.insert name (i, nCons) types))
        (0, #vars env, #types env)
        cons
  in Env { vars = vars, types = types, structTypes = #structTypes env }
  end

  fun lookupStructType (name : string) (Env env) : env =
    case StringMap.lookup name (#structTypes env) of
      SOME x => x
    | NONE => raise Fail ("unbound structure " ^ name)

  fun lookupCon ([name] : string list) (Env env) : int option =
        Option.map (fn (i, _) => i) (StringMap.lookup name (#types env))
    | lookupCon (structName :: names) env =
        lookupCon names (lookupStructType structName env)
    | lookupCon [] _ = raise Fail "lookupCon empty"

  fun nConstructors ([name] : string list) (Env env) : int =
        let val (_, n) = valOf (StringMap.lookup name (#types env))
        in n
        end
    | nConstructors (structName :: names) env =
        nConstructors names (lookupStructType structName env)
    | nConstructors [] _ = raise Fail "nConstructors empty"

  fun patternMatrix (arms : Syntax.pat list) : Syntax.pat list list =
  let val ts =
    foldl
      (fn (Syntax.PWild, tupleSize) => tupleSize
        | (Syntax.PVar _, tupleSize) => tupleSize
        | (Syntax.PCon _, _) => ~2
        | (Syntax.PInt _, _) => ~2
        | (Syntax.PTuple t, ~1) => length t
        | (Syntax.PTuple t, tupleSize) =>
            if tupleSize = length t
            then tupleSize
            else raise Fail "a type checker would have caught this")
    ~1
    arms
  in if ts < 0
  then [arms]
  else let
    val expandedArms =
      map
        (fn Syntax.PTuple t => t
          | _ => List.tabulate (ts, fn _ => Syntax.PWild))
        arms
    val cols = transpose expandedArms
  in List.concat (map patternMatrix cols)
  end
  end

  fun occurrenceVector (expr : Syntax.lexp) (arms : Syntax.pat list) : Syntax.lexp list =
  let
    val expandedArms =
      List.mapPartial
        (fn Syntax.PTuple t => SOME t
          | _ => NONE)
        arms
    val cols = transpose expandedArms
  in
    case cols of
      [] => [expr]
    | _ => List.concat (map (fn (i, col) => occurrenceVector (Syntax.LSelect (i, expr)) col) (enumerate cols))
  end

  fun patternBindings (expr : Syntax.lexp) (Syntax.PVar v) : (string * Syntax.lexp) list = [(v, expr)]
    | patternBindings expr (Syntax.PTuple t) =
        List.concat (map (fn (i, p) => patternBindings (Syntax.LSelect (i, expr)) p) (enumerate t))
    | patternBindings expr (Syntax.PCon (_, arg)) = patternBindings (Syntax.LSelect (1, expr)) arg
    | patternBindings _ _ = []

  fun swap1 0 (l : 'a list) : 'a list = l
    | swap1 n (first :: rest) =
        (case swap1 (n - 1) rest of
          x :: rest => x :: first :: rest
        | _ => raise Fail "swap1: index out of bounds")
    | swap1 _ _ = raise Fail "swap1: index out of bounds"

  fun swap (n : int) (patterns : Syntax.pat list list) (occurrences : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list =
    (map (swap1 n) patterns, swap1 n occurrences)

  fun specialize (env : env) (n : int) (patterns : Syntax.pat list list, occurrences : Syntax.lexp list, actions : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list * Syntax.lexp list =
    let
      val con1 =
        List.find
          (fn Syntax.PCon (name, _) => valOf (lookupCon name env) = n
            | _ => false)
          (map hd patterns)
      val newTupleSize =
        case con1 of
          SOME (Syntax.PCon (_, Syntax.PTuple t)) => length t
        | _ => 0
      val occHead = hd occurrences
      val occRest = tl occurrences
      val occurrences =
        if newTupleSize = 0
        then Syntax.LSelect (1, occHead) :: occRest
        else
          List.tabulate (newTupleSize, fn i => Syntax.LSelect (i, Syntax.LSelect (1, occHead))) @ occRest
      fun specializeRow (Syntax.PInt i :: rest) =
            if i = n then SOME (Syntax.PWild :: rest) else NONE
        | specializeRow (Syntax.PWild :: rest) = SOME (Syntax.PWild :: rest)
        | specializeRow (Syntax.PVar _ :: rest) = SOME (Syntax.PWild :: rest)
        | specializeRow (Syntax.PCon (con, Syntax.PTuple []) :: rest) =
            specializeRow (Syntax.PCon (con, Syntax.PTuple [Syntax.PWild]) :: rest)
        | specializeRow (Syntax.PCon (con, Syntax.PTuple args) :: rest) =
            if valOf (lookupCon con env) = n
            then SOME (args @ rest)
            else NONE
        | specializeRow (Syntax.PCon (con, obj) :: rest) =
            specializeRow (Syntax.PCon (con, Syntax.PTuple [obj]) :: rest)
        | specializeRow _ = raise Fail "unexpected pattern in the matrix"
      val (patterns, actions) =
        ListPair.unzip
          (List.mapPartial
            (fn (p, a) =>
              Option.map (fn p => (p, a)) (specializeRow p))
            (ListPair.zipEq (patterns, actions)))
    in (patterns, occurrences, actions)
    end

  fun default (patterns : Syntax.pat list list, occurrences : Syntax.lexp list, actions : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list * Syntax.lexp list =
    let
      val (patterns, actions) =
        ListPair.unzip
          (List.filter
            (fn (Syntax.PWild :: _, _) => true
              | (Syntax.PVar _ :: _, _) => true
              | _ => false)
            (ListPair.zipEq (patterns, actions)))
    in (patterns, occurrences, actions)
    end

  (* There's ambiguity between pattern variables and constructors that can only
  * be resolved by looking up each pattern variable in the constructors map *)
  fun fixConstructors (env : env) (expr : Syntax.pat) =
    case expr of
      Syntax.PTuple t => Syntax.PTuple (map (fixConstructors env) t)
    | Syntax.PVar v =>
        if isSome (lookupCon [v] env)
        then Syntax.PCon ([v], Syntax.PTuple [])
        else Syntax.PVar v
    | Syntax.PCon (s, p) => Syntax.PCon (s, fixConstructors env p)
    | _ => expr

  fun compilePatternMatching (env : env) ([] : Syntax.pat list list, _ : Syntax.lexp list, _ : Syntax.lexp list) : Syntax.lexp =
        raise Fail "nonexhaustive match"
    | compilePatternMatching env (patterns as firstRow :: rows, occurrences, actions) =
        let val refutablePattern =
          List.find
            (fn (_, Syntax.PInt _) => true
              | (_, Syntax.PCon _) => true
              | _ => false)
            (enumerate firstRow)
        in
          case refutablePattern of
            NONE => hd actions
          | SOME (i, _) =>
              let
                val (patterns, occurrences) =
                  if i = 0
                  then (patterns, occurrences)
                  else swap i patterns occurrences
                val firstCol = map hd patterns
                val signatures =
                  map (fn (x, _) => x)
                    (IntMap.toList
                      (foldl
                        (fn (Syntax.PInt i, acc) => IntMap.insert i true acc
                          | (Syntax.PCon (c, _), acc) => IntMap.insert (valOf (lookupCon c env)) true acc
                          | (_, acc) => acc)
                        IntMap.empty
                        firstCol))
                val nCons =
                  case List.find (fn (Syntax.PCon _) => true | _ => false) firstCol of
                    SOME (Syntax.PCon (name, _)) => nConstructors name env
                  | _ => ~1
                val defaultCase =
                  if length signatures = nCons
                  then NONE
                  else SOME (compilePatternMatching env (default (patterns, occurrences, actions)))
                val switchOperand =
                  if nCons < 0
                  then hd occurrences
                  else Syntax.LSelect (0, hd occurrences)
              in
                Syntax.LSwitch
                  ( switchOperand
                  , map
                      (fn i => (i, compilePatternMatching env (specialize env i (patterns, occurrences, actions))))
                      signatures
                  , defaultCase
                  )
              end
        end

  fun declBoundVars (Syntax.DVal (p, _)) : string list = map (fn (x, _) => x) (patternBindings (Syntax.LInt 0) p)
    | declBoundVars (Syntax.DValRec (p, _)) = map (fn (x, _) => x) (patternBindings (Syntax.LInt 0) p)
    | declBoundVars (Syntax.DFun (name, _)) = [name]
    | declBoundVars (Syntax.DDatatype (_, cons)) = map (fn (x, _) => x) cons
    | declBoundVars (Syntax.DStruct (name, _)) = [name]

  fun structBoundVars (decls : Syntax.dec list) : string list = List.concatMap declBoundVars decls

  fun bindStructType (name : string) (decls : Syntax.dec list) (Env env) : env =
    let
      val structEnv =
        foldl
          (fn (Syntax.DDatatype (_, cons), env) => bindDataCons cons env
            | (Syntax.DStruct (name, decls), env) => bindStructType name decls env
            | (_, env) => env)
          emptyEnv
          decls
      val structEnv =
        foldl
          (fn ((i, n), env) => bindVar n i env)
          structEnv
          (enumerate (structBoundVars decls))
    in
      Env { vars = #vars env, types = #types env, structTypes = StringMap.insert name structEnv (#structTypes env) }
    end

  fun actionVector (env : env) (expr : Syntax.lexp) (arms : (Syntax.pat * Syntax.expr) list) : Syntax.lexp list =
    map
      (fn (p, body) =>
        let
          val bindings = patternBindings expr p
          val env =
            foldl
              (fn ((name, _), env) =>
                bindVar name (Gensym.new ()) env)
              env
              bindings
        in
          foldl
            (fn ((name, binding), acc) =>
              Syntax.LApp (Syntax.LFn (lookupVar name env, acc), binding))
            (elab env body)
            bindings
        end)
      arms

  and elabCase (env : env) (expr : Syntax.lexp) (arms : (Syntax.pat * Syntax.expr) list) =
    let
      val arms = map (fn (p, e) => (fixConstructors env p, e)) arms
      val patterns = transpose (patternMatrix (map (fn (x, _) => x) arms))
      val occurrences = occurrenceVector expr (map (fn (x, _) => x) arms)
      val actions = actionVector env expr arms
      val actionFns =
        map
          (fn a => (Gensym.new (), Gensym.new (), a))
          actions
      val smallActions =
        map
          (fn (f, _, _) => Syntax.LApp (Syntax.LVar f, Syntax.LInt 0))
          actionFns
    in
      Syntax.LFix
        ( actionFns
        , compilePatternMatching env (patterns, occurrences, smallActions)
        )
    end

  and elab (env : env) (p : Syntax.expr) : Syntax.lexp =
    case p of
      Syntax.EIdent [i] => Syntax.LVar (lookupVar i env)
    | Syntax.EIdent (structName :: accessors) =>
        let
          val s = Syntax.LVar (lookupVar structName env)
          val env = lookupStructType structName env
          fun go env [i] acc = Syntax.LSelect (lookupVar i env, acc)
            | go env (accessor :: accessors) acc =
                let val env = lookupStructType accessor env
                in go env accessors (Syntax.LSelect (lookupVar accessor env, acc))
                end
            | go _ [] _ = raise Fail "go empty"
        in go env accessors s
        end
    | Syntax.EBuiltin builtin => Syntax.LPrim (primop builtin)
    | Syntax.EInt i => Syntax.LInt i
    | Syntax.EStr s => Syntax.LString s
    | Syntax.ETuple exprs => Syntax.LRecord (map (elab env) exprs)
    | Syntax.EList exprs =>
        foldr
          (fn (x, acc) =>
            Syntax.LRecord [elab env x, acc])
          (Syntax.LInt 0)
          exprs
    | Syntax.EApp (f, x) => Syntax.LApp (elab env f, elab env x)
    | Syntax.ETyped (e, _) => elab env e
    | Syntax.EAndAlso (_, _) => raise Fail "unimplemented"
    | Syntax.EOrElse (_, _) => raise Fail "unimplemented"
    | Syntax.ELet ([], body) => elab env body
    | Syntax.ELet (Syntax.DDatatype (name, cons) :: decls, body) =>
        let
          val env = bindDataCons cons env
          val funs =
            List.mapPartial
              (fn (_, (_, NONE)) => NONE
                | (i, (name, _)) =>
                    let val v = Gensym.new ()
                    in SOME (lookupVar name env, v, Syntax.LRecord [Syntax.LInt i, Syntax.LVar v])
                    end)
              (enumerate cons)
          val vals =
            List.mapPartial
              (fn (i, (name, NONE)) => SOME (lookupVar name env, Syntax.LRecord [Syntax.LInt i])
                | _ => NONE)
              (enumerate cons)
        in
          foldl
            (fn ((v, x), acc) =>
              Syntax.LApp (Syntax.LFn (v, acc), x))
            (Syntax.LFix (funs, elab env (Syntax.ELet (decls, body))))
            vals
        end
    | Syntax.ELet (Syntax.DVal (pat, v) :: decls, body) =>
        elab env (Syntax.ECase (v, [(pat, Syntax.ELet (decls, body))]))
    | Syntax.ELet (Syntax.DValRec (Syntax.PVar name, f as Syntax.ELambda _) :: decls, body) =>
        let
          val n = Gensym.new ()
          val env = bindVar name n env
          val (arg, fnBody) =
            case elab env f of
              Syntax.LFn x => x
            | _ => raise Fail "Syntax.ELambda should expand to Syntax.LFn"
        in
          Syntax.LFix ([(n, arg, fnBody)], elab env (Syntax.ELet (decls, body)))
        end
    | Syntax.ELet (Syntax.DValRec _ :: _, _) => raise Fail "invalid val rec"
    | Syntax.ELet (Syntax.DFun (name, cases) :: decls, body) =>
        let
          val (ps1, _) = hd cases
          val nPats = length ps1
        in if not (List.all (fn (ps, _) => length ps = nPats) cases)
        then raise Fail "clauses do not all have same number of patterns"
        else let
          val n = Gensym.new ()
          val temps = List.tabulate (nPats, fn _ => Gensym.new ())
          val env = bindVar name n env
          val t = Gensym.new ()
          val innerCase = elabCase env (Syntax.LVar t) (map (fn (ps, b) => (Syntax.PTuple ps, b)) cases)
        in
          Syntax.LFix
            ( [ ( n
                , hd temps
                , foldr
                    Syntax.LFn
                    (Syntax.LApp (Syntax.LFn (t, innerCase), Syntax.LRecord (map Syntax.LVar temps)))
                    (tl temps)
                )
              ]
            , elab env (Syntax.ELet (decls, body))
            )
        end
        end
    | Syntax.ELet (Syntax.DStruct (name, structDecls) :: decls, body) =>
        let
          val names = structBoundVars structDecls
          val tuple = elab env (Syntax.ELet (structDecls, Syntax.ETuple (map (fn n => Syntax.EIdent [n]) names)))
          val v = Gensym.new ()
          val env = bindStructType name structDecls env
          val env = bindVar name v env
        in Syntax.LApp (Syntax.LFn (v, elab env (Syntax.ELet (decls, body))), tuple)
        end
    | Syntax.ELambda body =>
        let val v = Gensym.new ()
        in Syntax.LFn (v, elabCase env (Syntax.LVar v) [body])
        end
    | Syntax.ECase (expr, arms) =>
        let val v = Gensym.new ()
        in Syntax.LApp (Syntax.LFn (v, elabCase env (Syntax.LVar v) arms), elab env expr)
        end

  fun elaborate (p : Syntax.expr) : Syntax.lexp = elab emptyEnv p
end