# LANGUAGE TypeOperators #
# LANGUAGE FlexibleContexts #
module KMC.Kleenex.Desugaring (desugarProg,desugarRegex) where

import           Control.Monad.Reader
import           Control.Monad.State
import qualified Data.ByteString as BS
import           Data.Char (ord)
import qualified Data.Map as M
import           Data.Maybe (fromMaybe)
import qualified Data.List as L
import           KMC.Kleenex.Actions
import           KMC.Kleenex.Approximation
import           KMC.Kleenex.Core
import           KMC.Kleenex.Syntax
import qualified KMC.RangeSet as RS
import qualified KMC.Syntax.External as E
import           KMC.Util.UTF8

-------------------
  Desugaring monad
-------------------

data DesugarState = DS { dsDecls    :: M.Map RIdent RTermAct
                       , dsRevDecls :: M.Map RTermAct RIdent
                       , dsFresh    :: RIdent
                       , dsApprox   :: [(RIdent, RIdent, Int)]
                       }

data DesugarContext = DC { dcIdents   :: M.Map (Ident, Bool) RIdent }

type Desugar = ReaderT DesugarContext (State DesugarState)

getFresh :: MonadState DesugarState m => m RIdent
getFresh = do
  i <- gets dsFresh
  modify $ \ds -> ds { dsFresh = dsFresh ds + 1 }
  return i

insertDecl :: MonadState DesugarState m => RIdent -> RTermAct -> m RIdent
insertDecl i t = do modify $ \ds -> ds { dsDecls    = M.insert i t (dsDecls ds)
                                       , dsRevDecls = M.insert t i (dsRevDecls ds)
                                       }
                    return i

insertApproxDecl :: MonadState DesugarState m => RIdent -> RIdent -> Int -> m RIdent
insertApproxDecl i1 i2 k = do
  _ <- insertDecl i2 $ RSeq [i1]
  modify $ \ds -> ds { dsApprox = (i1,i2,k):(dsApprox ds) }
  return i2

decl :: MonadState DesugarState m => RTermAct -> m RIdent
decl t = do
  rds <- gets dsRevDecls
  case M.lookup t rds of
    Just i -> return i
    Nothing -> do
      i <- getFresh
      insertDecl i t

lookupIdent :: MonadReader DesugarContext m => Ident -> Bool -> m RIdent
lookupIdent ident out =
  asks (M.lookup (ident, out) . dcIdents)
  >>= maybe (error $ "lookupIdent invariant broken: " ++ fromIdent ident) return

--------------------------------
-- Regular expression desugaring
--------------------------------

desugarRE :: Bool        -- ^ Indicates if input symbols are to be written as actions
  ^ Regular expression to 
  ^ Identifier of nonterminal representing desugared RE
