{-# LANGUAGE TypeOperators, FlexibleInstances #-}

module GLL.Combinators.MemInterface (
    SymbParser, IMParser,
    HasAlts(..), IsSymbParser(..), IsIMParser(..),
    parse, parseString,
    char, token, Token(..),
    epsilon, satisfy,
    many, some, optional,
    (<::=>),(<:=>),
    (<$>),
    (<$),
    (<*>),
    (<*),
    (<|>),
    (:.),
    memo, newMemoTable, MemoRef, MemoTable
    ) where

import Prelude hiding ((<*>), (<*), (<$>), (<$))

import GLL.Combinators.Options
import GLL.Combinators.Memoisation
import GLL.Types.Grammar hiding (epsilon)
import GLL.Types.Abstract
import GLL.Parser (gllSPPF, pNodeLookup, ParseResult(..))

import Control.Compose
import Control.Monad
import Data.List (unfoldr,intersperse)
import           Data.IORef
import qualified Data.IntMap as IM
import qualified Data.Map as M
import qualified Data.Set as S

type SymbVisit1 b = Symbol 
type SymbVisit2 b = M.Map Nt [Alt] -> M.Map Nt [Alt]
type SymbVisit3 b = PCOptions -> ParseContext -> SPPF -> Int -> Int -> IO [b]

type IMVisit1 b   = [Symbol] 
type IMVisit2 b   = M.Map Nt [Alt] -> M.Map Nt [Alt]
type IMVisit3 b   = PCOptions -> (Alt,Int) -> ParseContext -> SPPF -> Int -> Int -> IO [b]

type ParseContext = IM.IntMap (IM.IntMap (S.Set Nt))

data SymbParser b = SymbParser (SymbVisit1 b,SymbVisit2 b, SymbVisit3 b)
data IMParser b   = IMParser (IMVisit1 b, IMVisit2 b, IMVisit3 b)

parse' :: (IsSymbParser s) => PCOptions -> s a -> [Token] -> (Grammar, ParseResult, IO [a])
parse' opts p' input' =  
    let input               = input' ++ [Char 'z']
        SymbParser (Nt start,vpa2,vpa3) = toSymb (id <$> p' <* char 'z')
        snode               = (start, 0, m)
        m                   = length input
        rules               = vpa2 M.empty
        as                  = vpa3 opts IM.empty sppf 0 m
        grammar = Grammar start [] [ Rule x alts | (x, alts) <- M.assocs rules ]
        parse_res           = gllSPPF grammar input
        sppf                = sppf_result parse_res
    in (grammar, parse_res, as)

-- | The grammar of a given parser
grammar :: (IsSymbParser s) => s a -> Grammar
grammar p = (\(f,_,_) -> f) (parse' defaultOptions p [])

-- | The semantic results of a parser, given a string of Tokens
parse :: (IsSymbParser s) => s a -> [Token] -> IO [a]
parse = parseWithOptions defaultOptions 

-- | Change the behaviour of the parse using GLL.Combinators.Options
parseWithOptions :: (IsSymbParser s) => PCOptions -> s a -> [Token] -> IO [a]
parseWithOptions opts p = (\(_,_,t) -> t) . parse' opts p

-- | Parse a string of characters
parseString :: (IsSymbParser s) => s a -> String -> IO [a]
parseString = parseStringWithOptions defaultOptions

-- | Parse a string of characters using options
parseStringWithOptions :: (IsSymbParser s) => PCOptions -> s a -> String -> IO [a]
parseStringWithOptions opts p  = parseWithOptions opts p . map Char

-- | Get the SPPF produced by parsing the given input with the given parser
sppf :: (IsSymbParser s) => s a -> [Token] -> ParseResult
sppf p str =  (\(_,s,_) -> s) $ parse' defaultOptions p str

inParseContext :: ParseContext -> (Symbol, Int, Int) -> Bool
inParseContext ctx (Nt x, l, r) = maybe False inner $ IM.lookup l ctx
 where  inner = maybe False (S.member x) . IM.lookup r

toParseContext :: ParseContext -> (Nt, Int, Int) -> ParseContext
toParseContext ctx (x, l, r) = IM.alter inner l ctx
 where  inner mm = case mm of 
                    Nothing -> Just $ singleRX
                    Just m  -> Just $ IM.insertWith (S.union) r singleX m
        singleRX = IM.singleton r singleX
        singleX  = S.singleton x

infixl 2 <::=>
-- | Use this combinator on all combinators that might have an infinite
--  number of derivations for some input string. A non-terminal has
--  this property if and only if it is left-recursive and would be
--  left-recursive if all the right-hand sides of the productions of the
--  grammar are reversed.
(<::=>) :: (HasAlts b) => String -> b a -> SymbParser a 
x <::=> altPs' =
    let vas1 = [ va1 | va1 <- map (\(IMParser (f,_,_)) -> f) altPs ]
        alts  = map (Alt x) vas1    
        altPs = unO $ altsOf altPs' in SymbParser
        (Nt x
       ,\rules ->
           if x `M.member` rules 
            then rules 
            else foldr ($) (M.insert x alts rules) $ (map (\(IMParser (_,s,_)) -> s) altPs)
       ,\opts ctx sppf l r -> 
        let ctx' = ctx `toParseContext` (x,l,r)
            sems = zip alts (map (\(IMParser (_,_,t)) -> t) altPs) 
            seq (alt@(Alt _ rhs), va3) = va3 opts (alt,length rhs) ctx' sppf l r 
        in if ctx `inParseContext` (Nt x, l, r) 
                then return []
                else do ass <- forM sems seq
                        return (concatChoice opts ass)
       )

infixl 2 <:=>
-- | Use this combinator on all recursive non-terminals
(<:=>) :: (HasAlts b) => String -> b a -> SymbParser a 
x <:=> altPs' =
    let vas1 = [ va1 | va1 <- map (\(IMParser (f,_,_)) -> f) altPs ]
        alts  = map (Alt x) vas1    
        altPs = unO $ altsOf altPs' in SymbParser
        (Nt x
       ,\rules ->
           if x `M.member` rules 
            then rules 
            else foldr ($) (M.insert x alts rules) $ (map (\(IMParser (_,s,_)) -> s) altPs)
       ,\opts ctx sppf l r -> 
        let sems = zip alts (map (\(IMParser (_,_,t)) -> t) altPs)
            seq (alt@(Alt _ rhs), va3) = va3 opts (alt,length rhs) ctx sppf l r 
        in do   ass <- forM sems seq
                return (concatChoice opts ass)
       )

concatChoice :: PCOptions -> [[a]] -> [a]
concatChoice opts ress = if left_biased_choice opts
                            then firstRes ress
                            else concat ress
 where  firstRes []         = []
        firstRes ([]:ress)  = firstRes ress
        firstRes (res:_)    = res

infixl 4 <*>
(<*>) :: (IsIMParser i, IsSymbParser s) => i (a -> b) -> s a -> IMParser b
pl' <*> pr' = 
  let IMParser (vimp1,vimp2,vimp3) = toImp pl'
      SymbParser (vpa1,vpa2,vpa3)  = toSymb pr' in IMParser
  (vimp1++[vpa1]
  ,\rules ->  let rules1  = vpa2 rules
                  rules2  = vimp2 rules1 in rules2
  ,\opts (alt@(Alt x rhs),j) ctx sppf l r ->
    let ks     = maybe [] id $ sppf `pNodeLookup` ((alt,j), l, r)
        filter = maybe id id $ pivot_select opts
        seq k  = do     as      <- vpa3 opts ctx sppf k r
                        a2bs    <- vimp3 opts(alt,j-1) ctx sppf l k
                        return [ a2b a | a2b <- a2bs, a <- as ]
    in do   ass <- forM (filter ks) seq
            return (concat ass)
  )


infixl 4 <*
(<*) :: (IsIMParser i, IsSymbParser s) => i b -> s a -> IMParser b

pl' <* pr' = 
  let IMParser (vimp1,vimp2,vimp3) = toImp pl'
      SymbParser (vpa1,vpa2,vpa3)  = toSymb pr' in IMParser
  (vimp1++[vpa1]
  ,\rules ->
    let rules1  = vpa2 rules
        rules2  = vimp2 rules1
    in rules2
  ,\opts (alt@(Alt x rhs),j) ctx sppf l r ->
    let ks      = maybe [] id $ sppf `pNodeLookup` ((alt,j), l, r)
        filter  = maybe id id $ pivot_select opts
        seq k   = do    as <- vpa3 opts ctx sppf k r
                        bs <- vimp3 opts (alt,j-1) ctx sppf l k
                        return [ b | b <- bs, a <- as ]
    in do   ass <- forM (filter ks) seq
            return (concat ass)
  )

infixl 4 <$>
(<$>) :: (IsSymbParser s) => (a -> b) -> s a -> IMParser b
f <$> p' = 
    let SymbParser (vpa1,vpa2,vpa3) = toSymb p' in IMParser
      ([vpa1]
      ,\rules -> 
        vpa2 rules
      ,\opts (alt,j) ctx sppf l r ->
        let a   = vpa3 opts ctx sppf l r
            ks  = maybe [] id $ sppf `pNodeLookup` ((alt,1),l,r)
        in if null ks then return [] else do  res <- a
                                              return (map f res)
      )

infixl 4 <$
(<$) :: (IsSymbParser s) => b -> s a -> IMParser b
f <$ p' = 
    let SymbParser (vpa1,vpa2,vpa3) = toSymb p' in IMParser 
      ([vpa1]
      ,\rules -> 
        vpa2 rules
      ,\opts (alt,j) ctx sppf l r ->
        let a   = vpa3 opts ctx sppf l r
            ks  = maybe [] id $ sppf `pNodeLookup` ((alt,1),l,r)
        in if null ks then return [] else do  res <- a
                                              return (map (const f) res)
      )

infixr 3 <|>
(<|>) :: (IsIMParser i, HasAlts b) => i a -> b a -> ([] :. IMParser) a
l' <|> r' = let l = toImp l'
                r = altsOf r'
            in O (l : unO r)

memo :: (IsSymbParser s) => MemoRef [a] -> s a -> SymbParser a
memo ref p' = let   SymbParser (sym,rules,sem) = toSymb p' 
                    lhs_sem opts ctx sppf l r = do
                        tab <- readIORef ref
                        case memLookup (l,r) tab of
                            Just as -> return as
                            Nothing -> do   as <- sem opts ctx sppf l r
                                            modifyIORef ref (memInsert (l,r) as)
                                            return as
               in SymbParser (sym, rules, lhs_sem)

raw_parser :: Token -> (Token -> a) -> SymbParser a 
raw_parser t f = SymbParser (Term t, id,\_ _ _ _ _ -> return [f t])

token :: Token -> SymbParser Token
token t = raw_parser t id 

char :: Char -> SymbParser Char
char c = raw_parser (Char c) (\(Char c) -> c) 

epsilon :: SymbParser ()
epsilon = raw_parser (Epsilon) (\_ -> ()) 

satisfy :: a -> IMParser a
satisfy a = a <$ epsilon

many :: SymbParser a -> SymbParser [a]
many p = SymbParser f
 where  SymbParser (myx,_,_) = p
        SymbParser f = many_ ("(" ++ show myx ++ ")^") p

many_ x p = x <:=> (:) <$> p <*> many_ x p <|> [] <$ epsilon

some :: SymbParser a -> SymbParser [a]
some p = SymbParser f
 where  SymbParser (myx,_, _) = p
        SymbParser f = some_ ("(" ++ show myx ++ ")+") p 

some_ x p = x <:=> (:) <$> p <*> some_ x p <|> (:[]) <$> p

optional :: SymbParser a -> SymbParser (Maybe a)
optional p = SymbParser f
    where SymbParser (myx, _, _) = p 
          SymbParser f = optional_ ("(" ++ show myx ++ ")?") p

optional_ x p = x <:=> Just <$> p <|> (Nothing <$ epsilon)

class HasAlts a where
    altsOf :: a b -> ([] :. IMParser) b

instance HasAlts IMParser where
    altsOf = O . (:[])

instance HasAlts SymbParser where
    altsOf = altsOf . toImp

instance HasAlts ([] :. IMParser) where
    altsOf = id

class IsIMParser a where
    toImp :: a b -> IMParser b

instance IsIMParser IMParser where
    toImp = id

instance IsIMParser SymbParser where
    toImp p = id <$> p

instance IsIMParser ([] :. IMParser) where
    toImp = toImp . toSymb

class IsSymbParser a where
    toSymb :: a b -> SymbParser b

instance IsSymbParser IMParser where
    toSymb = toSymb . O . (:[]) 

instance IsSymbParser SymbParser where
    toSymb = id 

instance IsSymbParser ([] :. IMParser) where
    toSymb a = mkName <:=> a 
        where mkName = "_" ++ concat (intersperse "|" (map op (unO a)))
                where op (IMParser (rhs,_,_)) = concat (intersperse "*" (map show rhs))