module GLL.Parser (
        gllSPPF              -- run the parser
      , charS, charT, nT, epsilon   -- create terminals
      , pNodeLookup, ParseResult(..)
    ) where

import Data.Foldable hiding (forM_, toList, sum)
import Prelude  hiding (lookup, foldr, fmap, foldl, elem)
import Control.Monad
import Control.Applicative hiding (empty)
import Data.Map (Map(..), empty, insertWith, (!), toList, lookup)
import Data.Set (member, Set(..))
import qualified Data.IntMap as IM
import qualified Data.Map as M
import qualified Data.Array as A
import qualified Data.Set as S
import qualified Data.IntSet as IS

import GLL.Common
import GLL.Types.Abstract
import GLL.Types.Grammar

-- | Representation of the input string
type Input          =   A.Array Int Token

-- | Types for 
type LhsParams      =   (Nt     , Int   , GSSNode)
type RhsParams      =   (Slot   , Int   , GSSNode)

-- | The worklist and descriptor set
type Rcal           =   [(RhsParams,SPPFNode)]
type Ucal           =   IM.IntMap (IM.IntMap (S.Set (Slot, GSlot)))

-- | GSS representation
type GSS            =   IM.IntMap (M.Map GSlot [GSSEdge])
type GSSEdge        =   (GSSNode, SPPFNode)
type GSSNode        =   (GSlot, Int)
data GSlot          =   GSlot Slot
                    |   U0 
    deriving (Ord, Eq) 

-- | Pop-set
type Pcal           =   IM.IntMap (Map GSlot [Int])

-- | Connecting it all
type Mutable        =   (SPPF,Rcal, Ucal, GSS, Pcal)

-- | Monad for implicitly passing around 'context'
data GLL a          =   GLL (Mutable -> (a, Mutable))

addDescr        ::  SPPFNode -> RhsParams  ->   GLL ()
getDescr        ::  GLL (Maybe (RhsParams,SPPFNode))
addSPPFEdge     ::  SPPFNode    -> SPPFNode     ->  GLL ()
addPop          ::  GSSNode     -> Int          ->  GLL () 
getChildren     ::  GSSNode                     ->  GLL [GSSEdge]
addGSSEdge      ::  GSSNode     -> GSSEdge      ->  GLL ()
getPops         ::  GSSNode     -> GLL [Int]
joinSPPFs       ::  Slot -> SPPFNode -> Int -> Int -> Int 
                            -> GLL SPPFNode
runGLL :: GLL a -> Mutable -> Mutable
runGLL (GLL f) p = snd $ f p
addSPPFEdge f t = GLL $ \(sppf,r,u,gss,p) -> 
       ((), 
        (
        pMapInsert f t $
            sppf
        ,r,u,gss,p))

addDescr sppf alt@(slot,i,(gs,l)) = GLL $ \(dv,r,u,gss,p) -> 
    let new     = maybe True inner $ IM.lookup i u
          where inner m = maybe True (not . ((slot,gs) `S.member`)) $ IM.lookup l m  
        newU = IM.alter inner i u
         where inner mm = case mm of 
                             Nothing -> Just $ IM.singleton l single 
                             Just m  -> Just $ IM.insertWith (S.union) l single m
               single = S.singleton (slot,gs)
     in if new then ((), (dv, (alt,sppf):r, newU, gss , p))
               else ((), (dv, r, u, gss, p))

getDescr = GLL $ \(dv,r,u,gss,p) -> 
    case r of 
        []              -> (Nothing, (dv,r,u,gss,p))
        (next@(alt,sppf):rest)     -> 
            let res = (Just next, (dv,rest,u,gss,p))
            in res

addPop (gs,l) i = GLL $ \(dv,r,u,gss,p) ->
    let newP = IM.alter inner l p
         where inner mm = case mm of 
                            Nothing -> Just $ M.singleton gs [i]
                            Just m  -> Just $ M.insertWith (++) gs [i] m
    in ((), (dv,r,u,gss,newP))

getChildren (gs,l) = GLL $ \(dv,r,u,gss,p) ->
    let res = maybe [] inner $ IM.lookup l gss
         where inner m = maybe [] id $ M.lookup gs m
     in (res, (dv,r,u,gss,p))

addGSSEdge f@(gs,i) t = GLL $ \(dv,r,u,gss,p) -> 
    let newGSS = IM.alter inner i gss
         where inner mm = case mm of 
                            Nothing -> Just $ M.singleton gs [t] 
                            Just m  -> Just $ M.insertWith (++) gs [t] m
    in ((), (dv,r,u,newGSS,p))

getPops (gs,l) = GLL $ \ctx@(dv,r,u,gss,p) -> 
    let res = maybe [] inner $ IM.lookup l p
         where inner = maybe [] id .  M.lookup gs
    in (res, ctx)

instance Show GSlot where
    show (U0)       = "u0"
    show (GSlot gn) = show gn

instance Show SPPFNode where
    show (SNode (s, l, r))  = "(s: " ++ show s ++ ", " ++ show l ++ ", " ++ show r ++ ")"
    show (INode (s, l, r))  = "(i: " ++ show s ++ ", " ++ show l ++ ", " ++ show r ++ ")"
    show (PNode (p, l, k, r))  = "(p: " ++ show p ++ ", " ++ show l ++ ", " ++ show k ++ ", " ++ show r ++ ")"
    show Dummy              = "$"

instance Applicative GLL where
    (<*>) = ap
    pure  = return
instance Functor GLL where
    fmap  = liftM
