{-# LANGUAGE OverloadedStrings #-}

module Circuit.Parser.Token
  ( -- Tokenization
    tokenize,
    tokenizeLoop,
    -- Token patterns
    lowerLetter,
    upperLetter,
    letter,
    digit,
    word,
    number,
    punctuation,
    token,
    -- Vocabulary
    Vocabulary (..),
    buildVocabulary,
    lookupIndex,
    lookupToken,
    vocabularySize,
    filterVocabulary,
    takeTopN,
  )
where

import Circuit.Parser (Parser, These (..), Uncons (..), char, runParserIdentity, satisfy, some, (<|>))
import Data.Char (isAsciiLower, isAsciiUpper, isDigit)
import Data.Functor.Identity (Identity)
import Data.IntMap.Strict qualified as IntMap
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Text qualified as Text

-- ============================================================================
-- Token Patterns
-- ============================================================================

lowerLetter :: (Uncons f Char) => Parser Identity f Char Char
lowerLetter :: forall f. Uncons f Char => Parser Identity f Char Char
lowerLetter = (Char -> Bool) -> Parser Identity f Char Char
forall (m :: * -> *) f s.
(Monad m, Uncons f s) =>
(s -> Bool) -> Parser m f s s
satisfy Char -> Bool
isAsciiLower

upperLetter :: (Uncons f Char) => Parser Identity f Char Char
upperLetter :: forall f. Uncons f Char => Parser Identity f Char Char
upperLetter = (Char -> Bool) -> Parser Identity f Char Char
forall (m :: * -> *) f s.
(Monad m, Uncons f s) =>
(s -> Bool) -> Parser m f s s
satisfy Char -> Bool
isAsciiUpper

letter :: (Uncons f Char) => Parser Identity f Char Char
letter :: forall f. Uncons f Char => Parser Identity f Char Char
letter = Parser Identity f Char Char
forall f. Uncons f Char => Parser Identity f Char Char
lowerLetter Parser Identity f Char Char
-> Parser Identity f Char Char -> Parser Identity f Char Char
forall a.
Parser Identity f Char a
-> Parser Identity f Char a -> Parser Identity f Char a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser Identity f Char Char
forall f. Uncons f Char => Parser Identity f Char Char
upperLetter

digit :: (Uncons f Char) => Parser Identity f Char Char
digit :: forall f. Uncons f Char => Parser Identity f Char Char
digit = (Char -> Bool) -> Parser Identity f Char Char
forall (m :: * -> *) f s.
(Monad m, Uncons f s) =>
(s -> Bool) -> Parser m f s s
satisfy Char -> Bool
isDigit

word :: (Uncons f Char) => Parser Identity f Char [Char]
word :: forall f. Uncons f Char => Parser Identity f Char String
word = Parser Identity f Char Char -> Parser Identity f Char String
forall (m :: * -> *) f s a.
(Monad m, Uncons f s) =>
Parser m f s a -> Parser m f s [a]
some Parser Identity f Char Char
forall f. Uncons f Char => Parser Identity f Char Char
letter

number :: (Uncons f Char) => Parser Identity f Char [Char]
number :: forall f. Uncons f Char => Parser Identity f Char String
number = Parser Identity f Char Char -> Parser Identity f Char String
forall (m :: * -> *) f s a.
(Monad m, Uncons f s) =>
Parser m f s a -> Parser m f s [a]
some Parser Identity f Char Char
forall f. Uncons f Char => Parser Identity f Char Char
digit

punctuation :: (Uncons f Char) => Parser Identity f Char Char
punctuation :: forall f. Uncons f Char => Parser Identity f Char Char
punctuation = (Parser Identity f Char Char
 -> Parser Identity f Char Char -> Parser Identity f Char Char)
-> [Parser Identity f Char Char] -> Parser Identity f Char Char
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Parser Identity f Char Char
-> Parser Identity f Char Char -> Parser Identity f Char Char
forall a.
Parser Identity f Char a
-> Parser Identity f Char a -> Parser Identity f Char a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>) ((Char -> Parser Identity f Char Char)
-> String -> [Parser Identity f Char Char]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Parser Identity f Char Char
forall (m :: * -> *) f s.
(Monad m, Uncons f s, Eq s) =>
s -> Parser m f s s
char String
",.;:!?'\"-()[]{}@#$%&*+=<>/\\\\|~`")

