{-# LANGUAGE LambdaCase #-}

module Funcons.Substitution (
    Env(..), subsAndRewrite, envInsert, emptyEnv, Levelled(..), substitute,
    rewriteTermTo, stepTermTo,
    ) where

import Funcons.Types
import Funcons.MSOS
import Funcons.Exceptions

import Control.Monad
import Data.Monoid
import qualified Data.Map as M

-- | An environment mapping meta-variables to funcon terms.
-- This environment is used by a substitution procedure to transform
-- funcon terms from 'FTerm' representation to 'Funcons'.
type Env = M.Map MetaVar Levelled 
data Levelled   = ValueTerm Values
                | ValuesTerm [Values]
                | FunconTerm Funcons
                | FunconsTerm [Funcons]

-- | The empty substitution environment.
-- Bindings are inserted by pattern-matching.
emptyEnv = M.empty

envLookup :: Env -> MetaVar -> Rewrite Levelled
envLookup env k = case M.lookup k env of
                    Nothing -> internal ("undefined metavar: " <> k)
                    Just lf -> return lf

envInsert :: MetaVar -> Levelled -> Env -> Env
envInsert = M.insert

fsLevel :: Levelled -> Rewrite [Funcons]
fsLevel = \case FunconsTerm f       -> return f
                FunconTerm f        -> return [f]
                ValuesTerm vs       -> return (map FValue vs)
                ValueTerm v         -> return [FValue v]

fLevel k = \case    FunconsTerm [f]  -> return f
                    FunconsTerm fs 
                     | all isVal fs    -> return (FValue $ safe_tuple_val (map downcastValue fs))
                     | otherwise       -> return (FTuple fs)
                    FunconTerm f        -> return f
                    ValuesTerm vs       -> return (FValue (safe_tuple_val vs))
                    ValueTerm v         -> return (FValue v)


-- substitution
substitute :: FTerm ->  Env -> Rewrite Funcons
substitute (TVar k) env = envLookup env k >>= fLevel k
substitute (TName nm) env = return $ FName nm
substitute (TApp nm term) env = do
    -- if its anything but a sequence-var or a tuple, make it a singleton sequence
    fs <- case term of TVar k   -> envLookup env k >>= fsLevel
                       TTuple l -> substitute term env >>= \case
                                    FTuple fs -> return fs  
                                    _ -> error "subsitute assert"
                       _        -> (:[]) <$> substitute term env
    return $ FApp nm (FTuple fs)
substitute (TTuple terms) env = FTuple <$> subsFlatten terms env 
-- e.g. print(V+:values+) -- standard-out![V+] -> ()
substitute (TList terms) env = FList <$> subsFlatten terms env 
substitute (TSet terms)  env  = FSet <$> subsFlatten terms env
substitute (TMap terms)  env  = FMap <$> mapM (flip substitute env) terms
substitute (TFuncon f) env = return f
substitute (TSortUnion t1 t2) env = FSortUnion <$> substitute t1 env <*> substitute t2 env
substitute (TSortSeq t1 op) env = flip FSortSeq op <$> substitute t1 env
substitute (TSortComputes t1) env = FSortComputes <$> substitute t1 env
substitute (TSortComputesFrom t1 t2) env = 
    FSortComputesFrom <$> substitute t1 env <*> substitute t2 env

-- flatten out sequence-variables
subsFlatten :: [FTerm] -> Env -> Rewrite [Funcons]
subsFlatten terms env = concat <$> (forM terms $ \case 
                            TVar k  -> envLookup env k >>= fsLevel
                            term    -> (:[]) <$> substitute term env)

subsAndRewrite :: FTerm -> Env -> Rewrite Values
subsAndRewrite term env = do
    f <- substitute term env 
    rewriteFuncons f >>= \case
        ValTerm v -> return v
        CompTerm _ _    -> rewrite_throw (SideCondFail "premise evaluation requires step")


-- | Variant of 'rewriteTo' that applies substitution.
rewriteTermTo :: FTerm -> Env -> Rewrite Rewritten
rewriteTermTo fterm env = substitute fterm env >>= rewriteTo

-- | Variant of 'stepTo' that applies substitution.
stepTermTo :: FTerm -> Env -> MSOS Funcons
stepTermTo fterm env = do   
    (liftRewrite $ substitute fterm env) >>= stepTo