content: -- Copyright (c) 2018, Michael D. Adams
  --
  -- All rights reserved.
  --
  -- Redistribution and use in source and binary forms, with or without
  -- modification, are permitted provided that the following conditions are met:
  --
  --     * Redistributions of source code must retain the above copyright
  --       notice, this list of conditions and the following disclaimer.
  --
  --     * Redistributions in binary form must reproduce the above
  --       copyright notice, this list of conditions and the following
  --       disclaimer in the documentation and/or other materials provided
  --       with the distribution.
  --
  --     * Neither the name of Michael D. Adams nor the names of other
  --       contributors may be used to endorse or promote products derived
  --       from this software without specific prior written permission.
  --
  -- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  -- "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  -- LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  -- A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  -- OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  -- SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  -- LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  -- DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  -- THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  -- (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  -- OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  {-# LANGUAGE OverloadedLists #-}
  {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE MultiWayIf #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE ViewPatterns #-}
  module Data.TreeGrammar
  ( Grammar(..)
  , Rhs(..)
  , Productions
  , IsGrammar
  , grammar
  , lookup
  , lookup'
  , union
  , intersection
  , epsilonClosure
  , dropUnreachable
  , dropUnproductive
  , productive
  , determinize
  , subsetOf
  , subsetOf'
  , toSubterms
  , fromSubterms
  , isEmpty
  , minimize
  , identical
  , equivalenceClasses
  , hashTerminals
  , hashDepth
  )
  where
  
  import           Prelude hiding (lookup,repeat)
  
  import           Control.Arrow
  import           Control.Monad.State
  import           Control.Monad.Writer
  import           Control.Monad.Reader
  import qualified Control.Monad.Trans.UnionFind as U
  
  import           Data.Hashable
  import           Data.HashMap.Lazy (HashMap)
  import qualified Data.HashMap.Lazy as Map
  import           Data.HashSet (HashSet)
  import qualified Data.HashSet as Set
  import           Data.Foldable
  import qualified Data.List as L
  import           Data.Order
  import           Data.Sequence (Seq)
  import qualified Data.Abstract.Either as A
  import qualified Data.Abstract.Boolean as A
  import           Data.Utils(powComplementPick,forAll,exists)
  import           Data.Functor.Identity
  
  import           Data.TreeGrammar.OrdMap(OrdMap)
  import qualified Data.TreeGrammar.OrdMap as O
  import           Data.TreeGrammar.NonTerminal
  import           Data.TreeGrammar.Terminal (Terminal)
  import qualified Data.TreeGrammar.Terminal as T
  import           Data.Empty
  
  import           Text.Printf
  
  type Identifiable a = (Hashable a, Eq a)
  
  data Grammar n t = Grammar { start :: n, productions :: Productions n t }
  type Productions n t = HashMap n (Rhs n t)
  data Rhs n t = Rhs { cons :: t n, eps :: HashSet n } deriving (Eq)
  type IsGrammar n t = (Terminal t, Monoid (t n), NonTerminal n, Identifiable n)
  
  grammar :: (IsGrammar n t, Monoid (t String))
          => String -> [(String,t String)] -> [(String,[String])] -> Grammar n t
  grammar s prods epsilon
    =  generate
    $  rename
    $  Grammar s
    $  Map.fromListWith (<>)
    $  [ (n,Rhs c mempty) | (n,c) <- prods ]
    ++ [ (n,Rhs mempty (Set.fromList es)) | (n,es) <- epsilon]
  
  
  rename :: (Identifiable n, IsGrammar n' t) => Grammar n t -> State (Gen n') (Grammar n' t)
  rename g = do
    i <- get
    let ((s'',p'),(i',_)) = runState (do
          s' <- go (start g)
          productions' <- traverse (\(n,r) -> do
                        n' <- go n
                        c' <- T.traverse go (cons r)
                        e' <- traverse go (Set.toList (eps r))
                        return (n',Rhs c' (Set.fromList e')))
                      (Map.toList (productions g))
          return (s',Map.fromList productions'))
          (i,mempty)
    put i'
    return $ Grammar s'' p'
    where
      go :: (Identifiable n, NonTerminal n') => n -> State (Gen n', HashMap n n') n'
      go n = do
        (_,rmap) <- get
        case Map.lookup n rmap of
          Just x -> return x
          Nothing -> do
            n' <- fresh' Nothing
            modify $ second $ Map.insert n n'
            return n'
  
  subsetOf :: (IsGrammar n1 t, IsGrammar n2 t) => Grammar n1 t -> Grammar n2 t -> Bool
  subsetOf m1 m2 = fst $ subsetOf' m1 m2 O.empty
  
  subsetOf' :: forall n1 n2 t. (IsGrammar n1 t, IsGrammar n2 t)
           => Grammar n1 t -> Grammar n2 t -> OrdMap [n1] [n2] -> (Bool,OrdMap [n1] [n2])
  subsetOf' g1 g2 m = case runStateT (go ([[start g1]],[[start g2]])) m of
    Just ((),m') -> (True,m')
    Nothing -> (False,O.insertNotLeq [start g1] [[start g2]] m)
    where
      go :: (HashSet [n1],HashSet [n2]) -> StateT (OrdMap [n1] [n2]) Maybe ()
      go ([[]],[[]]) = return ()
      go (l1,l2) = forM_ l1 $ \m1 -> do
        seen <- get
        case O.leq m1 l2 seen of
          A.True  -> return ()
          A.False -> fail ""
          A.Top   -> do
            put $ O.insertLeq m1 l2 seen
            -- Try the cheap subset check first. If it did not work, try the exponential check.
            exists l2 (\ys ->
               forAll (zip m1 ys) $ \(x,y) -> do
                 T.subsetOf go (lookup x g1) (lookup y g2)
                 modify $ O.insertLeq [x] [[y]])
             `mplus`
             forAll (powComplementPick (L.transpose (Set.toList l2))) (\l2' ->
               exists (zip m1 l2') $ \(x,ys) -> do
                 T.subsetOf go (lookup x g1) (lookup' ys g2)
                 modify $ O.insertLeq [x] [ys])
  
  
  union :: (IsGrammar n t) => Grammar n t -> Grammar n t -> Grammar n t
  union = union'
  
  union' :: forall n1 n2 n' t. (IsGrammar n1 t, IsGrammar n2 t, IsGrammar n' t)
         => Grammar n1 t -> Grammar n2 t -> Grammar n' t
  union' g1 g2 = generate $ do
    g1' <- rename g1
    g2' <- rename g2
    s <- fresh Nothing
    return $ Grammar s (Map.insert s (Rhs mempty [start g1',start g2']) $ Map.union (productions g1') (productions g2'))
  
  intersection :: forall n1 n2 n' t. (IsGrammar n1 t, IsGrammar n2 t, IsGrammar n' t)
               => Grammar n1 t -> Grammar n2 t -> Grammar n' t
  intersection g1 g2 = generate $ do
    i <- get
    let (s',(_,(ps,_))) = runState (go (A.LeftRight (start g1) (start g2))) (i,(mempty,emptyRenameMap))
    return $ Grammar s' ps
    where
      go :: A.Either n1 n2 -> State (Gen n', (Productions n' t, RenameMap n1 n2 n')) n'
      go n =
        case n of
          A.Left  n1 -> addProduction n $ T.traverse (go . A.Left) (lookup n1 g1)
          A.Right n2 -> addProduction n $ T.traverse (go . A.Right) (lookup n2 g2)
          A.LeftRight n1 n2 -> do
            m <- getRenameMap
            case (O.leq [n1] [[n2]] (ordMap1 m), O.leq [n2] [[n1]] (ordMap2 m)) of
              (A.True,_)        -> go (A.Left n1)
              (_,A.True)        -> go (A.Right n2)
              (A.False,A.False) -> addProduction n
                                 $ T.intersection (go . uncurry A.LeftRight) (lookup n1 g1) (lookup n2 g2)
              (_,_) -> do
                updateOrdMap g1 g2 n1 n2
                go (A.LeftRight n1 n2)
  
  -- | Inlines all productions reachable via epsilon rules.
  epsilonClosure :: (IsGrammar n t) => Grammar n t -> Grammar n t
  epsilonClosure g = g { productions = Map.mapWithKey close (productions g) }
    where
      close name _ =
        Rhs { cons = lookup name g
            , eps = Set.empty
            }
  
  lookup :: forall n t. IsGrammar n t => n -> Grammar n t -> t n
  lookup n = lookup' ([n] :: [n])
  
  lookup' :: (Foldable f, IsGrammar n t) => f n -> Grammar n t -> t n
  lookup' n g@(Grammar _ p) = cons $ foldMap (\k -> fold (Map.lookup k p)) $ epsilonReachable n g
  
  -- | Computes the set of non-terminals reachable via epsilon rules.
  epsilonReachable :: (Foldable f, IsGrammar n t) => f n -> Grammar n t -> HashSet n
  epsilonReachable ns (Grammar _ p) = foldl' go Set.empty ns
    where
      go r n
        | Set.member n r = r
        | otherwise      = foldl' go (Set.insert n r) (eps (fold (Map.lookup n p)))
  
  -- | Removes productions that are not reachable from the start symbol.
  dropUnreachable :: (IsGrammar n t) => Grammar n t -> Grammar n t
  dropUnreachable g =
    let reach = reachable' (start g) g
    in g { productions = Map.filterWithKey (\k _ -> Set.member k reach) (productions g) }
  
  reachable :: (Foldable f, IsGrammar n t) => f n -> Grammar n t -> HashSet n
  reachable ns g = foldl' go Set.empty ns
    where
      go r n = let rhs = fold (Map.lookup n (productions g))
               in foldl' go (Set.insert n r) ((T.nonTerminals (cons rhs) <> eps rhs) `Set.difference` r)
  
  reachable' :: forall t n. IsGrammar n t => n -> Grammar n t -> HashSet n
  reachable' n = reachable ([n] :: [n])
  
  dropUnproductive :: IsGrammar n t => Grammar n t -> Grammar n t
  dropUnproductive g =
    g { productions = Map.mapMaybeWithKey
                    (\k r -> if Set.member k (productive g)
                              then Just (r { cons = T.filter (`Set.member` productive g) (cons r)
                                           , eps  = Set.intersection (eps r) (productive g) })
                              else Nothing)
                    (productions g)
      }
  
  isEmpty :: (Eq (t n), IsGrammar n t) => Grammar n t -> Bool
  isEmpty g = let g' = dropUnproductive g
              in lookup (start g') g' == mempty
  
  -- | Returns all productive nonterminals in the given grammar.
  productive :: forall n t. IsGrammar n t => Grammar n t -> HashSet n
  productive (Grammar _ prods) = go $ Set.fromList [ n | (n, rhs) <- Map.toList prods, T.productive mempty (cons rhs)]
    where
      go :: HashSet n -> HashSet n
      go p =
        let p' = Set.union p $ Set.fromList [ n | (n, rhs) <- Map.toList prods, rhsProductive rhs p ]
        in if p == p' then p else go p'
  
      -- A right hand side is productive if one of its epsilon rules is
      -- productive or a all non-terminals of a constructor are productive.
      rhsProductive :: Rhs n t -> HashSet n -> Bool
      rhsProductive rhs prod = any (`Set.member` prod) (eps rhs)
                            || T.productive prod (cons rhs)
  
  determinize :: (IsGrammar n t) => Grammar n t -> Grammar n t
  determinize = determinize'
  
  determinize' :: forall n n' t. (IsGrammar n t, IsGrammar n' t) => Grammar n t -> Grammar n' t
  determinize' g = generate $ do
    i <- get
    let (s',(_,(ps,_))) = runState (go [start g]) (i,(mempty,mempty))
    return $ Grammar s' ps
   where
     -- go {N1,N2,N3}
     --    [ N1 -> foo(A1,B1) | bar(B1)
     --      N2 -> foo(A2,B2) | bar(B2)
     --      N3 -> foo(A3,B3) | baz(B3)
     --      ...
     --    ]
     -- = [ N1 U N2 U N3 -> foo(A1 U A2 U A3, B1 U B2 U B3) | bar(B1 U B2) | biz(B3) ... ]
     go :: HashSet n -> State (Gen n', (Productions n' t,HashMap (HashSet n) n')) n'
     go ns = do
       rmap <- getRenameMap
       case Map.lookup ns rmap of
         Just x -> return x
         Nothing -> do
           n <- fresh' Nothing
           putRenameMap $ Map.insert ns n rmap
           ctrs <- T.determinize go (lookup' ns g)
           result $ Map.insert n (Rhs { cons = ctrs, eps = mempty})
           return n
  
  -- | Optimizes the grammar by joining non-terminals that produce the
  -- same language. This function has a runtime complexity /O(n*log(n))/
  -- where /n/ is the number of reachable non-terminals.
  minimize :: (IsGrammar n t) => Grammar n t -> Grammar n t
  minimize = minimize'
  
  minimize' :: forall n n' t. (IsGrammar n t, IsGrammar n' t) => Grammar n t -> Grammar n' t
  minimize' (dropUnreachable -> g) = generate $ do
    rmap <- foldM ( \m ns -> do
                      -- Map all members of the same equivalence class to a fresh symbol.
                      n' <- fresh Nothing
                      return $ foldl (\m' n -> Map.insert n n' m') m ns
                  ) Map.empty (equivalenceClasses g)
    return (mapNonTerms (rmap Map.!) g)
  
  mapNonTerms :: (IsGrammar n t, IsGrammar n' t) => (n -> n') -> Grammar n t -> Grammar n' t
  mapNonTerms f g = Grammar
    { start = f (start g)
    , productions = Map.fromList [ (f k, Rhs (T.map f (lookup k g)) mempty) | k <- Map.keys (productions g)]
    }
  
  -- | Returns the equivalence classes of non-terminals that produce the
  -- same language. This function has a runtime complexity /O(n*log(n))/
  -- where /n/ is the number of non-terminals.
  equivalenceClasses :: IsGrammar n t => Grammar n t -> HashMap n (Seq n)
  equivalenceClasses g = flip evalState O.empty $ U.runUnionFind $ do
    points <- Map.fromList <$> mapM (\n -> (n,) <$> U.fresh n) nonTerms
  
    forM_ (triangle nonTerms) $ \(n1,n2) -> do
      let g1 = g {start = n1}; p1 = points Map.! n1
          g2 = g {start = n2}; p2 = points Map.! n2
      b1 <- lift' (subsetOf' g1 g2)
      b2 <- lift' (subsetOf' g2 g1)
      when (b1 && b2) (U.union p1 p2)
  
    foldM (\equiv (n,p) -> do
        representative <- U.descriptor =<< U.repr p
        return $ Map.insertWith (<>) representative (pure n) equiv
      ) Map.empty (Map.toList points)
  
     where
       nonTerms = Map.keys (productions g)
  
       lift' :: (MonadTrans t) => (s -> (b,s)) -> t (State s) b
       lift' f = lift (state f)
  
       triangle :: [x] -> [(x,x)]
       triangle [] = []
       triangle (x:xs) = (x,x) : fmap (x,) xs ++ triangle xs
  
  toSubterms :: (IsGrammar n t, Monoid (t Int)) => Grammar n t -> t (Grammar n t)
  toSubterms g = T.map (\n -> dropUnreachable (g {start = n})) (lookup (start g) g)
  
  -- OPT: deduplicate grammar.
  fromSubterms :: forall n n' t. (IsGrammar n t, IsGrammar n' t) => t (Grammar n t) -> Grammar n' t
  fromSubterms t = generate $ do
    (s :: n') <- fresh Nothing
    (t',p) <- runWriterT $ T.traverse (\b -> do
                g <- lift (rename b)
                tell (productions g)
                return (start g)
              ) t
    return $ Grammar s $ Map.insert s (Rhs t' mempty) p
  
  identical :: (Eq n, Eq (t n)) => Grammar n t -> Grammar n t -> Bool
  identical (Grammar s1 ps1) (Grammar s2 ps2) = s1 == s2 && ps1 == ps2
  
  instance IsGrammar n t => Eq (Grammar n t) where
    -- First check if the grammars are identical. If this is not the case, perform
    -- the slower subset checks.
    b1 == b2 = subsetOf b1 b2 && subsetOf b2 b1
  
  type Salt = Int
  
  hashTerminals :: (IsGrammar n t, Monoid (t Int)) => Salt -> Grammar n t -> Int
  hashTerminals s (dropUnreachable -> Grammar _ prods) = runIdentity
    $ T.hashWithSalt (\h _ -> return h) s
    $ mconcat
    $ T.map (const (0 :: Int)) . cons <$> Map.elems prods
  
  type Depth = Int
  
  hashDepth ::  forall n t. (IsGrammar n t, Monoid (t Int)) => Depth -> Salt -> Grammar n t -> Int
  hashDepth depth s0 b =
    let g = determinize' b :: Grammar Int t
    in runReader (go g s0 (start g)) depth
    where
      go :: forall n' t'. IsGrammar n' t' => Grammar n' t' -> Int -> n' -> Reader Int Int
      go g s n = do
        fuel <- ask
        if fuel <= 0
        then return $ hashWithSalt s (0 :: Int)
        else local (\f -> f - 1) (T.hashWithSalt (go g) s (lookup n g))
  
  instance forall n t. (IsGrammar n t, Monoid (t Int)) => Hashable (Grammar n t) where
    hashWithSalt =
      hashTerminals
      -- hashDepth 3
  
  instance IsGrammar n t => PreOrd (Grammar n t) where
    (⊑) = subsetOf
  
  instance IsGrammar n t => Complete (Grammar n t) where
    g1 ⊔ g2 = minimize (union g1 g2 :: Grammar n t)
  
  instance IsGrammar n t => CoComplete (Grammar n t) where
    g1 ⊓ g2 = intersection g1 g2
  
  instance (UpperBounded (t String), Monoid (t String), IsGrammar n t) => UpperBounded (Grammar n t) where
    top = grammar "S" [("S",top)] []
  
  instance (Identifiable n, Semigroup (t n)) => Semigroup (Rhs n t) where
    Rhs c1 e1 <> Rhs c2 e2 = Rhs (c1 <> c2) (e1 <> e2)
  instance (Identifiable n, Monoid (t n)) => Monoid (Rhs n t) where
    mappend = (<>)
    mempty = Rhs { cons = mempty, eps = mempty }
  
  instance (Show n, Show (t n)) => Show (Grammar n t) where
    show (Grammar s p) = printf "{start = %s, prods = [%s]}" (show s)
                                (L.intercalate ", " [ printf "%s -> %s" (show n) (show t) | (n,t) <- Map.toList p ])
  
  instance (Show n, Show (t n)) => Show (Rhs n t) where
    show rhs = L.intercalate " | " $ show (cons rhs) : map show (Set.toList (eps rhs))
  
  instance (Identifiable n, IsEmpty n, IsEmpty (t n)) => IsEmpty (Grammar n t) where
    empty = Grammar empty (Map.singleton empty empty)
  
  instance IsEmpty (t n) => IsEmpty (Rhs n t) where
    empty = Rhs { cons = empty, eps = empty }
  
  ----- Helper Functions -----
  
  data RenameMap n1 n2 n' = RenameMap
    { ordMap1   :: OrdMap [n1] [n2]
    , ordMap2   :: OrdMap [n2] [n1]
    , renameMap :: HashMap (A.Either n1 n2) n'
    }
  
  emptyRenameMap :: RenameMap n1 n2 n'
  emptyRenameMap = RenameMap O.empty O.empty Map.empty
  
  getRenameMap :: State (a,(b,rmap)) rmap
  getRenameMap = gets (snd . snd)
  
  putRenameMap :: rmap -> State (a,(b,rmap)) ()
  putRenameMap r = modify (second (second (const r)))
  
  result :: (res -> res) -> State (a,(res,rmap)) ()
  result f = modify (second (first f))
  
  addProduction :: (NonTerminal n', Identifiable n1, Identifiable n2, Identifiable n') => A.Either n1 n2 -> State (Gen n', (Productions n' t, RenameMap n1 n2 n')) (t n') -> State (Gen n', (Productions n' t, RenameMap n1 n2 n')) n'
  addProduction n g = do
    m <- getRenameMap
    case Map.lookup n (renameMap m) of
      Just n' -> return n'
      Nothing -> do
        n' <- fresh' Nothing
        putRenameMap $ m { renameMap = Map.insert n n' (renameMap m) }
        ctrs <- g
        result $ Map.insert n' (Rhs { cons = ctrs, eps = mempty})
        return n'
  
  updateOrdMap :: (IsGrammar n1 t, IsGrammar n2 t) => Grammar n1 t -> Grammar n2 t -> n1 -> n2 -> State (a, (b, RenameMap n1 n2 n')) ()
  updateOrdMap g1 g2 n1 n2 = do
    m <- getRenameMap
    let g1' = g1 { start = n1 }
        g2' = g2 { start = n2 }
        (_,o1) = subsetOf' g1' g2' (ordMap1 m)
        (_,o2) = subsetOf' g2' g1' (ordMap2 m)
    putRenameMap $ m { ordMap1 = o1
                     , ordMap2 = o2
                     }
  