token :: (Uncons f Char) => Parser Identity f Char [Char]
token :: forall f. Uncons f Char => Parser Identity f Char String
token = Parser Identity f Char String
forall f. Uncons f Char => Parser Identity f Char String
word Parser Identity f Char String
-> Parser Identity f Char String -> Parser Identity f Char String
forall a.
Parser Identity f Char a
-> Parser Identity f Char a -> Parser Identity f Char a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser Identity f Char String
forall f. Uncons f Char => Parser Identity f Char String
number Parser Identity f Char String
-> Parser Identity f Char String -> Parser Identity f Char String
forall a.
Parser Identity f Char a
-> Parser Identity f Char a -> Parser Identity f Char a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Char -> String)
-> Parser Identity f Char Char -> Parser Identity f Char String
forall a b.
(a -> b) -> Parser Identity f Char a -> Parser Identity f Char b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Char -> String -> String
forall a. a -> [a] -> [a]
: []) Parser Identity f Char Char
forall f. Uncons f Char => Parser Identity f Char Char
punctuation

-- ============================================================================
-- Core Tokenizer
-- ============================================================================

-- | Tokenize text into a list of tokens using Circuit.Parser patterns.
tokenize :: Text -> [Text]
tokenize :: Text -> [Text]
tokenize Text
input =
  let str :: String
str = Text -> String
Text.unpack Text
input
      tokens :: [String]
tokens = String -> [String] -> [String]
tokenizeLoop String
str []
   in (String -> Text) -> [String] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map String -> Text
Text.pack [String]
tokens

-- | Recursively extract tokens using runParser
tokenizeLoop :: String -> [String] -> [String]
tokenizeLoop :: String -> [String] -> [String]
tokenizeLoop [] [String]
acc = [String] -> [String]
forall a. [a] -> [a]
reverse [String]
acc
tokenizeLoop String
s [String]
acc =
  case Parser Identity String Char String -> String -> These String String
forall {k} f (s :: k) a. Parser Identity f s a -> f -> These a f
runParserIdentity Parser Identity String Char String
forall f. Uncons f Char => Parser Identity f Char String
token String
s of
    These String
tok String
rest ->
      if String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
tok
        then String -> [String] -> [String]
tokenizeLoop String
rest [String]
acc
        else String -> [String] -> [String]
tokenizeLoop String
rest (String
tok String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
acc)
    This String
tok ->
      if String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
tok
        then [String] -> [String]
forall a. [a] -> [a]
reverse [String]
acc
        else [String] -> [String]
forall a. [a] -> [a]
reverse (String
tok String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
acc)
    That String
_ ->
      case String
s of
        (Char
_ : String
rest) -> String -> [String] -> [String]
tokenizeLoop String
rest [String]
acc

-- ============================================================================
-- Vocabulary Building
-- ============================================================================

data Vocabulary = Vocabulary
  { Vocabulary -> Map Text Int
vocabTokenToIndex :: !(Map.Map Text Int),
    Vocabulary -> IntMap Text
vocabIndexToToken :: !(IntMap.IntMap Text),
    Vocabulary -> Int
vocabSize :: !Int
  }
  deriving (Vocabulary -> Vocabulary -> Bool
(Vocabulary -> Vocabulary -> Bool)
-> (Vocabulary -> Vocabulary -> Bool) -> Eq Vocabulary
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Vocabulary -> Vocabulary -> Bool
== :: Vocabulary -> Vocabulary -> Bool
$c/= :: Vocabulary -> Vocabulary -> Bool
/= :: Vocabulary -> Vocabulary -> Bool
Eq, Int -> Vocabulary -> String -> String
[Vocabulary] -> String -> String
Vocabulary -> String
(Int -> Vocabulary -> String -> String)
-> (Vocabulary -> String)
-> ([Vocabulary] -> String -> String)
-> Show Vocabulary
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Vocabulary -> String -> String
showsPrec :: Int -> Vocabulary -> String -> String
$cshow :: Vocabulary -> String
show :: Vocabulary -> String
$cshowList :: [Vocabulary] -> String -> String
showList :: [Vocabulary] -> String -> String
Show)

buildVocabulary :: [Text] -> Vocabulary
buildVocabulary :: [Text] -> Vocabulary
buildVocabulary [Text]
tokens =
  let (Map Text Int
fwd, IntMap Text
rev, Int
finalIdx) = ((Map Text Int, IntMap Text, Int)
 -> Text -> (Map Text Int, IntMap Text, Int))
-> (Map Text Int, IntMap Text, Int)
-> [Text]
-> (Map Text Int, IntMap Text, Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Map Text Int, IntMap Text, Int)
-> Text -> (Map Text Int, IntMap Text, Int)
forall {a}.
Ord a =>
(Map a Int, IntMap a, Int) -> a -> (Map a Int, IntMap a, Int)
insertUnique (Map Text Int
forall k a. Map k a
Map.empty, IntMap Text
forall a. IntMap a
IntMap.empty, Int
0) [Text]
tokens
   in Map Text Int -> IntMap Text -> Int -> Vocabulary
