module Eval where

import Parser
import Research

import Data.Map (Map)
import qualified Data.Map as Map
import Data.List (foldl')
import qualified Data.Set as Set

evalSingle :: Map String T -> SaplingAST -> Map String T
evalSingle env term = case term of
  SFunc name [] body ->
    let
        lineNoLambda = eliminateLambda body
        result = evalAST env lineNoLambda
    in Map.insert name result env
  SLambda _ body ->
    let result = evalAST env body
    in Map.insert "__result" result env
  SApp func arg ->
    let result = apply (evalAST env func) (evalAST env arg)
    in Map.insert "__result" result env
  SVar name ->
    case Map.lookup name env of
      Just value -> Map.insert "__result" value env
      Nothing -> error $ "Variable " ++ name ++ " not defined"
  _ ->
    let result = evalAST env term
    in Map.insert "__result" result env

evalSapling :: Map String T -> [SaplingAST] -> Map String T
evalSapling env [] = env
evalSapling env [lastLine] =
    let lastLineNoLambda = eliminateLambda lastLine
        updatedEnv = evalSingle env lastLineNoLambda
    in Map.insert "__result" (result updatedEnv) updatedEnv
evalSapling env (line:rest) =
    let lineNoLambda = eliminateLambda line
        updatedEnv = evalSingle env lineNoLambda
    in evalSapling updatedEnv rest

evalAST :: Map String T -> SaplingAST -> T
evalAST env term = case term of
    SVar name -> case Map.lookup name env of
        Just value -> value
        Nothing -> error $ "Variable " ++ name ++ " not defined"
    TLeaf -> Leaf
    TStem t -> Stem (evalAST env t)
    TFork t1 t2 -> Fork (evalAST env t1) (evalAST env t2)
    SApp t1 t2 -> apply (evalAST env t1) (evalAST env t2)
    SStr str -> toString str
    SInt num -> toNumber num
    SList elems -> toList (map (evalAST Map.empty) elems)
    SFunc name args body ->
        error $ "Unexpected function definition " ++ name
        ++ " in evalAST; define via evalSingle."
    SLambda {} -> error "Internal error: SLambda found in evalAST after elimination."

eliminateLambda :: SaplingAST -> SaplingAST
eliminateLambda (SLambda (v:vs) body)
  | null vs = lambdaToT v (eliminateLambda body)
  | otherwise = eliminateLambda (SLambda [v] (SLambda vs body))
eliminateLambda (SApp f arg) = SApp (eliminateLambda f) (eliminateLambda arg)
eliminateLambda (TStem t) = TStem (eliminateLambda t)
eliminateLambda (TFork l r) = TFork (eliminateLambda l) (eliminateLambda r)
eliminateLambda (SList xs) = SList (map eliminateLambda xs)
eliminateLambda other = other

lambdaToT :: String -> SaplingAST -> SaplingAST
lambdaToT x (SVar y)
  | x == y = tI
lambdaToT x (SVar y)
  | x /= y = SApp tK (SVar y)
lambdaToT x t
  | not (isFree x t) = SApp tK t
lambdaToT x (SApp n u)
  | not (isFree x (SApp n u)) = SApp tK (SApp (eliminateLambda n) (eliminateLambda u))
lambdaToT x (SApp n u) = SApp (SApp tS (lambdaToT x (eliminateLambda n))) (lambdaToT x (eliminateLambda u))
lambdaToT x body
  | not (isFree x body) = SApp tK body
  | otherwise = SApp (SApp tS (lambdaToT x body)) TLeaf

freeVars :: SaplingAST -> Set.Set String
freeVars (SVar v) = Set.singleton v
freeVars (SInt _) = Set.empty
freeVars (SStr _) = Set.empty
freeVars (SList xs) = foldMap freeVars xs
freeVars (SApp f arg) = freeVars f <> freeVars arg
freeVars TLeaf = Set.empty
freeVars (SFunc _ _ b) = freeVars b
freeVars (TStem t) = freeVars t
freeVars (TFork l r) = freeVars l <> freeVars r
freeVars (SLambda vs b) = foldr Set.delete (freeVars b) vs

isFree :: String -> SaplingAST -> Bool
isFree x = Set.member x . freeVars

toAST :: T -> SaplingAST
toAST Leaf = TLeaf
toAST (Stem a) = TStem (toAST a)
toAST (Fork a b) = TFork (toAST a) (toAST b)

tI :: SaplingAST
tI = SApp (SApp TLeaf (SApp TLeaf (SApp TLeaf TLeaf))) TLeaf

tK :: SaplingAST
tK = SApp TLeaf TLeaf

tS :: SaplingAST
tS = SApp (SApp TLeaf (SApp TLeaf (SApp (SApp TLeaf TLeaf) TLeaf))) TLeaf

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