{-# LANGUAGE MonadComprehensions #-}
{-# LANGUAGE FlexibleContexts #-}

module ParserW (
               module Parse,
               WeightedList,
               step,
               parse,
               parseA,
               parse',
               parseA',
               -- parseSel,
               -- selectSize,
               shift,
               shift',
               -- shifts,
               connect,
               addHead,
               addDep) where

import Lexicon
import Grammar
import Parse
import Node

import Data.List
import Data.Maybe
import Data.Sequence
import Data.Foldable
import Toolbox
import Control.Monad
import Control.Monad.Trans.State
import Control.Monad.Trans.Class
import Control.Monad.Plus
import Base
import Weighted

  
shift :: (MonadPlus m, Lexicon m κ δ) => δ -> String -> Parse τ κ -> m (Parse τ κ)
shift d w p = [ p << Step (nextId p) c [] [] | c <- dic d w ]
  where
    nextId [] = 1
    nextId (Step i _ _ _:_) = i + 1


shift' :: (MonadPlus m, Lexicon m κ δ) => δ -> String -> Parse τ (m κ) -> m (Parse τ (m κ))
shift' d w p = return $ p << Step (nextId p) (dic d w) [] []
  where
    nextId [] = 1
    nextId (Step i _ _ _:_) = i + 1


                              
-- shift :: (MonadPlus m, Lexicon m κ δ) => δ -> String -> Parse τ κ -> StateT Int m (Parse τ κ)
-- shift d w p = do nextId <- get
--                  c <- lift (dic d w)
--                  let ret = p << Step nextId c [] []
--                  modify (+ 1)
--                  return ret

-- where
  --   nextId [] = 1
  --   nextId (Step i _ _ _:_) = i + 1


-- shifts :: (MonadPlus m) => [String] -> Parse -> m Parse
-- shifts (s:ss) p = shift s p >>= shifts ss
-- shifts [] p = return p


addHead, addDep :: (MonadPlus m, Grammar m τ κ γ) => γ -> Parse τ κ -> m (Parse τ κ)
addHead g p = [ p +-> Head r (ind v') | let v = lastNode p, headless v, (r,v') <- heads g v ]
addDep  g p = [ p +<- Dep  r (ind v') | let v = lastNode p, (r,v') <- deps g v, headless v' ]

addHead', addDep' :: (MonadPlus m, Grammar' m τ κ γ) => γ -> Parse τ (m κ) -> m (Parse τ (m κ))
addHead' g p = [ p +-> Head r (ind v') | let v = lastNode p, headless v, (r,v') <- heads' g v ]
addDep'  g p = [ p +<- Dep  r (ind v') | let v = lastNode p, (r,v') <- deps' g v, headless v' ]

                                                
connect :: (MonadPlus m, Grammar m τ κ γ) => γ -> Parse τ κ -> m (Parse τ κ)
connect g = (addDep g >=> connect g) <> addHead g <> return

connect' :: (MonadPlus m, Grammar' m τ κ γ) => γ -> Parse τ (m κ) -> m (Parse τ (m κ))
connect' g = (addDep' g >=> connect' g) <> addHead' g <> return

          
step :: (MonadPlus m, Lexicon m κ δ, Grammar m τ κ γ) => δ -> γ -> Wordform -> Parse τ κ -> m (Parse τ κ)
step d g w = shift d w >=> connect g |? (<= 4) . trees

step' :: (MonadPlus m, Lexicon m κ δ, Grammar' m τ κ γ) => δ -> γ -> Wordform -> Parse τ (m κ) -> m (Parse τ (m κ))
step' d g w = shift' d w >=> connect' g |? (<= 4) . trees

steps :: (MonadPlus m, Lexicon m κ δ, Grammar m τ κ γ) => δ -> γ -> [Wordform] -> m (Parse τ κ)
steps d g = foldM (flip (step d g)) mzero

steps' :: (MonadPlus m, Lexicon m κ δ, Grammar' m τ κ γ) => δ -> γ -> [Wordform] -> m (Parse τ (m κ))
steps' d g = foldM (flip (step' d g)) mzero

parser :: (MonadPlus m, Lexicon m κ δ, Grammar m τ κ γ) => δ -> γ -> [Wordform] -> m (Parse τ κ)
parser d g = steps d g -- |? (==1) . trees

parser' :: (MonadPlus m, Lexicon m κ δ, Grammar' m τ κ γ) => δ -> γ -> [Wordform] -> m (Parse τ (m κ))
parser' d g = steps' d g -- |? (==1) . trees
-- parser = steps |? pass.lastNode 

-- parse :: (Lexicon WeightedList κ δ, Grammar WeightedList τ κ γ) => δ -> γ -> String -> WeightedList (Parse τ κ)
parse d g = (parser d g |? (==1) . trees |? pass g . lastNode) . words

parse' :: (Lexicon WeightedList κ δ, Grammar' WeightedList τ κ γ) => δ -> γ -> String -> WeightedList (Parse τ (WeightedList κ))
parse' d g = (parser' d g |? (==1) . trees) . words