desugarRE out re =
  case re of
  E.One            -> decl $ RSeq []
  E.Dot            -> decl $ RRead RS.universe out
  E.Chr a          -> sequence [ decl $ RRead (RS.singleton b) out | b <- encodeChar a ] >>= decl . RSeq
  E.Group _ e      -> desugarRE out e
  E.Concat e1 e2   -> sequence [ desugarRE out e1, desugarRE out e2 ] >>= decl . RSeq
  E.Branch e1 e2   -> sequence [ desugarRE out e1, desugarRE out e2 ] >>= decl . RSum
  E.Class pos ers  -> let rs  = (if pos then id else RS.complement)
                                $ RS.rangeSet [ (toEnum (ord lo), toEnum (ord hi)) | (lo, hi) <- ers ]
                      in decl $ RRead rs out
  E.Star e         -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         i <- getFresh
                         iloop <- decl $ RSeq [ie, i]
                         insertDecl i $ RSum [iloop, ieps]
  E.LazyStar e     -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         i <- getFresh
                         iloop <- decl $ RSeq [ie, i]
                         insertDecl i $ RSum [ieps, iloop]
  E.Plus e         -> do ie <- desugarRE out e
                         istar <- desugarRE out (E.Star e)
                         decl $ RSeq [ie, istar]
  E.LazyPlus e     -> do ie <- desugarRE out e
                         istar <- desugarRE out (E.LazyStar e)
                         decl $ RSeq [ie, istar]
  E.Question e     -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         decl $ RSum [ie, ieps]
  E.LazyQuestion e -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         decl $ RSum [ieps, ie]
  E.Range e n m    -> do
    ie <- desugarRE out e
    case m of
      Nothing -> do istar <- desugarRE out (E.Star e)
                    decl $ RSeq (replicate n ie ++ [istar])
      Just m' -> if n == m' then
                   decl $ RSeq (replicate n ie)
                 else
                   do iquest <- desugarRE out (E.Question e)
                      decl $ RSeq (replicate n ie ++ replicate m' iquest)
  E.Suppress e     -> desugarRE False e
  E.NamedSet _ _   -> error "named sets not supported"
  E.LazyRange _ _ _ -> error "lazy ranges not supported"


---------------------
  Kleenex desugaring
---------------------

desugarTerm :: Bool        -- ^ Indicates if input symbols are to be written as actions
            -> Term i      -- ^ Term to be desugared
            -> Desugar RIdent -- ^ Identifier of nonterminal representing desugared term
desugarTerm out t =
  case t of
  (Var ident)   -> lookupIdent ident out
  (Constant bs) -> mapM (decl . RConst . Left) (if out then BS.unpack bs else []) >>= decl . RSeq
  (RE e)        -> desugarRE out e
  (Seq _ _)     -> mapM (desugarTerm out) (flattenSeq t) >>= decl . RSeq
  (Sum _ _)     -> mapM (desugarTerm out) (flattenSum t) >>= decl . RSum
  (Star t1)     -> do it <- desugarTerm out t1
                      ieps <- decl $ RSeq []
                      i <- getFresh
                      iloop <- decl $ RSeq [it, i]
                      insertDecl i $ RSum [iloop, ieps]
  (Plus t1)     -> do it <- desugarTerm out t1
                      istar <- desugarTerm out (Star t1)
                      decl $ RSeq [it, istar]
  (Question t1) -> do it <- desugarTerm out t1
                      ieps <- decl $ RSeq []
                      decl $ RSum [it, ieps]
  (Approx e t1)  -> if not out then desugarTerm out t1 else
                      do i1 <- desugarTerm out t1
                         i2 <- getFresh
                         insertApproxDecl i1 i2 e
  (Range m' n' t1) -> do it <- desugarTerm out t1
                         let m = fromMaybe 0 m'
                         case n' of
                           Nothing -> do
                             istar <- desugarTerm out (Star t1)
                             decl $ RSeq (replicate m it ++ [istar])
                           Just n | n < m -> error $ "invalid range: {" ++ show m ++ "," ++ show n ++ "}"
                                  | n == m -> decl $ RSeq (replicate m it)
                                  | otherwise -> do
                                      iquest <- desugarTerm out (Question t1)
                                      decl $ RSeq (replicate m it ++ replicate (n-m) iquest)
  (SuppressOutput t1) -> desugarTerm False t1
  One                 -> decl $ RSeq []
  UpdateReg r str     -> mapM (decl . RConst) (Right Push
                                               :(str >>= dsUpdateSym)
                                               ++ [Right $ Pop r]) >>= decl . RSeq
  WriteReg r          -> decl $ RConst (Right (Write r))
  RedirectReg r t1    -> do it <- desugarTerm out t1
                            ~[ipush, ipop] <- mapM (decl . RConst . Right) [Push, Pop r]
                            decl $ RSeq [ipush,it,ipop]
  TermInfo _ t1       -> desugarTerm out t1
  where
    flattenSeq (Seq t1 t2) = flattenSeq t1 ++ flattenSeq t2
    flattenSeq t'          = [t']
    flattenSum (Sum t1 t2) = flattenSum t1 ++ flattenSum t2
    flattenSum t'          = [t']
    dsUpdateSym (Left r)   = [Right $ Write r]
    dsUpdateSym (Right bs) = map Left $ BS.unpack bs


