content: {-# LANGUAGE
   BangPatterns
   #-}
  module Context where
  
  import Names
  import AST
  import Substitution
  import Data.Monoid
  import Data.Functor
  import qualified Data.Map as M
  import qualified Data.Set as S
  import Data.Map (Map)
  import Data.Set (Set)
  import Control.Monad.State.Strict (StateT, runStateT, modify, get, put)
  import Control.Monad.RWS.Strict (RWST, ask, local, censor, runRWST, get, put,listen)
  import Control.Monad.Trans (lift)
  import Control.Monad.Trans.Cont
  import Choice
  import Data.List
  import Debug.Trace
  
  --------------------
  ---  context map ---
  --------------------
  type ContextMap = Map Name ((Bool,Integer),Type)
  type ContextMapT = Map Name Type
  
  
  --------------------------------
  ---  constraint context list ---
  --------------------------------
  data Binding = Binding { elmQuant :: Quant
                         , elmName :: Name
                         , elmType :: Type
                         , elmPrev :: Maybe Name
                         , elmNext :: Maybe Name
                         } deriving (Show)
                 
  instance Subst Binding where
    substFree sub f b = b { elmType = substFree sub f $! elmType b }
                      
  data Context = Context { ctxtHead :: Maybe Name
                         , ctxtMap  :: Map Name Binding
                         , ctxtTail :: Maybe Name
                         } deriving (Show)
                                    
  instance Subst Context where               
    substFree sub f b = b { ctxtMap = substFree sub f <$> ctxtMap b }
  
  lookupWith s a ctxt = case M.lookup a ctxt of
    Just r -> r
    Nothing -> error s
  
  
  
  emptyContext = Context Nothing mempty Nothing
  
  -- assumes the element is not already in the context, or it is and the only thing that is changing is it's type.
  addToContext :: String -> Context -> Binding -> Context
  addToContext s (Context Nothing ctxt Nothing) elm@(Binding _ nm _ Nothing Nothing) | M.null ctxt = checkContext (s++"\naddToCtxt N N: ") $ 
                                                                                                   Context (Just nm) (M.singleton nm elm) (Just nm)
  addToContext s c (Binding _ _ _ Nothing Nothing) = error $ "context not empty so can't add to tail: "++show c
  addToContext s c@(Context h ctxt t) elm@(Binding _ nm _ t'@(Just p) Nothing) | t' == t = checkContext (s++"\naddToCtxt J N: "++show elm ++ "\n\tOLD CONTEXT: "++show c) $ 
    Context h (M.insert p t'val $ M.insert nm elm $ ctxt) (Just nm)
    where t'val = (lookupWith (s++" looking up p ctxt") p ctxt) { elmNext = Just nm }
  addToContext s _ (Binding _ _ _ _ Nothing) = error "can't add this to tail"
  addToContext s (Context h ctxt t) elm@(Binding _ nm _ Nothing h'@(Just n)) | h' == h = checkContext (s++"\naddToCtxt N J: ") $ 
    Context (Just nm) (M.insert n h'val $ M.insert nm elm $ ctxt) t
    where h'val = (lookupWith (s++" looking up n ctxt") n ctxt) { elmPrev = Just nm }
  addToContext s _ (Binding _ _ _ Nothing _) = error "can't add this to head"
  addToContext s ctxt@Context{ctxtMap = cmap} elm@(Binding _ nm _ (Just p) (Just n)) = checkContext (s++"\naddToCtxt J J: ") $ 
    ctxt { ctxtMap = M.insert n n'val $ M.insert p p'val $ M.insert nm elm $ cmap }
    where n'val = (lookupWith (s++" looking up n cmap") n cmap) { elmPrev = Just nm }
          p'val = (lookupWith (s++" looking up p cmap") p cmap) { elmNext = Just nm }
    
  removeFromContext :: Name -> Context -> Context
  removeFromContext nm ctxt@(Context h cmap t) = case M.lookup nm cmap of
    Nothing -> checkContext "removing: nothing" $ ctxt
    Just Binding{ elmPrev = Nothing, elmNext = Nothing } -> emptyContext
    Just Binding{ elmPrev = Nothing, elmNext = Just n } -> isSane (Just nm == h) $ checkContext "removing: N J" $ Context (Just n) (M.insert n h' $ M.delete nm cmap) t
      where h' = (lookupWith "attempting to find new head" n cmap) { elmPrev = Nothing }
    Just Binding{ elmPrev = Just p, elmNext = Nothing } -> isSane (Just nm == t) $ checkContext "removing: J N" $ Context h (M.insert p t' $ M.delete nm cmap) (Just p)
      where t' = (lookupWith "attempting to find new tail" p cmap) { elmNext = Nothing }
    Just Binding{elmPrev = Just cp, elmNext = Just cn } -> case () of
      _ | h == t -> checkContext "removing: J J | h == t " $ Context Nothing mempty Nothing
      _ | h == Just nm -> checkContext "removing: J J | h == Just nm  " $ Context (Just cn) (n' $ M.delete nm cmap) t
      _ | t == Just nm -> checkContext "removing: J J | t == Just nm  " $ Context h   (p' $ M.delete nm cmap) (Just cp)
      _ -> checkContext ("removing: J J | h /= t \n\t"++show ctxt) $ Context h (n' $ p' $ M.delete nm cmap) t
      where n' = M.insert cn $ (lookupWith "looking up a cmap for n'" cn cmap) { elmPrev = Just cp }
            p' = M.insert cp $ (lookupWith "looking up a cmap for p'" cp cmap ) { elmNext = Just cn }
    where isSane bool a = if bool then a else error "This doesn't match intended binding"
  
  addAfter s t quant nm tp ctxt@Context{ctxtMap = cmap} = checkContext "addAfter" $ do
    let top = lookupWith ("looking up "++t++"\n\t in context: "++show cmap++"\n\t"++s) t cmap
    addToContext ("ADD_AFTER: "++s) ctxt $ Binding quant nm tp (Just $ elmName top) (elmNext top)
  
  addBefore s b quant nm tp ctxt@Context{ctxtMap = cmap} = checkContext "addBefore" $ do
    let bot = lookupWith ("looking up "++b++"\n\t in context: "++show cmap++"\n\t"++s) b cmap
    addToContext ("ADD_BEFORE: "++s) ctxt $ Binding quant nm tp (elmPrev bot) (Just $ elmName bot)
  
  addToHead s quant nm tp ctxt@Context{ctxtMap = cmap} = case M.lookup nm cmap of 
    Nothing -> addToContext ("ATH: "++s) ctxt $ Binding quant nm tp Nothing (ctxtHead ctxt)
    Just (Binding{ elmQuant = quant', elmType = tp'}) | quant' == quant && tp' == tp && quant == Forall -> 
      addToContext ("ATH2: "++s) ctxt' $ Binding quant nm tp Nothing (ctxtHead ctxt')
      where ctxt' = removeFromContext nm ctxt
    _ -> error $ "Can't add to head, already in context: "++show nm++" : "++show tp++"\n@"++show ctxt
    
  addToTail s quant nm tp ctxt@Context{ctxtMap = cmap} = case M.lookup nm cmap of
    Nothing -> addToContext ("ATT: "++s) ctxt $ Binding quant nm tp (ctxtTail ctxt) Nothing
    Just (Binding{ elmQuant = quant', elmType = tp'}) | quant' == quant && tp' == tp && quant == Forall -> ctxt
    _ -> error $ "Can't add to tail, already in context: "++show nm++" : "++show tp++"\n@"++show ctxt
  
  removeHead ctxt = case ctxtHead ctxt of 
    Nothing -> ctxt
    Just a -> removeFromContext a ctxt
  
  removeTail ctxt = case ctxtTail ctxt of 
    Nothing -> ctxt
    Just a -> removeFromContext a ctxt
  
  getTail (Context _ ctx (Just t)) = lookupWith "getting tail" t ctx
  getTail (Context _ _ Nothing) = error "no tail!"
  
  getHead (Context (Just h) ctx _) = lookupWith "getting head" h ctx
  getHead (Context Nothing _ _) = error "no head"
  
  -- gets the list of bindings after (below) a given binding
  getAfter s bind ctx = tail $ getAfterInclusive s bind ctx            
  getAfterInclusive s bind ctx@(Context{ ctxtMap = ctxt }) = gb bind
    where gb ~(Binding quant nm ty _ n) = (quant, (nm,ty)):case n of
            Nothing -> []
            Just n -> gb $ case M.lookup n ctxt of 
              Nothing -> error $ "element "++show n++" not in map \n\twith ctxt: "++show ctx++" \n\t for bind: "++show bind++"\n\t"++s
              Just c -> c
  
  -- gets the list of bindings before (above) a given binding
  getBefore s bind ctx = tail $ getBeforeInclusive s bind ctx            
  getBeforeInclusive s bind ctx@(Context{ ctxtMap = ctxt }) = gb bind
    where gb ~(Binding quant nm ty p _) = (quant, (nm,ty)):case p of
            Nothing -> []
            Just p -> gb $ case M.lookup p ctxt of 
              Nothing -> error $ "element "++show p++" not in map \n\twith ctxt: "++show ctx++" \n\t for bind: "++show bind++"\n\t"++s
              Just c -> c
  
  checkContext _ c = c
  {-
  checkContext _ c@(Context Nothing _ Nothing) = c
  checkContext s ctx = foldr (\v c -> seq (checkEquals v) c) ctx $ zip st (reverse $ ta)
    where st = getBeforeInclusive s (getTail ctx) ctx
          ta = getAfterInclusive s (getHead ctx) ctx
          checkEquals (a,b) | (a == b) = ()
          checkEquals (a,b) = error $ s++" \n\tNOT THE SAME" ++show (a,b) ++ " \n\t IN "++show ctx
  -}
  
  -------------------------
  ---  Traversal Monad  ---
  -------------------------
  data ContextState = ContextState { stateNum :: !Integer
                                   , stateCtxt :: Context 
                                   }
                      
  emptyState = ContextState 0 emptyContext
  
  
  
  instance ValueTracker ContextState where
    putValue i c = c { stateNum = i }
    takeValue c = stateNum c
    
  type Env = RWST ContextMap Constraint ContextState Choice
  
  isolateForFail m = do
    s <- get
    c <- m
    case c of
      Nothing -> do
        put s
        return Nothing
      _ -> return c
  
  ------------------------        
  -- env with a context --        
  ------------------------
  
  
  getElm :: String -> Name -> Env (Either Binding Spine)
  getElm _ !x | isChar x = do
    return $ Right $ var "char"
  getElm s !x = do
    ty <- lookupConstant x
    case ty of
      Nothing -> Left <$> (\ctxt -> lookupWith ("looking up "++x++"\n\t in context: "++show ctxt++"\n\t"++s) x ctxt) <$> ctxtMap <$> stateCtxt <$> get
      Just a -> return $ Right a
    
  getBindings :: Binding -> Env [(Name,Type)]
  getBindings bind = fmap snd <$> getQuantBindings bind
  
  getBindingsInclusive :: Binding -> Env [(Name,Type)]
  getBindingsInclusive bind = fmap snd <$> getQuantBindings bind
  
  -- | This gets all the bindings outside of a given bind and returns them in a list (not including that binding).
  getQuantBindings :: Binding -> Env [(Quant, (Name,Type))]
  getQuantBindings bind = do
    ctx <- stateCtxt <$> get
    return $ getBefore "IN: getQuantBindings" bind ctx
    
  getQuantBindingsInclusive :: Binding -> Env [(Quant, (Name,Type))]
  getQuantBindingsInclusive bind = do
    ctx <- stateCtxt <$> get
    return $ getBeforeInclusive "IN: getQuantBindingsInclusive" bind ctx  
      
    
  getBindingsBetween :: Binding -> Binding -> Env [(Name,Type)]
  getBindingsBetween top bottom = do
    above <- getBindings top
    target <- getBindings bottom
    let aboveSet = S.insert (elmName top) $ S.fromList $ fst <$> above 
    return $ filter (\n -> not $ S.member (fst n) aboveSet) target
    
  getAnExist :: Env (Maybe (Name,Type))
  getAnExist = do
    ctx <- stateCtxt <$> get
    let til = getTail ctx
        last = (elmQuant til, (elmName til, elmType til))
    return $ case ctx of
      Context _ _ Nothing -> Nothing
      _ -> snd <$> find (\(q,_) -> q == Exists) (last:getBefore "IN: getAnExist" til ctx)
  
  -- | `getAllBindings` gets all bindings, listed from tightest to loosest
  -- ie, ∀a:t∃b:t will return [(∃,(b,t)), (∀,(a,t))]
  getAllBindings = do
    ctx <- stateCtxt <$> get
    case ctx of
      Context _ _ Nothing -> return []
      _ -> getQuantBindingsInclusive $ getTail ctx
  
      
  getForalls :: Env ContextMap
  getForalls = do
    ctx <- ctxtMap <$> stateCtxt <$> get
    return $ anonymous <$> elmType <$> M.filter (\q -> elmQuant q == Forall) ctx
  
  getForallsAfter :: Binding -> Env (S.Set Name)
  getForallsAfter bind = do
    ctx <- stateCtxt <$> get
    return $ S.fromList $ map (fst . snd) $ filter ((== Forall) . fst) $ getAfter "IN: getForalls" bind ctx
    
  getExistsAfter :: Binding -> Env (S.Set Name)
  getExistsAfter bind = do
    ctx <- stateCtxt <$> get
    return $ S.fromList $ map (fst . snd) $ filter ((== Exists) . fst) $ getAfter "IN: getExistsAfter" bind ctx      
  
  getExists :: Env ContextMap
  getExists = do
    ctx <- ctxtMap <$> stateCtxt <$> get
    return $ anonymous <$> elmType <$> M.filter (\q -> elmQuant q == Exists) ctx
  
  getConstants :: Env ContextMap
  getConstants = ask  
  
  clearContext :: Env ()
  clearContext = do
    ContextState i _ <- get
    put $ ContextState i emptyContext
  
  
  getFullCtxt :: Env ContextMap
  getFullCtxt = do
    constants <- getConstants
    ctx <- ctxtMap <$> stateCtxt <$> get
    return $ M.union (anonymous <$> elmType <$> ctx) constants
  
  getVariablesBeforeExists :: Name -> Env ContextMap
  getVariablesBeforeExists nm = do
    constants <- getConstants
    ctx <- stateCtxt <$> get  
    let bind = ctxtMap ctx M.! nm
    return $ M.union constants 
           $ M.fromList $ (\(nm,v) -> (nm, anonymous v)) <$> snd <$> getBefore "IN: getVariablesBeforeExists" bind ctx
    
  
  modifyCtxt :: (Context -> Context) -> Env ()
  modifyCtxt f = modify $ \m -> m { stateCtxt = f $ stateCtxt m }
  
  
  -------------------------
  ---  traversal monads ---
  -------------------------
  lookupConstant :: Name -> Env (Maybe Type)
  lookupConstant x = fmap snd <$> (M.lookup x) <$> ask 
  
  type TypeChecker = ContT Spine Env
  
  typeCheckToEnv :: TypeChecker Spine -> Env (Spine,Constraint)
  typeCheckToEnv m = listen $ runContT m return
  
  
  
  
  addToEnv :: (Name -> Type -> Constraint -> Constraint) -> Name -> Type -> TypeChecker a -> TypeChecker a
  addToEnv e x ty = mapContT (censor $ e x ty) . liftLocal ask local (M.insert x $ anonymous ty)
  