instance Monad GLL where
    return a = GLL $ \p -> (a, p)
    (GLL m) >>= f  = GLL $ \p -> let (a, p')  = m p
                                     (GLL m') = f a
                                    in m' p'
gllSPPF :: Grammar -> [Token] -> ParseResult 
gllSPPF grammar@(Grammar start _ _ ) input = 
    let (mutable,_,_,_) = gll m False grammar input
        m = length input
    in resultFromMutable mutable (Nt start, 0, m)

gll :: Int -> Bool -> Grammar -> [Token] -> (Mutable, [Alt], SelectMap, FollowMap)
gll m debug (Grammar start _ rules) input' = 
    (runGLL (pLhs (start, 0, (U0,0))) context, prs, selects, follows)
 where 
    prs     = [ alt | Rule _ alts _ <- rules, alt <- (reverse alts) ]
    context = (emptySPPF, [], IM.empty, IM.empty, IM.empty)
    input   = A.array (0,m) $ zip [0..] $ input' ++ [EOS]

    dispatch    ::                                  GLL ()
    pLhs        :: LhsParams                    ->  GLL () 
    pRhs        :: RhsParams    ->  SPPFNode    ->  GLL ()

    dispatch = do
        mnext <- getDescr
        case mnext of
            Nothing            -> return () -- no continuation
            Just (next,sppf)   -> pRhs next sppf
    pLhs (bigx, i, gn) = do 
        let     alts  =  [  (Slot bigx [] beta, i, gn) | (Alt bigx beta) <- altsOf bigx
                         ,  select (input A.! i) beta bigx
                         ]
        forM_ alts (addDescr Dummy)
        dispatch 

    pRhs (Slot bigx [] [Term Epsilon], i, (gs,l)) _  = do
        root <- joinSPPFs slot Dummy l i i
        pRhs (slot, i, (gs,l)) root
     where  slot    = Slot bigx [Term Epsilon] []

    pRhs (Slot bigx alpha ((Term tau):beta), i, (gs,l)) sppf = 
     if (input A.! i == tau) 
      then do -- token test 
        root <-  joinSPPFs slot sppf l i (i+1) 
        pRhs (slot, i+1, (gs,l)) root 
      else
        dispatch
     where  slot       = Slot bigx (alpha++[Term tau]) beta

    pRhs (Slot bigx alpha ((Nt bigy):beta), i, (gs, l)) sppf = 
      if (select (input A.! i) ((Nt bigy):beta) bigx) 
        then do
          addGSSEdge ret ((gs,l), sppf) 
          rs <- getPops ret     -- has ret been popped?
          forM_ rs $ \r -> do   -- yes, use given extents
                          root <- joinSPPFs slot sppf l i r
                          addDescr root (slot, r, (gs,l))
          pLhs (bigy, i, ret)
        else
          dispatch
     where  ret      = (GSlot slot, i)
            slot     = Slot bigx (alpha++[Nt bigy]) beta

    pRhs (Slot bigy alpha [], i, (U0,0)) sppf = dispatch 

    pRhs (Slot bigy alpha [], i, gn@(GSlot slot,l)) ynode = do
        addPop gn i
        returns <- getChildren gn
        forM_ returns $ \((gs',l'),sppf) -> do  
            root <- joinSPPFs slot sppf l' l i  -- create SPPF for lhs
            addDescr root (slot, i, (gs',l'))   -- add new descriptors
        dispatch

    (prodMap,_,_,follows,selects)   = fixedMaps start prs
    follow x          = follows ! x
    select t rhs x    = t `member` (selects ! (x,rhs))
    altsOf x          = prodMap ! x
    merge m1 m2 = IM.unionWith inner m1 m2
     where inner  = IM.unionWith S.union 

joinSPPFs (Slot bigx alpha beta) sppf l k r =
    case (sppf, beta) of
--        (Dummy, _:_)    ->  return snode
        (Dummy, [])     ->  do  addSPPFEdge xnode pnode
                                addSPPFEdge pnode snode
                                return xnode
        (_, [])         ->  do  addSPPFEdge xnode pnode
                                addSPPFEdge pnode sppf
                                addSPPFEdge pnode snode
                                return xnode
        _               ->  do  addSPPFEdge inode pnode
                                addSPPFEdge pnode sppf
                                addSPPFEdge pnode snode
                                return inode
 where  x       =   last alpha  -- symbol before the dot
        snode   =   SNode (x, k, r)     
        xnode   =   SNode (Nt bigx, l, r)
        inode   =   INode ((Slot bigx alpha beta), l, r)
        pnode   =   PNode ((Slot bigx alpha beta), l, k, r)

data ParseResult = ParseResult  { sppf_result       :: SPPF
                                , success           :: Bool
                                , nr_descriptors    :: Int
                                , nr_sppf_edges     :: Int
                                , nr_gss_nodes      :: Int
                                }

resultFromMutable :: Mutable -> SNode -> ParseResult
resultFromMutable (sppf@(_,_,_,eMap,_),_,u,gss,_) s_node =
    ParseResult sppf success usize sppf_edges gsssize
 where  success     = sppf `sNodeLookup` s_node
        usize       = sum  [ S.size s   | (l, r2s) <- IM.assocs u
                                        , (r,s) <- IM.assocs r2s ]
        sppf_edges  = sum [ S.size ts | (_, ts) <- M.assocs eMap ]
        gsssize     = 1 + sum [ length $ M.keys x2s| (l,x2s) <- IM.assocs gss] 

instance Show ParseResult where
    show res = unlines $
        [   "Success: "     ++ show (success res)
        ,   "Descriptors: " ++ show (nr_descriptors res)
        ,   "SPPFEdges: "   ++ show (nr_sppf_edges res)
        ,   "GSSNodes: "    ++ show (nr_gss_nodes res)
        ]