-----------------------------------------------------------------------------
-- |
-- Module      :  CPLSystem
-- Copyright   :  (c) Masahiro Sakai 2004,2009

-- License     :  BSD-style
-- 
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  portable
--
-----------------------------------------------------------------------------

module CPLSystem
    ( VarTable
    , System (..)
    , emptySystem
    , parseExp
    , checkName
    , parseCDT
    , parseDef
    , addCDT
    , letExp
    , simp
    ) where

import CDT
import qualified CDTParser
import qualified ExpParser
import Exp
import qualified AExp
import qualified Simp
import Type
import qualified Typing
import Typing (Typing(..))

-- FIXME
import qualified FE
import qualified Subst

import Control.Monad
import Data.Maybe
import Data.List
import Text.ParserCombinators.Parsec
import qualified Data.Map as Map

type CDTEnv = [CDT]

type VarTable = Map.Map Id ([Id], Exp, FType)

data System
    = System
    { objects  :: !CDTEnv
    , varTable :: !VarTable
    , trace    :: !Bool
    , lastExp  :: !(Maybe Exp)
    }

emptySystem :: System
emptySystem =
    System{ objects  = []
          , varTable = Map.empty
          , trace    = False
          , lastExp  = Nothing
          }

parseExp :: System -> String -> Either String (Int, Typing)
parseExp sys str =
  case parse ExpParser.exp "" str of
    Left err -> fail (show err)
    Right e ->
      case ExpParser.evalExp (objects sys) (arityEnv sys) e of
        Nothing -> Left "invalid expression" -- FIXME
        Just e ->
          Typing.runTI (tiEnv sys) $ do
            t <- Typing.inferType (substIt sys e)
            t <- Typing.appSubst t
            return $ AExp.quantify t

type Def = (Id, [Id], AExp.AExp, FType)

parseDef :: System -> String -> Either String Def
parseDef sys str =
  case parse ExpParser.def "" str of
    Left err -> fail (show err)
    Right (name,ps,e) ->
      case ExpParser.evalExp (objects sys) (Map.union (Map.fromList (zip ps (repeat 0))) (arityEnv sys)) e of
        Nothing -> Left "invalid definition" -- FIXME
        Just e1 ->
          Typing.runTI (tiEnv sys) $ do
            (ts, t) <- Typing.inferType2 ps e1
            ts <- Typing.appSubst ts
            t@(ae :! t') <- Typing.appSubst t
            let vars = Subst.tv (t' : ts)
                s = zip vars [FE.Var i | i<-[0..]]
            return (name, ps, Subst.apply s ae, FType (length vars) (Subst.apply s ts) (Subst.apply s t'))

arityEnv :: System -> Map.Map Id Int
arityEnv sys = if isJust (lastExp sys) then Map.insert "it" 0 env0 else env0
  where env0 = Map.map (\(ps,_,_) -> length ps) (varTable sys)

tiEnv :: System -> Map.Map Id (Either FType Type)
tiEnv sys = Map.map (\(ps, body, t) -> Left t) (varTable sys)

substIt :: System -> Exp -> Exp
substIt sys e = 
  case lastExp sys of
    Just it -> f it e
    Nothing -> e
  where
    f it Identity         = Identity
    f it (Comp a b)       = Comp (f it a) (f it b)
    f it e@(Nat _)        = e
    f it (Fact obj args)  = Fact obj $ map (f it) args
    f it (Funct obj args) = Funct obj $ map (f it) args
    f it (Var "it" [])    = it
    f it (Var v args)     = Var v $ map (f it) args

checkName :: Monad m => System -> String -> m ()
checkName sys name =
  if name `elem` names
    then fail ("\"" ++ name ++ "\" is already used")
    else return ()
  where
    vt = varTable sys
    objs = objects sys
    names = Map.keys vt ++
            map CDT.functName objs ++
            map CDT.factName objs  ++
            concatMap (map CDT.natName . CDT.nats) objs

parseCDT :: Monad m => System -> String -> m CDT.CDT
parseCDT sys src = case parse CDTParser.cdtDecl "" src of
                   Left err   -> fail (show err)
                   Right decl -> CDTParser.evalCDTDecl (objects sys) decl

addCDT :: Monad m => System -> CDT.CDT -> m System
addCDT sys obj =
    if CDT.isComputable obj
    then do checkName sys (CDT.functName obj)
            checkName sys (CDT.factName obj)
            mapM_ (checkName sys . CDT.natName) (CDT.nats obj)
            return sys{ objects = obj : objects sys }
    else fail "not a computable object"

compile :: System -> Exp -> Simp.CompiledExp
compile sys e = Simp.compile env e
  where env = Map.map (\(ps,body,_) -> (ps, body)) (varTable sys)

letExp :: Monad m => System -> Def -> m System
letExp sys (name,ps,e,ftype) = do
  checkName sys name
  return sys{ varTable = Map.insert name (ps, AExp.skelton e, ftype) (varTable sys) }

simp :: System -> Bool -> Exp -> [(Int, Simp.CompiledExp, Simp.CompiledExp)]
simp sys full exp =
    let exp' = compile sys exp
        traces =
            if trace sys
               then Simp.simpWithTrace full exp'
               else [(0, compile sys Identity, Simp.simp full exp')]
    in traces