module Eval where

import Parser
import Research

import Data.Map (Map)
import Data.List (foldl')

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

evalSingle :: Map String T -> TricuAST -> 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 $ eliminateLambda func) (evalAST env $ eliminateLambda 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

evalTricu :: Map String T -> [TricuAST] -> Map String T
evalTricu env list = evalTricu' env (filter (/= SEmpty) list)
  where
  evalTricu' :: Map String T -> [TricuAST] -> Map String T
  evalTricu' env [] = env
  evalTricu' env [lastLine] =
    let lastLineNoLambda = eliminateLambda lastLine
        updatedEnv = evalSingle env lastLineNoLambda
    in Map.insert "__result" (result updatedEnv) updatedEnv
  evalTricu' env (line:rest) =
    let lineNoLambda = eliminateLambda line
        updatedEnv = evalSingle env lineNoLambda
    in evalTricu updatedEnv rest

evalAST :: Map String T -> TricuAST -> 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 -> ofString str
  SInt num -> ofNumber num
  SList elems -> ofList (map (evalAST Map.empty) elems)
  SEmpty -> Leaf
  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 :: TricuAST -> TricuAST
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

-- https://github.com/barry-jay-personal/typed_tree_calculus/blob/main/typed_program_analysis.pdf
-- Chapter 4: Lambda-Abstraction
lambdaToT :: String -> TricuAST -> TricuAST
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 :: TricuAST -> 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 -> TricuAST -> Bool
isFree x = Set.member x . freeVars

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

-- We need the SKI operators in an unevaluated TricuAST tree form so that we
-- can keep the evaluation functions straightforward
tI :: TricuAST
tI = SApp (SApp TLeaf (SApp TLeaf (SApp TLeaf TLeaf))) TLeaf

tK :: TricuAST
tK = SApp TLeaf TLeaf

tS :: TricuAST
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"