module Eval where

import Parser
import Research

import Data.Map  (Map)

import qualified Data.Map as Map
import qualified Data.Set as Set

evalSingle :: Env -> TricuAST -> Env
evalSingle env term
  | SFunc name [] body <- term =
      if
        | Map.member name env -> 
            errorWithoutStackTrace $ 
              "Error: Identifier '" ++ name ++ "' is already defined."
        | otherwise -> 
            let res = evalAST env body
            in Map.insert "__result" res (Map.insert name res env)
  | SApp func arg <- term =
      let res = apply (evalAST env func) (evalAST env arg)
      in Map.insert "__result" res env
  | SVar name <- term =
      case Map.lookup name env of
        Just v  -> Map.insert "__result" v env
        Nothing -> errorWithoutStackTrace $ "Variable " ++ name ++ " not defined"
  | otherwise =
      Map.insert "__result" (evalAST env term) env

evalTricu :: Env -> [TricuAST] -> Env
evalTricu env []     = env
evalTricu env [x]    =
  let updatedEnv = evalSingle env x
  in Map.insert "__result" (result updatedEnv) updatedEnv
evalTricu env (x:xs) =
  evalTricu (evalSingle env x) xs

evalAST :: Env -> TricuAST -> T
evalAST env term
  | SLambda _ _ <- term = evalAST env (elimLambda term)
  | SVar   name <- term = evalVar name
  | TLeaf       <- term = Leaf
  | TStem  t    <- term = Stem (evalAST env t)
  | TFork  t u  <- term = Fork (evalAST env t) (evalAST env u)
  | SApp   t u  <- term = apply (evalAST env t) (evalAST env u)
  | SStr   s    <- term = ofString s
  | SInt   n    <- term = ofNumber n
  | SList  xs   <- term = ofList (map (evalAST env) xs)
  | SEmpty      <- term = Leaf
  | otherwise           = errorWithoutStackTrace "Unexpected AST term"
    where
      evalVar name = Map.findWithDefault
        (errorWithoutStackTrace $ "Variable " ++ name ++ " not defined")
        name env

-- https://github.com/barry-jay-personal/typed_tree_calculus/blob/main/typed_program_analysis.pdf
-- Chapter 4: Lambda-Abstraction
elimLambda :: TricuAST -> TricuAST
elimLambda = go
  where
    go (SLambda (v:vs) body)
      | null vs              = toSKI v (elimLambda body)
      | otherwise            = elimLambda (SLambda [v] (SLambda vs body))
    go (SApp f g)            = SApp (elimLambda f) (elimLambda g)
    go x                     = x

    toSKI x (SVar y)
      | x == y           = _I
      | otherwise        = SApp _K (SVar y)
    toSKI x t@(SApp n u)
      | not (isFree x t) = SApp _K t
      | otherwise        = SApp (SApp _S (toSKI x n)) (toSKI x u)
    toSKI x t
      | not (isFree x t) = SApp _K t
      | otherwise        = SApp (SApp _S (toSKI x t)) TLeaf

    _S = parseSingle "t (t (t t t)) t"
    _K = parseSingle "t t"
    _I = parseSingle "t (t (t t)) t"
    
    isFree x = Set.member x . freeVars
    freeVars (SVar    v    ) = Set.singleton v
    freeVars (SInt    _    ) = Set.empty
    freeVars (SStr    _    ) = Set.empty
    freeVars (SList   s    ) = foldMap freeVars s
    freeVars (SApp    f a  ) = freeVars f <> freeVars a
    freeVars (TLeaf        ) = Set.empty
    freeVars (SFunc   _ _ b) = freeVars b
    freeVars (TStem   t    ) = freeVars t
    freeVars (TFork   l r  ) = freeVars l <> freeVars r
    freeVars (SLambda v b  ) = foldr Set.delete (freeVars b) v

result :: Env -> T
result r = case Map.lookup "__result" r of
  Just a -> a
  Nothing -> errorWithoutStackTrace "No __result field found in provided environment"