Vocabulary Map Text Int
fwd IntMap Text
rev Int
finalIdx
  where
    insertUnique :: (Map a Int, IntMap a, Int) -> a -> (Map a Int, IntMap a, Int)
insertUnique (Map a Int
fwd, IntMap a
rev, Int
idx) a
tok =
      if a -> Map a Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member a
tok Map a Int
fwd
        then (Map a Int
fwd, IntMap a
rev, Int
idx)
        else
          ( a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
tok Int
idx Map a Int
fwd,
            Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
idx a
tok IntMap a
rev,
            Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          )

-- ============================================================================
-- Vocabulary Lookups
-- ============================================================================

lookupIndex :: Text -> Vocabulary -> Maybe Int
lookupIndex :: Text -> Vocabulary -> Maybe Int
lookupIndex Text
tok Vocabulary
vocab = Text -> Map Text Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
tok (Vocabulary -> Map Text Int
vocabTokenToIndex Vocabulary
vocab)

lookupToken :: Int -> Vocabulary -> Maybe Text
lookupToken :: Int -> Vocabulary -> Maybe Text
lookupToken Int
idx Vocabulary
vocab = Int -> IntMap Text -> Maybe Text
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
idx (Vocabulary -> IntMap Text
vocabIndexToToken Vocabulary
vocab)

vocabularySize :: Vocabulary -> Int
vocabularySize :: Vocabulary -> Int
vocabularySize = Vocabulary -> Int
vocabSize

-- ============================================================================
-- Vocabulary Filtering
-- ============================================================================

filterVocabulary :: (Text -> Bool) -> Vocabulary -> Vocabulary
filterVocabulary :: (Text -> Bool) -> Vocabulary -> Vocabulary
filterVocabulary Text -> Bool
predicate Vocabulary
vocab =
  let filteredFwd :: Map Text Int
filteredFwd =
        (Text -> Int -> Bool) -> Map Text Int -> Map Text Int
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey
          (\Text
tok Int
_ -> Text -> Bool
predicate Text
tok)
          (Vocabulary -> Map Text Int
vocabTokenToIndex Vocabulary
vocab)
      tokenList :: [Text]
tokenList = Map Text Int -> [Text]
forall k a. Map k a -> [k]
Map.keys Map Text Int
filteredFwd
      (Map Text Int
newFwd, IntMap Text
newRev, Int
newSize) =
        ((Map Text Int, IntMap Text, Int)
 -> Text -> (Map Text Int, IntMap Text, Int))
-> (Map Text Int, IntMap Text, Int)
-> [Text]
-> (Map Text Int, IntMap Text, Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          ( \(Map Text Int
f, IntMap Text
r, Int
i) Text
tok ->
              (Text -> Int -> Map Text Int -> Map Text Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
tok Int
i Map Text Int
f, Int -> Text -> IntMap Text -> IntMap Text
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
i Text
tok IntMap Text
r, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          )
          (Map Text Int
forall k a. Map k a
Map.empty, IntMap Text
forall a. IntMap a
IntMap.empty, Int
0)
          [Text]
tokenList
   in Map Text Int -> IntMap Text -> Int -> Vocabulary
Vocabulary Map Text Int
newFwd IntMap Text
newRev Int
newSize

takeTopN :: Int -> Vocabulary -> Vocabulary
takeTopN :: Int -> Vocabulary -> Vocabulary
takeTopN Int
n Vocabulary
vocab =
  let tokensToKeep :: [Text]
tokensToKeep =
        [ Text
tok
        | Int
i <- [Int
0 .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Vocabulary -> Int
vocabSize Vocabulary
vocab Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)],
          Just Text
tok <- [Int -> IntMap Text -> Maybe Text
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
i (Vocabulary -> IntMap Text
vocabIndexToToken Vocabulary
vocab)]
        ]
      newFwd :: Map Text Int
newFwd = [(Text, Int)] -> Map Text Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([Text] -> [Int] -> [(Text, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
tokensToKeep [Int
0 ..])
      newRev :: IntMap Text
newRev = [(Int, Text)] -> IntMap Text
forall a. [(Int, a)] -> IntMap a
IntMap.fromList ([Int] -> [Text] -> [(Int, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Text]
tokensToKeep)
   in Map Text Int -> IntMap Text -> Int -> Vocabulary
Vocabulary Map Text Int
newFwd IntMap Text
newRev ([Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
tokensToKeep)