{-# LANGUAGE OverloadedStrings, LambdaCase #-}

-- |
-- This module implements HGMPification of funcons based on Berger et al. (2017)
-- (First-order.)
module Funcons.MetaProgramming where

import Funcons.MSOS
import Funcons.EDSL
import Funcons.Types
import Funcons.Patterns
import Funcons.RunOptions
import Funcons.Simulation

import Data.Text (pack, unpack)
import qualified Data.Map as M
import qualified Data.Set as S

-- | This function implements the ==CT=> relation.
-- Compiling programs to executable funcons terms,
-- removing occurrences of `meta-up`, `meta-down` and `meta-let` and `meta-bound`.
ctRel :: Funcons -> MSOS Funcons
ctRel f = case f of
  FName nm                        -> return f 
  FApp "meta-up" [m]              -> ulRel m
  FApp "meta-down" [m]            -> staticEval m
  FApp "meta-let" [FValue nm,m,n] | isString_ nm -> do 
    v <- evalRel =<< ctRel m
    [menv] <- getInh env_entity 
    let env' = case menv of 
                Map env -> Map (M.insert nm [v] env)
                _       -> Map (M.singleton nm [v])
    withInh env_entity [env'] (ctRel n)
  FApp nm arg                     -> FApp nm <$> mapM ctRel arg
--  FList fs                        -> FList <$> mapM ctRel fs
  FSet fs                         -> FSet <$> mapM ctRel fs
  FMap fs                         -> FMap <$> mapM ctRel fs
  FValue v                        -> return (FValue v)
  _                               -> liftRewrite (sortErr f ("ctRel not defined"))
  where staticEval m = ctRel m >>= evalRel >>= liftRewrite . dlRel  

-- | This function implements the ==UL=> relation.
-- Translating a funcon into its meta-representation
ulRel :: Funcons -> MSOS Funcons
ulRel f = case f of
  FName nm              -> return $ ast_ [string_ (unpack nm)]
  FApp "meta-down" [f]  -> ctRel f
  FApp "meta-up" [m]    -> ulRel m >>= ulRel
  FApp nm fs            -> ast_ . (string_ (unpack nm):) <$> mapM ulRel fs
--  FList fs              -> ast_ . (string_ "list":) <$> mapM ulRel fs
  FSet fs               -> ast_ . (string_ "set":) <$> mapM ulRel fs
  FMap fs               -> ast_ . (string_ "map":) <$> mapM ulRel fs
  FValue v              -> return $ ast_ [type_ (tyOf v), FValue v]
  -- What TODO with type annotations? 
  _                     -> liftRewrite (sortErr f ("ulRel not defined"))

-- | This function implements the ==DL=> relation.
-- Translating a meta-representation of a program into the actual program
dlRel :: Values -> Rewrite Funcons
dlRel v = case v of 
  VMeta tag -> dlRel' tag
  _         -> sortErr (meta_down_ (fvalues [v])) "meta-down not applied to a meta-representation"
 where
    dlRel' :: TaggedSyntax -> Rewrite Funcons 
    dlRel' t = case t of 
      TagType ty lit                  -> return (FValue lit)
      TagName nm []                   -> return $ FName nm
      TagName "set" vs                -> FSet <$> mapM dlRel vs
      TagName "map" vs                -> FMap <$> mapM dlRel vs
      TagName nm vs                   -> FApp nm <$> mapM dlRel vs

evalRel :: Funcons -> MSOS Values
evalRel f = evalFuncons f >>= \case 
  Right [v]   -> return v
  Right vs    -> liftRewrite $ internal "meta evaluation yields a sequence of values"
  Left f'     -> evalRel f'
  where setGlobal f ctxt = ctxt { ereader = (ereader ctxt) { global_fct = f } }

compile :: FunconLibrary -> TypeRelation -> Funcons -> Funcons -> Funcons
compile lib tyenv fenv f = 
  case runSimIO (runMSOS process (cmp_MSOSReader lib tyenv f) cmp_MSOSState) M.empty of
    ((Left ie , _,_), _) -> error ("failed to compile\n" ++ showIException ie)
    ((Right f, _, _), _) -> f  
  where process = do  
          env <- evalRel fenv
          putMut atom_gen_entity (Set S.empty)
          putMut store_entity (Map M.empty)
          withInh env_entity [env] (ctRel f)

env_entity = "environment"
store_entity = "store"
atom_gen_entity = "used-atom-set"
 
translationStep :: ([Funcons] -> Funcons) -> StrictFuncon
translationStep f vs = compstep $ do  fs <- liftRewrite (mapM dlRel vs)
                                      Left <$> ulRel (f fs)

cmp_MSOSReader lib env f = MSOSReader cmp_RewriteReader M.empty M.empty (fread True)
  where cmp_RewriteReader = RewriteReader lib env defaultRunOptions f f  
--cmp_MSOSWriter = MSOSWriter mempty mempty mempty
cmp_MSOSState = MSOSState M.empty M.empty (RewriteState)

library :: FunconLibrary
library = libFromList [
--    ("meta-up", NonStrictFuncon step_meta_up)  -- static funcon
--  , ("meta-down", StrictFuncon step_meta_down) -- static funcon
    ("eval", StrictFuncon step_meta_eval)
  , ("code", NonStrictFuncon step_code)
  , ("ast", StrictFuncon step_ast)
  , ("type-of", StrictFuncon step_ty_of)
  ]

meta_up_ = applyFuncon "meta-up"
meta_down_ = applyFuncon "meta-down"
meta_let_ = applyFuncon "meta-let"

eval_ = applyFuncon "eval"
code_ = applyFuncon "code"
step_meta_eval :: [Values] -> Rewrite Rewritten
step_meta_eval [v] = dlRel v >>= rewriteTo
step_meta_eval fs = sortErr (eval_ (fvalues fs)) "eval not applied to one argument"

step_code :: [Funcons] -> Rewrite Rewritten
step_code [f] = compstep (toStepRes <$> ulRel f)
step_code fs = sortErr (code_ fs) "code not applied to a single term"

ast_ = applyFuncon "ast"
step_ast :: [Values] -> Rewrite Rewritten
step_ast vs@(s:args) | isString_ s = 
  Funcons.Patterns.isInTupleType args [(ASTs, Just StarOp)] >>= \case
    True  -> rewriteTo $ FValue $ VMeta (TagName (pack (unString s)) args)
    False -> sortErr (ast_ (fvalues vs)) "'ast' is not applied to a string and a sequence of meta-representations"
step_ast vs@[ComputationType (Type ty), v] = Funcons.Patterns.isInType v ty >>= \case
  True  -> rewriteTo $ FValue $ VMeta (TagType ty v)
  False -> sortErr (ast_ (fvalues vs)) "type-checking failed during evaluation of 'ast'"
step_ast vs = sortErr (ast_ (fvalues vs)) "'ast' is not applied to a string or not to a type and a member of that type"

type_of_ = applyFuncon "type-of"
step_ty_of :: [Values] -> Rewrite Rewritten
step_ty_of [v] = rewriteTo $ type_ $ tyOf v
step_ty_of vs = sortErr (type_of_ (fvalues vs)) "type-of not applied to a single value"

--TODO perhaps the parser should be extended to recognise "ast-" prefixes
-- which are translated into applications of VMeta