{-# Language LambdaCase, Strict #-}
module Assignment where

import Chi
import Data.Functor ( (<&>) )
import Control.Monad.Identity (Identity(runIdentity))

-- Task 3
subst :: Variable -> Exp -> Exp -> Exp
subst var e = \case
  Apply e1 e2 -> Apply (subst var e e1) (subst var e e2)
  Lambda x e' -> Lambda x $ if var /= x then subst var e e' else e'
  Var x -> if var == x then e else Var x
  Const c es -> Const c $ map (subst var e) es
  Rec x e' -> Rec x $ if var /= x then subst var e e' else e'
  Case e' bs -> Case (subst var e e') $ map substBr bs
    where
      substBr :: Br -> Br
      substBr (Branch c vs e') = Branch c vs $ if var `notElem` vs then subst var e e' else e'

-- Task 5
eval :: Exp -> Exp
eval = runIdentity . eval'
  where
    eval' :: Exp -> Identity Exp
    eval' = \case
      e@(Apply e1 e2) -> eval' e1 >>= \case
        Lambda x e' -> eval' e2 >>= eval' . flip (subst x) e'
        _ -> error $ "Function was not function in evaluation: " <> show e
      Const c es -> mapM eval' es <&> Const c
      Rec x e -> eval' $ subst x (Rec x e) e
      Case e bs -> eval' e >>= \case
        Const c vs -> lookupBranch c bs >>= \case
          (xs,e') ->
            if length vs /= length xs then error "Not the same amount of arguments in case" else
            eval' $ foldr (uncurry subst) e' (zip xs vs)
        e -> error $ "Non const in case: " <> show e
      x -> pure x
    
    lookupBranch :: Constructor -> [Br] -> Identity ([Variable], Exp)
    lookupBranch c [] = error "No matching branch"
    lookupBranch c ((Branch c' bs e):brs) =
      if c == c'
      then pure (bs,e)
      else lookupBranch c brs


main :: IO ()
main = getLine >>= print . eval . parse