module Let where

infixl :$:
infixl :$$:

-- We use a representation that differentiates between free and bound variables.
-- A valid Term should not contain any dangling bound variables, i.e. no Bound i should have less
-- than (i + 1) lambda abstractions surrounding it.
-- We define application using an infix constructor :$: for convenience.
-- Like application the constructor is left-associative.
-- The lambda abstraction takes a String such that we don't have to come
-- up with names for the bound variables.
data Term = Free String | Bound Int | Term :$: Term | Lam String Term

-- We use the semantics that Let binds a *free* variable. Thus, binding for Let and Lam are independent.
-- Using 'Data Types á la carte' we could add the Let constructor to Term in a more
-- elegant way but I digress.
data LTerm =
  LFree String | LBound Int | LTerm :$$: LTerm | LLam String LTerm |
  Let String LTerm LTerm

instance Show LTerm where
  show = show' []
    where
      show' _ (LFree s) = "(F " ++ s ++ ")"
      show' bns (LBound i) = "(B " ++ bns !! i ++ ")"
      show' bns (t1 :$$: t2@(_ :$$: _)) = show' bns t1 ++ " (" ++ show' bns t2 ++ ")"
      show' bns (t1 :$$: t2) = show' bns t1 ++ " " ++ show' bns t2
      show' bns (LLam bn t) = "(λ" ++ bn ++ ". " ++ show' (bn : bns) t ++ ")"
      show' bns (Let v t1 t2) = "(let " ++ v ++ " = " ++ show' bns t1 ++ " in " ++ show' bns t2 ++ ")" 

fromTerm :: Term -> LTerm
fromTerm (Free v) = LFree v
fromTerm (Bound i) = LBound i
fromTerm (t1 :$: t2) = fromTerm t1 :$$: fromTerm t2
fromTerm (Lam n t) = LLam n (fromTerm t)

instance Eq Term where
  (==) (Free v) (Free w) = v == w
  (==) (Bound i) (Bound j) = i == j
  (==) (t1 :$: t2) (t1' :$: t2') = t1 == t1' && t2 == t2'
  -- Since the String argument for Lam is only a name hint, we can ignore it.
  (==) (Lam _ t) (Lam _ t') = t == t' 
  (==) _ _ = False

instance Show Term where
  show = show . fromTerm

-- lift and subst are defined like exactly like in the tutorial.
lift :: Int -> Term -> Term
lift _ (Free v) = Free v
lift l (Bound i) = Bound $ if i < l then i else i + 1
lift l (t1 :$: t2) = lift l t1 :$: lift l t2
lift l (Lam n t) = Lam n $ lift (l + 1) t

subst :: Term -> Int -> Term -> Term
subst _ _ (Free v) = Free v
subst s j (Bound i)
  | i < j = Bound i
  | i == j = s
  | otherwise = Bound $ i - 1
subst s j (t1 :$: t2) = subst s j t1 :$: subst s j t2
subst s j (Lam n t) = Lam n $ subst (lift 0 s) (j + 1) t

-- This converts a free variable into the first Bound variable.
-- Since we are only working on closed terms, 
-- it is safe to use the first bound variable (it can't occur in the term beforehand).
bind :: String -> Int -> Term -> Term
bind v i (Free w) = if v == w then Bound i else Free w
bind _ _ (Bound i) = Bound i
bind v i (t1 :$: t2) = bind v i t1 :$: bind v i t2
bind v i (Lam n t) = Lam n (bind v (i + 1) t)

-- Putting everything together.
expandLets :: LTerm -> Term
expandLets (LFree s) = Free s
expandLets (LBound i) = Bound i
expandLets (t1 :$$: t2) = expandLets t1 :$: expandLets t2
expandLets (LLam n t) = Lam n (expandLets t)
-- As per the task description it should hold that:
--   (let v = t1 in t2) = (λv. t2) t1
-- and (λv. t2) t1 is equal to t2[t1/v].
-- Since we need to replace a free variable, we convert it to a bound variable
-- with bind before performing the actual substitution.
expandLets (Let v t1 t2) = subst (expandLets t1) 0 (bind v 0 $ expandLets t2)


-- Terms and Tests
lfa, lfb, lfc, lfd, lfe :: LTerm
[lfa, lfb, lfc, lfd, lfe] = map (LFree . pure) ['a'..'e']

fa, fb, fc, fd, fe :: Term 
[fa, fb, fc, fd, fe] = map (Free . pure) ['a'..'e']

tt1 = expandLets (Let "a" lfb (lfc :$$: lfd)) == fc :$: fd
tt11 = expandLets (Let "a" lfb (LLam "a" $ LBound 0)) == Lam "a" (Bound 0)
tt2 = expandLets (Let "a" lfe (lfc :$$: lfa :$$: lfb)) == fc :$: fe :$: fb
tt3 = expandLets (Let "a" (Let "b" (LLam "x" (LBound 0)) (lfb :$$: lfb)) (lfa :$$: lfb))
      == let tid = Lam "x" (Bound 0) in tid :$: tid :$: fb 
tt4 = expandLets (Let "a" (LLam "x" $ lfe) $ Let "b" (lfa :$$: lfb) (lfa :$$: lfb :$$: lfc))
      == Lam "x" fe :$: (Lam "x" fe :$: fb) :$: fc