{-# LANGUAGE StandaloneDeriving, FlexibleInstances, RecordWildCards #-}
{-# OPTIONS_GHC -Wall -fno-warn-orphans #-}
module NLP.ToyDottedCky where

import Control.Applicative ((<$>),(<*>))
import Control.Exception (assert)
import Control.Monad (guard)
import Data.List (intercalate)
import qualified Data.Map as Map
import Data.Maybe (maybeToList, mapMaybe)
import Data.Tree

import Data.DList (DList(..))
import qualified Data.DList as DL

import NLP.ChartParser
import NLP.CfgStuff

go :: Cfg -> String -> String
go cfg s = intercalate "\n\n" . map drawTree . parseTrees cfg wmap
         $ parse cfg s
 where
  wmap = toWordMap . words $ s

goS :: Cfg -> String -> Bool
goS cfg s = successful config $ parse cfg s
 where
  config = toyConfig cfg wmap
  wmap = toWordMap . words $ s

-- ----------------------------------------------------------------------
-- configuration of the chart parser
-- ----------------------------------------------------------------------

data ToyItem = ToyItem { symbol    :: String
                       , beforeDot :: DList String
                       , afterDot  :: [String]
                       , left   :: Int
                       , right  :: Int }
 deriving (Show, Eq, Ord)

completed :: ToyItem -> Bool
completed = null . afterDot

instance Show a => Show (DList a) where
  show = showString "fromList " . show . DL.toList

instance Eq a => Eq (DList a) where
  d1 == d2 = DL.toList d1 == DL.toList d2

instance Ord a => Ord (DList a) where
  d1 < d2 = DL.toList d1 < DL.toList d2


parse :: Cfg -> String -> SimpleStatus ToyItem
parse cfg s = chartParse (toyConfig cfg wmap) initialSt
 where
  initialSt = Status { agenda = axioms cfg wmap
                     , chart  = [] }
  wmap = toWordMap . words $ s

-- ----------------------------------------------------------------------
-- configuration of NLP.ChartParser
-- ----------------------------------------------------------------------

toyConfig :: Cfg -> WordMap -> SimpleConfig ToyItem
toyConfig gram wmap = Config
 { inferenceRules  = ckyInferenceRules wmap
 , isGoalItem      = \s -> completed s
                        && symbol s == goal gram
                        && left  s == 0
                        && right s == rmax + 1
 }
 where
  rmax = maximum . Map.keys $ wmap

type WordMap = Map.Map Int String

toWordMap :: [String] -> Map.Map Int String
toWordMap = Map.fromList . zip [0..]

-- ----------------------------------------------------------------------
-- axioms
-- ----------------------------------------------------------------------

axioms :: Cfg -> WordMap -> [ToyItem]
axioms cfg wmap =
 mkItem <$> Map.keys wmap <*> rules cfg

mkItem :: Int -> CfgRule -> ToyItem
mkItem i (CfgRule l rs) =
 ToyItem { symbol = l
         , beforeDot = DL.empty
         , afterDot  = rs
         , left = i
         , right = i
         }

-- ----------------------------------------------------------------------
-- inference rules
-- ----------------------------------------------------------------------

ckyInferenceRules :: WordMap -> [ToyItem] -> ToyItem -> [ToyItem]
ckyInferenceRules wmap chrt aItem =
    mapMaybe (     completeR aItem) cItems ++
    mapMaybe (flip completeR aItem) cItems ++
    scanR wmap aItem
 where
    cItems = chrt

-- | we ignore the chart item for this inference rule
scanR :: WordMap -> ToyItem -> [ToyItem]
scanR wmap aItem = maybeToList $ do
  w           <- Map.lookup (right aItem) wmap
  (w2,after2) <- maybeHeadTail (afterDot aItem)
  guard (w == w2)
  return $ aItem { beforeDot = beforeDot aItem `DL.snoc` w
                 , afterDot  = after2
                 , right     = right aItem + 1 }

completeR :: ToyItem -> ToyItem -> Maybe ToyItem
completeR bItem cItem =
 if right bItem == left cItem && completed cItem
 then do (bc,brest) <- maybeHeadTail (afterDot bItem)
         guard (bc == symbol cItem)
         return $ bItem { beforeDot = beforeDot bItem `DL.snoc` bc
                        , afterDot  = brest
                        , right     = right cItem
                        }
 else Nothing

-- ----------------------------------------------------------------------
-- returning results
-- ----------------------------------------------------------------------

-- I have a sneaking suspicion that there is a far simpler way to
-- write this... intuition tells me that some of the checks here are
-- just redundant code
parseTrees :: Cfg -> WordMap -> SimpleStatus ToyItem -> [Tree String]
parseTrees (Cfg start _) wmap st =
   assert (all (\r -> length r == 1) result) -- exactly 1
 $ map head result
 where
  -- showCChart = unlines . map show $ cchart
  result = helper start [] 0 (rmax + 1)
  lchart = map (uncurry mkLexItem) . Map.toList $ wmap
  cchart = filter completed . chart $ st
  rmax   = maximum $ Map.keys wmap
  helper :: String -> [String] -> Int -> Int -> [[Tree String]]
  helper c cs l r =
   do -- trace (unwords $ ["at", show l, show r, c] ++ cs) $ guard True
      let isMatch = if null cs then isLastNode else isLeftNode
      ti <- filter isLastNode lchart ++ filter isMatch cchart
      -- trace (unwords $ ["--", show l, show r, c, show ti]) $ guard True
      let rti :: Int
          rti = right ti
      guard $ case cs of -- avoid left recursion loop
                []     -> rti == r
                (c2:_) -> any (\t -> left t == rti && symbol t == c2) cchart
      -- trace (unwords $ ["--", show l, show r, c, "kids"]) $ guard True
      kids <- case DL.toList (beforeDot ti) of
                []       -> return []
                (kc:kcs) -> helper kc kcs l rti
      -- trace (unwords $ ["--", show l, show r, c, "siblings"]) $ guard True
      siblings <- case cs of
                    []       -> return []
                    (c2:cs2) -> helper c2 cs2 rti r
      return $ Node c kids : siblings
   where
    isLeftNode t = left t == l && right t <  r && symbol t == c
    isLastNode t = left t == l && right t == r && symbol t == c

mkLexItem :: Int -> String -> ToyItem
mkLexItem i w =
 ToyItem { symbol = w
         , beforeDot = DL.empty
         , afterDot  = []
         , left = i
         , right = i+1
         }

-- ----------------------------------------------------------------------
-- returning results
-- ----------------------------------------------------------------------

maybeCons :: Maybe a -> [a] -> [a]
maybeCons = maybe id (\x -> (x :))

maybeHeadTail :: [a] -> Maybe (a,[a])
maybeHeadTail [] = Nothing
maybeHeadTail (x:xs) = Just (x,xs)