parseA :: (Lexicon WeightedList κ δ, Grammar WeightedList τ κ γ) => δ -> γ -> String -> WeightedList (Parse τ κ)
parseA d g = parser d g . words

parseA' :: (Lexicon WeightedList κ δ, Grammar' WeightedList τ κ γ) => δ -> γ -> String -> WeightedList (Parse τ (WeightedList κ))
parseA' d g = parser' d g . words


type WeightedList = WeightedT []

-- type WeightedList = StateT Int (WeightedT [])

-- LEXER I PARSER ODDZIELNIE

-- step :: (Ind,String) -> [Step]
-- step (i,w) = [ Step i c [] [] | c <- dic w ]

-- steps :: [(Ind,String)] -> [[Step]]
-- steps []    = [[]]
-- _steps ((i,w):t)  = [ Step i c [] [] : ss | c <- dic w, ss <- _steps t ]

-- _lexer = words >>> zip [1..] >>> _steps

-- -- lexer3 s = map (Node []) (stes s)

-- -- back = (lng >=> back) <?> pure
-- _parser = connect >=> (rng >=> _parser) <?> pure

-- -- --parser'' ws = shifts ws ep >>= back >>= pa

-- _parse = _lexer >=> _parser









-- shift :: (MonadPlus m) => Wordform -> Parse -> m Parse
-- shift w p = [ p << Step (nextId p) c [] [] | c <- mfromList (dic w) ]
--   where
--     nextId [] = 1
--     nextId (Step i _ _ _:_) = i + 1

-- shifts :: (MonadPlus m) => [String] -> Parse -> m Parse
-- shifts (s:ss) p = shift s p >>= shifts ss
-- shifts [] p = return p

-- addHead, addDep :: (MonadPlus m) => Parse -> m Parse
-- addHead p = [ p +-> Arc r (ind v') | let v = lastNode p, headless v, (r,v') <- mfromList (heads v) ]
-- addDep  p = [ p +<- Arc r (ind v') | let v = lastNode p, (r,v') <- mfromList (deps v), headless v' ]

-- connect :: (MonadPlus m) => Parse -> m Parse
-- connect = (addDep >=> connect) <> addHead <> return

-- step :: (MonadPlus m) => Wordform -> Parse -> m Parse
-- step w = shift w >=> connect

-- steps :: (MonadPlus m) => [Wordform] -> m Parse
-- steps = foldM (flip step) mzero

-- parser :: (MonadPlus m) => [Wordform] -> m Parse
-- parser = steps |? pass.lastNode |? (==1).trees

-- parse :: String-> [Parse]
-- parse = parser . words


        
-- stepSel :: (Parse -> [Parse]) -> Wordform -> Parse -> [Parse]
-- stepSel s w = (shift w >=> connect >=> s)

-- stepsSel :: (Parse -> [Parse]) -> [Wordform] -> [Parse]
-- stepsSel s = foldM (flip (stepSel s)) []

-- parserSel :: (Parse -> [Parse]) -> [Wordform] -> [Parse]
-- parserSel s = stepsSel s >=> selectWith (pass.lastNode) >=> selectSize (==1)

-- parseSel :: (Parse -> [Parse]) -> String-> [Parse]
-- parseSel s = parserSel s . words








-- EXTENSIONS

-- type Selector = Parse -> [Parse]

-- sAny :: (Parse -> [Parse])
-- sAny v = pure v

-- selectWith :: (Parse -> Bool) -> (Parse -> [Parse])
-- selectWith f v = if f v then [v] else []

-- selectSize :: (Int -> Bool) -> (Parse -> [Parse])
-- selectSize f = selectWith (f . trees)
-- sForestSize f p = if f (trees p) then [p] else []

-- shift' :: (Ind,Word) -> Parse -> [Parse]
-- shift' (i,w) p = [ p << (i,c) | c <- dic w]

-- addHead', addDep' :: Parse -> [Parse]
-- addHead' n = [ n +-> (r,ind n') | headless n, (r,n') <- heads' n ]
-- addDep'  n = [ n +<- (r,ind n') | (r,n') <- deps' n, headless n' ]

-- connect' :: Parse -> [Parse]
-- -- connect = ((addDep <> addHead) >=> connect) <> pure
-- connect' = addDep' <> addHead' <> pure

-- step' :: (Parse -> [Parse]) -> (Ind,Word) -> Parse -> [Parse]
-- step' s w = shift' w >=> connect >=> s

-- step'''' w = shift' w >=> connect

-- stepS s w = step'''' w >=> s

-- parser' s = foldM (flip (stepS s)) ep >=> accept'
-- parse' s = words >>> zip [1..] >>> parser' s

-- parserRL s = foldrM ((step' s)) ep >=> accept'
-- parseRL s = words >>> zip [1..] >>> parserRL s

-- type Filter   = [Parse] -> [Parse]


-- -- takeWith :: (Parse -> Bool) -> ([Parse] -> [Parse])
-- -- takeWith f = filter f

-- -- takeBestN = mapM (sForestSize)