{-# LANGUAGE FlexibleInstances, TypeOperators, MultiParamTypeClasses,
  TypeFamilies, UndecidableInstances #-}

module Examples.LLBasics
  (module Examples.LLBasics, module Examples.TermBase,
   module Examples.InnerBase, module Examples.LLBase,
   local, asks, (<$>), (<*>)) where

import Examples.TermBase (Type(..))
import Examples.InnerBase (Inner)
import qualified Examples.InnerBase as I
import Examples.LLBase

import qualified Control.Arrow as Arrow
import qualified Data.Set as Set; import Data.Set (Set)
import qualified Data.IntMap as IM; import Data.IntMap (IntMap)
import Data.Maybe (fromMaybe)

import Control.Monad.Reader.Class (MonadReader(..), asks)
import Control.Monad (ap)
import Control.Applicative (Applicative(pure, (<*>)), (<$>))



type Frees = Set Int; type Rename = IntMap Int

bump :: Frees -> Frees
bump = Set.map (subtract 1) . Set.filter (> 0)

lookupRN :: Rename -> Int -> Int
lookupRN rn i = fromMaybe i $ IM.lookup i rn

prepend :: Rename -> Rename -> Rename
prepend f = IM.unionWith const f .
            IM.fromDistinctAscList . map ((+ 1) Arrow.*** (+ 1)) . IM.toAscList



fvs :: Inner -> Frees
fvs (I.Lam ty tm) = Set.map (subtract 1) $ Set.filter (> 0) $ fvs tm
fvs (I.Var i) = Set.singleton i
fvs (I.App tm1 tm2) = fvs tm1 `Set.union` fvs tm2



newtype Mnd a = Mnd {runMnd :: ([Type], Rename, Int) -> (a, [TLD])}
instance Functor Mnd where fmap f = (>>= return . f)
instance Applicative Mnd where pure = return; (<*>) = ap
instance Monad Mnd where
  return a = Mnd $ \_ -> (a, [])
  m >>= k = Mnd $ \e@(tys, rn, sh) -> case runMnd m e of
    ~(a, w) -> Arrow.second (w ++) $ runMnd (k a) (tys, rn, sh + length w)
instance (e ~ ([Type], Rename)) => MonadReader e Mnd where
  ask = Mnd $ \ ~(x, y, _) -> ((x, y), []);
  local f (Mnd g) =
    Mnd $ \ ~(x, y, z) -> case f (x, y) of
      ~(x', y') -> g (x', y', z)

numEmissions :: Mnd Int
numEmissions = Mnd $ \ ~(_, _, z) -> (z, [])

emit :: [TLD] -> Mnd ()
emit w = Mnd $ \_ -> ((), w)

newTLD :: Type -> Frees -> Mnd Term -> Mnd Term
-- NB could check if such a TLD already exists
newTLD ty fvs m = do
  (rho, rn) <- ask; sh <- numEmissions
  let fvs' = reverse $ Set.toAscList $ bump fvs
  m >>= \tm -> emit [(map (rho !!) fvs', ty, tm)]
  return $ foldl ((. Var . lookupRN rn) . App) (DVar sh) fvs'