desugarProg :: (Show i) => Prog i -> ApproxMetric -> ApproxMode -> Bool -> RProgAct
desugarProg prog m mode ite =
  RProg { rprogPipeline = [ lu ident | ident <- progPipeline prog ]
          , rprogDecls    = applyApproximation ds m mode ite
          }
  where
    ds        = execState (runReaderT (go (progDecls prog)) initDC) initDS
    identMap  = M.fromList $ zip [ (declIdent d, out) | d <- progDecls prog, out <- [True, False] ] [0..]
    nextFresh = M.size identMap + 1
    initDS    = DS { dsDecls = M.empty, dsRevDecls = M.empty, dsFresh = nextFresh, dsApprox = []}
    initDC    = DC { dcIdents = identMap }

    declIdent (Decl ident _) = ident
    declIdent (DeclInfo _ d) = declIdent d

    go [] = return ()
    go (DeclInfo _ d:decls) = go (d:decls)
    go (Decl ident t:decls) = do
      i <- lookupIdent ident True
      j <- lookupIdent ident False
      i' <- desugarTerm True t
      j' <- desugarTerm False t
      _ <- insertDecl i (RSeq [i'])
      _ <- insertDecl j (RSeq [j'])
      go decls

    lu ident = maybe (error $ "identifier in pipeline with no declaration: " ++ fromIdent ident) id
                     (M.lookup (ident, True) identMap)

applyApproximation :: DesugarState -> ApproxMetric -> ApproxMode -> Bool -> M.Map RIdent RTermAct
applyApproximation ds m mode ite = fst $ foldl go (dsDecls ds, dsFresh ds) $ dsApprox ds
  where
    applyOnce decls pl k c = rprogDecls $
      (if ite then approxProgIt else approxProg) (stdToCore (RProg [pl] decls)) k c m mode
    addApproxStms dold c i dnew = (M.insert i (RSeq [c]) (M.union dold dnew) , c + M.size dnew)
    approxIds = map (\(_,i,_) -> i) $ dsApprox ds
    go (decls,c) (i1,i2,k) = if k < 1 then (decls,c) else
      if null (L.intersect approxIds (calculateReach i1 decls []))
      then addApproxStms decls c i2 (applyOnce decls i1 k c)
      else error "Approximated sub-programs cannot contain approximation terms"

calculateReach :: RIdent -> M.Map RIdent RTermAct -> [RIdent] -> [RIdent]
calculateReach r terms rs =
  case M.findWithDefault (RSeq []) r terms of
    RSum ids -> foldl go rs ids
    RSeq ids -> foldl go rs ids
    _        -> rs
  where
    go rs' r' = if (elem r' rs') then rs
                else calculateReach r' terms (r':rs')

desugarRegex :: E.Regex -> RProgAct
desugarRegex re =
  RProg { rprogPipeline = [ initIdent ]
          , rprogDecls    = dsDecls ds
          }
  where
    (initIdent, ds) = runState (runReaderT (desugarRE True re) initDC) initDS
    initDC = DC { dcIdents = M.empty }
    initDS = DS { dsDecls = M.empty, dsRevDecls = M.empty, dsFresh = 0, dsApprox = [] }

-----------------
-----------------
------------------------------
 Regular expression desugaring
------------------------------
 ^ Indicates if input symbols are to be written as actions
-------------------
-------------------
 ^ Indicates if input symbols are to be written as actions
 ^ Term to be desugared
 ^ Identifier of nonterminal representing desugared term
# LANGUAGE TypeOperators #
# LANGUAGE FlexibleContexts #
module KMC.Kleenex.Desugaring (desugarProg,desugarRegex) where

import           Control.Monad.Reader
import           Control.Monad.State
import qualified Data.ByteString as BS
import           Data.Char (ord)
import qualified Data.Map as M
import           Data.Maybe (fromMaybe)
import qualified Data.List as L
import           KMC.Kleenex.Actions
import           KMC.Kleenex.Approximation
import           KMC.Kleenex.Core
import           KMC.Kleenex.Syntax
import qualified KMC.RangeSet as RS
import qualified KMC.Syntax.External as E
import           KMC.Util.UTF8

  Desugaring monad

data DesugarState = DS { dsDecls    :: M.Map RIdent RTermAct
                       , dsRevDecls :: M.Map RTermAct RIdent
                       , dsFresh    :: RIdent
                       , dsApprox   :: [(RIdent, RIdent, Int)]
                       }

data DesugarContext = DC { dcIdents   :: M.Map (Ident, Bool) RIdent }

type Desugar = ReaderT DesugarContext (State DesugarState)

getFresh :: MonadState DesugarState m => m RIdent
getFresh = do
  i <- gets dsFresh
  modify $ \ds -> ds { dsFresh = dsFresh ds + 1 }
  return i

insertDecl :: MonadState DesugarState m => RIdent -> RTermAct -> m RIdent
insertDecl i t = do modify $ \ds -> ds { dsDecls    = M.insert i t (dsDecls ds)
                                       , dsRevDecls = M.insert t i (dsRevDecls ds)
                                       }
                    return i

insertApproxDecl :: MonadState DesugarState m => RIdent -> RIdent -> Int -> m RIdent
insertApproxDecl i1 i2 k = do
  _ <- insertDecl i2 $ RSeq [i1]
  modify $ \ds -> ds { dsApprox = (i1,i2,k):(dsApprox ds) }
  return i2

decl :: MonadState DesugarState m => RTermAct -> m RIdent
decl t = do
  rds <- gets dsRevDecls
  case M.lookup t rds of
    Just i -> return i
    Nothing -> do
      i <- getFresh
      insertDecl i t

lookupIdent :: MonadReader DesugarContext m => Ident -> Bool -> m RIdent
lookupIdent ident out =
  asks (M.lookup (ident, out) . dcIdents)
  >>= maybe (error $ "lookupIdent invariant broken: " ++ fromIdent ident) return


  ^ Regular expression to 
  ^ Identifier of nonterminal representing desugared RE
desugarRE out re =
  case re of
  E.One            -> decl $ RSeq []
  E.Dot            -> decl $ RRead RS.universe out
  E.Chr a          -> sequence [ decl $ RRead (RS.singleton b) out | b <- encodeChar a ] >>= decl . RSeq
  E.Group _ e      -> desugarRE out e
  E.Concat e1 e2   -> sequence [ desugarRE out e1, desugarRE out e2 ] >>= decl . RSeq
  E.Branch e1 e2   -> sequence [ desugarRE out e1, desugarRE out e2 ] >>= decl . RSum
  E.Class pos ers  -> let rs  = (if pos then id else RS.complement)
                                $ RS.rangeSet [ (toEnum (ord lo), toEnum (ord hi)) | (lo, hi) <- ers ]
                      in decl $ RRead rs out
  E.Star e         -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         i <- getFresh
                         iloop <- decl $ RSeq [ie, i]
                         insertDecl i $ RSum [iloop, ieps]
  E.LazyStar e     -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         i <- getFresh
                         iloop <- decl $ RSeq [ie, i]
                         insertDecl i $ RSum [ieps, iloop]
  E.Plus e         -> do ie <- desugarRE out e
                         istar <- desugarRE out (E.Star e)
                         decl $ RSeq [ie, istar]
  E.LazyPlus e     -> do ie <- desugarRE out e
                         istar <- desugarRE out (E.LazyStar e)
                         decl $ RSeq [ie, istar]
  E.Question e     -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         decl $ RSum [ie, ieps]
  E.LazyQuestion e -> do ie <- desugarRE out e
                         ieps <- decl $ RSeq []
                         decl $ RSum [ieps, ie]
  E.Range e n m    -> do
    ie <- desugarRE out e
    case m of
      Nothing -> do istar <- desugarRE out (E.Star e)
                    decl $ RSeq (replicate n ie ++ [istar])
      Just m' -> if n == m' then
                   decl $ RSeq (replicate n ie)
                 else
                   do iquest <- desugarRE out (E.Question e)
                      decl $ RSeq (replicate n ie ++ replicate m' iquest)
  E.Suppress e     -> desugarRE False e
  E.NamedSet _ _   -> error "named sets not supported"
  E.LazyRange _ _ _ -> error "lazy ranges not supported"


  Kleenex desugaring

desugarTerm out t =
  case t of
  (Var ident)   -> lookupIdent ident out
  (Constant bs) -> mapM (decl . RConst . Left) (if out then BS.unpack bs else []) >>= decl . RSeq
  (RE e)        -> desugarRE out e
  (Seq _ _)     -> mapM (desugarTerm out) (flattenSeq t) >>= decl . RSeq
  (Sum _ _)     -> mapM (desugarTerm out) (flattenSum t) >>= decl . RSum
  (Star t1)     -> do it <- desugarTerm out t1
                      ieps <- decl $ RSeq []
                      i <- getFresh
                      iloop <- decl $ RSeq [it, i]
                      insertDecl i $ RSum [iloop, ieps]
  (Plus t1)     -> do it <- desugarTerm out t1
                      istar <- desugarTerm out (Star t1)
                      decl $ RSeq [it, istar]
  (Question t1) -> do it <- desugarTerm out t1
                      ieps <- decl $ RSeq []
                      decl $ RSum [it, ieps]
  (Approx e t1)  -> if not out then desugarTerm out t1 else
                      do i1 <- desugarTerm out t1
                         i2 <- getFresh
                         insertApproxDecl i1 i2 e
  (Range m' n' t1) -> do it <- desugarTerm out t1
                         let m = fromMaybe 0 m'
                         case n' of
                           Nothing -> do
                             istar <- desugarTerm out (Star t1)
                             decl $ RSeq (replicate m it ++ [istar])
                           Just n | n < m -> error $ "invalid range: {" ++ show m ++ "," ++ show n ++ "}"
                                  | n == m -> decl $ RSeq (replicate m it)
                                  | otherwise -> do
                                      iquest <- desugarTerm out (Question t1)
                                      decl $ RSeq (replicate m it ++ replicate (n-m) iquest)
  (SuppressOutput t1) -> desugarTerm False t1
  One                 -> decl $ RSeq []
  UpdateReg r str     -> mapM (decl . RConst) (Right Push
                                               :(str >>= dsUpdateSym)
                                               ++ [Right $ Pop r]) >>= decl . RSeq
  WriteReg r          -> decl $ RConst (Right (Write r))
  RedirectReg r t1    -> do it <- desugarTerm out t1
                            ~[ipush, ipop] <- mapM (decl . RConst . Right) [Push, Pop r]
                            decl $ RSeq [ipush,it,ipop]
  TermInfo _ t1       -> desugarTerm out t1
  where
    flattenSeq (Seq t1 t2) = flattenSeq t1 ++ flattenSeq t2
    flattenSeq t'          = [t']
    flattenSum (Sum t1 t2) = flattenSum t1 ++ flattenSum t2
    flattenSum t'          = [t']
    dsUpdateSym (Left r)   = [Right $ Write r]
    dsUpdateSym (Right bs) = map Left $ BS.unpack bs


desugarProg :: (Show i) => Prog i -> ApproxMetric -> ApproxMode -> Bool -> RProgAct
desugarProg prog m mode ite =
  RProg { rprogPipeline = [ lu ident | ident <- progPipeline prog ]
          , rprogDecls    = applyApproximation ds m mode ite
          }
  where
    ds        = execState (runReaderT (go (progDecls prog)) initDC) initDS
    identMap  = M.fromList $ zip [ (declIdent d, out) | d <- progDecls prog, out <- [True, False] ] [0..]
    nextFresh = M.size identMap + 1
    initDS    = DS { dsDecls = M.empty, dsRevDecls = M.empty, dsFresh = nextFresh, dsApprox = []}
    initDC    = DC { dcIdents = identMap }

    declIdent (Decl ident _) = ident
    declIdent (DeclInfo _ d) = declIdent d

    go [] = return ()
    go (DeclInfo _ d:decls) = go (d:decls)
    go (Decl ident t:decls) = do
      i <- lookupIdent ident True
      j <- lookupIdent ident False
      i' <- desugarTerm True t
      j' <- desugarTerm False t
      _ <- insertDecl i (RSeq [i'])
      _ <- insertDecl j (RSeq [j'])
      go decls

    lu ident = maybe (error $ "identifier in pipeline with no declaration: " ++ fromIdent ident) id
                     (M.lookup (ident, True) identMap)

applyApproximation :: DesugarState -> ApproxMetric -> ApproxMode -> Bool -> M.Map RIdent RTermAct
applyApproximation ds m mode ite = fst $ foldl go (dsDecls ds, dsFresh ds) $ dsApprox ds
  where
    applyOnce decls pl k c = rprogDecls $
      (if ite then approxProgIt else approxProg) (stdToCore (RProg [pl] decls)) k c m mode
    addApproxStms dold c i dnew = (M.insert i (RSeq [c]) (M.union dold dnew) , c + M.size dnew)
    approxIds = map (\(_,i,_) -> i) $ dsApprox ds
    go (decls,c) (i1,i2,k) = if k < 1 then (decls,c) else
      if null (L.intersect approxIds (calculateReach i1 decls []))
      then addApproxStms decls c i2 (applyOnce decls i1 k c)
      else error "Approximated sub-programs cannot contain approximation terms"

calculateReach :: RIdent -> M.Map RIdent RTermAct -> [RIdent] -> [RIdent]
calculateReach r terms rs =
  case M.findWithDefault (RSeq []) r terms of
    RSum ids -> foldl go rs ids
    RSeq ids -> foldl go rs ids
    _        -> rs
  where
    go rs' r' = if (elem r' rs') then rs
                else calculateReach r' terms (r':rs')

desugarRegex :: E.Regex -> RProgAct
desugarRegex re =
  RProg { rprogPipeline = [ initIdent ]
          , rprogDecls    = dsDecls ds
          }
  where
    (initIdent, ds) = runState (runReaderT (desugarRE True re) initDC) initDS
    initDC = DC { dcIdents = M.empty }
    initDS = DS { dsDecls = M.empty, dsRevDecls = M.empty, dsFresh = 0, dsApprox = [] }
