content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TemplateHaskellQuotes #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  -- |
  -- SPDX-License-Identifier: BSD-3-Clause
  --
  -- A random collection of small, useful functions that are (or could
  -- be) used throughout the code base.
  module Swarm.Util (
    -- * Miscellaneous utilities
    (?),
    sortPair,
    maxOn,
    maximum0,
    enumeratedMap,
    cycleEnum,
    enumerateNonEmpty,
    showEnum,
    indexWrapNonEmpty,
    chunksOfNE,
    rangeNE,
    uniq,
    binTuples,
    histogram,
    findDup,
    both,
    allEqual,
    tails1,
    prependList,
    deleteKeys,
    lookupEither,
    applyWhen,
    applyJust,
    unsnocNE,
    applyN,
  
    -- * Directory utilities
    readFileMay,
    readFileMayT,
    findAllWithExt,
    acquireAllWithExt,
  
    -- * Text utilities
    isIdentChar,
    replaceLast,
    failT,
    showT,
    showLowT,
  
    -- * English language utilities
    reflow,
    quote,
    squote,
    bquote,
    parens,
    brackets,
    commaList,
    indefinite,
    indefiniteQ,
    singularSubjectVerb,
    plural,
    number,
  
    -- * Validation utilities
    holdsOr,
    isJustOr,
    isRightOr,
    isSuccessOr,
  
    -- * Template Haskell utilities
    liftText,
  
    -- * Lens utilities
    (%%=),
    (<%=),
    (<+=),
    (<<.=),
    (<>=),
    _NonEmpty,
  
    -- * Set utilities
    removeSupersets,
    smallHittingSet,
  ) where
  
  import Control.Arrow ((***))
  import Control.Carrier.Throw.Either
  import Control.Effect.State (State, modify, state)
  import Control.Lens (ASetter', Lens', LensLike, LensLike', Over, lens, (<&>), (<>~))
  import Control.Monad (filterM, unless)
  import Data.Bifunctor (Bifunctor (bimap), first)
  import Data.Char (isAlphaNum, toLower)
  import Data.Either.Extra (maybeToEither)
  import Data.Either.Validation
  import Data.Foldable (Foldable (..))
  import Data.Foldable qualified as Foldable
  import Data.IntMap.Strict (IntMap)
  import Data.IntMap.Strict qualified as IM
  import Data.Ix (Ix, range, rangeSize)
  import Data.List (maximumBy, partition)
  import Data.List qualified as List
  import Data.List.Extra (enumerate)
  import Data.List.NonEmpty (NonEmpty ((:|)))
  import Data.List.NonEmpty qualified as NE
  import Data.Map (Map)
  import Data.Map qualified as M
  import Data.Maybe (fromMaybe)
  import Data.Ord (comparing)
  import Data.Set (Set)
  import Data.Set qualified as S
  import Data.Text (Text, toUpper)
  import Data.Text qualified as T
  import Data.Text.IO qualified as T
  import Data.Tuple (swap)
  import Data.Yaml
  import Language.Haskell.TH
  import Language.Haskell.TH.Syntax (lift)
  import NLP.Minimorph.English qualified as MM
  import NLP.Minimorph.Util ((<+>))
  import System.Clock (TimeSpec)
  import System.Directory (doesDirectoryExist, doesFileExist, listDirectory)
  import System.FilePath (takeExtension, (</>))
  import System.IO.Error (catchIOError)
  import Witch (from)
  import Prelude hiding (Foldable (..))
  
  infixr 1 ?
  infix 4 %%=, <+=, <%=, <<.=, <>=
  
  -- | A convenient infix flipped version of 'fromMaybe': @Just a ? b =
  --   a@, and @Nothing ? b = b@. It can also be chained, as in @x ? y ?
  --   z ? def@, which takes the value inside the first @Just@,
  --   defaulting to @def@ as a last resort.
  (?) :: Maybe a -> a -> a
  (?) = flip fromMaybe
  
  -- | Ensure the smaller value in a pair is the first element.
  sortPair :: Ord b => (b, b) -> (b, b)
  sortPair (x, y) = if x <= y then (x, y) else (y, x)
  
  -- | Find the maximum of two values, comparing them according to a
  --   custom projection function.
  maxOn :: Ord b => (a -> b) -> a -> a -> a
  maxOn f x y
    | f x > f y = x
    | otherwise = y
  
  -- | Find the maximum of a list of numbers, defaulting to 0 if the
  --   list is empty.
  maximum0 :: (Num a, Ord a) => [a] -> a
  maximum0 [] = 0
  maximum0 xs = maximum xs
  
  enumeratedMap :: Int -> [a] -> IntMap a
  enumeratedMap startIdx = IM.fromList . zip [startIdx ..]
  
  -- | Take the successor of an 'Enum' type, wrapping around when it
  --   reaches the end.
  cycleEnum :: (Eq e, Enum e, Bounded e) => e -> e
  cycleEnum e
    | e == maxBound = minBound
    | otherwise = succ e
  
  -- | See
  -- https://hackage.haskell.org/package/relude-1.2.1.0/docs/Relude-Enum.html#v:universeNonEmpty
  enumerateNonEmpty :: (Enum e, Bounded e) => NonEmpty e
  enumerateNonEmpty = minBound :| drop 1 enumerate
  
  -- | We know by the syntax rules of Haskell that constructor
  --  names must consist of one or more symbols!
  showEnum :: (Show e, Enum e) => e -> NonEmpty Char
  showEnum = NE.fromList . show
  
  -- | Guaranteed to yield an element of the list.
  --
  -- This is true even if the supplied @index@ is negative,
  -- since 'mod' always satisfies @0 <= a `mod` b < b@
  -- when @b@ is positive
  -- (see <comment https://github.com/swarm-game/swarm/pull/1181#discussion_r1151177735>).
  indexWrapNonEmpty ::
    Integral b =>
    NonEmpty a ->
    -- | index
    b ->
    a
  indexWrapNonEmpty list idx =
    NE.toList list !! fromIntegral wrappedIdx
   where
    wrappedIdx = idx `mod` fromIntegral (NE.length list)
  
  -- | Split the given list into chunks of the given size; the last
  --   chunk may be smaller.  If @n <= 0@, chunks of size 1 will be
  --   returned.
  chunksOfNE :: Int -> NonEmpty a -> NonEmpty (NonEmpty a)
  chunksOfNE n xs
    | n <= 0 = fmap NE.singleton xs
    | otherwise = NE.unfoldr takeN xs
   where
    takeN :: NonEmpty a -> (NonEmpty a, Maybe (NonEmpty a))
    takeN = (NE.fromList *** NE.nonEmpty) . NE.splitAt n
  
  -- Calling NE.fromList is safe here, since we know n > 0, so
  -- NE.splitAt n will generate a first chunk of length n > 0.
  
  -- | Like 'range', except that at least the first index is always
  --   returned, even if the indices are out of order.  For example,
  --   'range (4,0) = []' but 'rangeNE (4,0) = [4]'.
  rangeNE :: Ix a => (a, a) -> NonEmpty a
  rangeNE rng@(lo, _)
    | rangeSize rng == 0 = NE.singleton lo
    | otherwise = NE.fromList $ range rng
  
  -- | Drop repeated elements that are adjacent to each other.
  --
  -- >>> uniq []
  -- []
  -- >>> uniq [1..5]
  -- [1,2,3,4,5]
  -- >>> uniq (replicate 10 'a')
  -- "a"
  -- >>> uniq "abbbccd"
  -- "abcd"
  uniq :: Eq a => [a] -> [a]
  uniq = \case
    [] -> []
    (x : xs) -> x : uniq (dropWhile (== x) xs)
  
  -- | Place the second element of the tuples into bins by
  -- the value of the first element.
  binTuples ::
    (Foldable t, Ord a) =>
    t (a, b) ->
    Map a (NE.NonEmpty b)
  binTuples = foldr f mempty
   where
    f = uncurry (M.insertWith (<>)) . fmap pure
  
  -- | Count occurrences of a value
  histogram ::
    (Foldable t, Ord a) =>
    t a ->
    Map a Int
  histogram = foldl' (\m k -> M.insertWith (+) k 1 m) M.empty
  
  -- | Find a duplicate element within the list, if any exists.
  findDup :: Ord a => [a] -> Maybe a
  findDup = go S.empty
   where
    go _ [] = Nothing
    go seen (a : as)
      | a `S.member` seen = Just a
      | otherwise = go (S.insert a seen) as
  
  both :: Bifunctor p => (a -> d) -> p a a -> p d d
  both f = bimap f f
  
  allEqual :: (Ord a) => [a] -> Bool
  allEqual [] = True
  allEqual (x : xs) = all (== x) xs
  
  -- | Taken from here:
  -- https://hackage.haskell.org/package/ghc-9.8.1/docs/GHC-Data-FiniteMap.html#v:deleteList
  deleteKeys :: Ord key => [key] -> Map key elt -> Map key elt
  deleteKeys ks m = foldl' (flip M.delete) m ks
  
  -- | Convenience function to indicate which key
  -- was not found in the map.
  lookupEither :: Ord k => k -> Map k v -> Either k v
  lookupEither k = maybeToEither k . M.lookup k
  
  ------------------------------------------------------------
  -- Backported functions
  
  -- | The 'tails1' function takes a 'NonEmpty' stream @xs@ and returns all the
  -- non-empty suffixes of @xs@, starting with the longest.
  --
  -- > tails1 (1 :| [2,3]) == (1 :| [2,3]) :| [2 :| [3], 3 :| []]
  -- > tails1 (1 :| []) == (1 :| []) :| []
  --
  -- @since 4.18
  tails1 :: NonEmpty a -> NonEmpty (NonEmpty a)
  tails1 =
    -- fromList is an unsafe function, but this usage should be safe, since:
    -- \* `tails xs = [xs, tail xs, tail (tail xs), ..., []]`
    -- \* If `xs` is nonempty, it follows that `tails xs` contains at least one nonempty
    --   list, since `head (tails xs) = xs`.
    -- \* The only empty element of `tails xs` is the last one (by the definition of `tails`)
    -- \* Therefore, if we take all but the last element of `tails xs` i.e.
    --   `init (tails xs)`, we have a nonempty list of nonempty lists
    NE.fromList . Prelude.map NE.fromList . List.init . List.tails . Foldable.toList
  
  -- | Attach a list at the beginning of a 'NonEmpty'.
  -- @since 4.16
  prependList :: [a] -> NonEmpty a -> NonEmpty a
  prependList ls ne = case ls of
    [] -> ne
    (x : xs) -> x :| xs <> NE.toList ne
  
  -- Note, once we upgrade to an LTS version that includes
  -- base-compat-0.13, we should switch to using 'applyWhen' from there.
  applyWhen :: Bool -> (a -> a) -> a -> a
  applyWhen True f x = f x
  applyWhen False _ x = x
  
  -- |
  -- Equivalent to `fromMaybe id`.
  applyJust :: Maybe (a -> a) -> a -> a
  applyJust Nothing x = x
  applyJust (Just f) x = f x
  
  -- | Like 'unsnoc', but for 'NonEmpty' so without the 'Maybe'
  --
  -- Taken from Cabal-syntax Distribution.Utils.Generic.
  --
  -- Example:
  -- >>> import Data.List.NonEmpty (NonEmpty ((:|)))
  -- >>> unsnocNE (1 :| [2, 3])
  -- ([1,2],3)
  --
  -- >>> unsnocNE (1 :| [])
  -- ([],1)
  unsnocNE :: NonEmpty a -> ([a], a)
  unsnocNE (x :| xs) = go x xs
   where
    go y [] = ([], y)
    go y (z : zs) = let ~(ws, w) = go z zs in (y : ws, w)
  
  -- | Iteratively apply a function a specific number of times.
  --
  -- >>> applyN 4 succ 'a'
  -- 'e'
  -- >>> applyN 0 succ 'a'
  -- 'a'
  -- >>> applyN 5 (*2) 1
  -- 32
  applyN :: Int -> (a -> a) -> a -> a
  applyN 0 _ = id
  applyN n f = applyN (n - 1) f . f
  
  ------------------------------------------------------------
  -- Directory stuff
  
  -- | Safely attempt to read a file.
  readFileMay :: FilePath -> IO (Maybe String)
  readFileMay = catchIO . readFile
  
  -- | Safely attempt to (efficiently) read a file.
  readFileMayT :: FilePath -> IO (Maybe Text)
  readFileMayT = catchIO . T.readFile
  
  -- | Recursively acquire all files in the given directory with the
  --   given extension, but does not read or open the file like 'acquireAllWithExt'.
  findAllWithExt :: FilePath -> String -> IO [FilePath]
  findAllWithExt dir ext = do
    paths <- listDirectory dir <&> map (dir </>)
    filePaths <- filterM (\path -> doesFileExist path <&> (&&) (hasExt path)) paths
    -- recurse
    sub <- filterM doesDirectoryExist paths
    transChildren <- concat <$> mapM (`findAllWithExt` ext) sub
    return $ filePaths <> transChildren
   where
    hasExt path = takeExtension path == ("." ++ ext)
  
  -- | Recursively acquire all files in the given directory with the
  --   given extension, and their contents.
  acquireAllWithExt :: FilePath -> String -> IO [(FilePath, String)]
  acquireAllWithExt dir ext = findAllWithExt dir ext >>= mapM addContent
   where
    addContent :: FilePath -> IO (FilePath, String)
    addContent path = (,) path <$> readFile path
  
  -- | Turns any IO error into Nothing.
  catchIO :: IO a -> IO (Maybe a)
  catchIO act = (Just <$> act) `catchIOError` (\_ -> return Nothing)
  
  ------------------------------------------------------------
  -- Some Text-y stuff
  
  -- | Predicate to test for characters which can be part of a valid
  --   identifier: alphanumeric, underscore, or single quote.
  --
  -- >>> isIdentChar 'A' && isIdentChar 'b' && isIdentChar '9'
  -- True
  -- >>> isIdentChar '_' && isIdentChar '\''
  -- True
  -- >>> isIdentChar '$' || isIdentChar '.' || isIdentChar ' '
  -- False
  isIdentChar :: Char -> Bool
  isIdentChar c = isAlphaNum c || c == '_' || c == '\''
  
  -- | @replaceLast r t@ replaces the last word of @t@ with @r@.
  --
  -- >>> :set -XOverloadedStrings
  -- >>> replaceLast "foo" "bar baz quux"
  -- "bar baz foo"
  -- >>> replaceLast "move" "(make"
  -- "(move"
  replaceLast :: Text -> Text -> Text
  replaceLast r t = T.append (T.dropWhileEnd isIdentChar t) r
  
  -- | Fail with a Text-based message, made out of phrases to be joined
  --   by spaces.
  failT :: MonadFail m => [Text] -> m a
  failT = fail . from @Text . T.unwords
  
  -- | Show a value, but as Text.
  --
  --   Note: Data.Text.show was added in text-2.1.2.  Eventually we can
  --   require that version of text and get rid of showT.
  showT :: Show a => a -> Text
  showT = from @String . show
  
  -- | Show a value in all lowercase, but as Text.
  showLowT :: Show a => a -> Text
  showLowT = from @String . map toLower . show
  
  ------------------------------------------------------------
  -- Some language-y stuff
  
  -- | Reflow text by removing newlines and condensing whitespace.
  reflow :: Text -> Text
  reflow = T.unwords . T.words
  
  -- | Prepend a noun with the proper indefinite article (\"a\" or \"an\").
  indefinite :: Text -> Text
  indefinite w = MM.indefiniteDet w <+> w
  
  -- | Prepend a noun with the proper indefinite article, and surround
  --   the noun in single quotes.
  indefiniteQ :: Text -> Text
  indefiniteQ w = MM.indefiniteDet w <+> squote w
  
  -- | Combine the subject word with the simple present tense of the verb.
  --
  -- Only some irregular verbs are handled, but it should be enough
  -- to scrap some error message boilerplate and have fun!
  --
  -- >>> :set -XOverloadedStrings
  -- >>> singularSubjectVerb "I" "be"
  -- "I am"
  -- >>> singularSubjectVerb "he" "can"
  -- "he can"
  -- >>> singularSubjectVerb "The target robot" "do"
  -- "The target robot does"
  singularSubjectVerb :: Text -> Text -> Text
  singularSubjectVerb sub verb
    | verb == "be" = case toUpper sub of
        "I" -> "I am"
        "YOU" -> sub <+> "are"
        _ -> sub <+> "is"
    | otherwise = sub <+> (if is3rdPerson then verb3rd else verb)
   where
    is3rdPerson = toUpper sub `notElem` ["I", "YOU"]
    verb3rd
      | verb == "have" = "has"
      | verb == "can" = "can"
      | otherwise = fst $ MM.defaultVerbStuff verb
  
  -- | Pluralize a noun.
  plural :: Text -> Text
  plural = MM.defaultNounPlural
  
  -- For now, it is just MM.defaultNounPlural, which only uses heuristics;
  -- in the future, if we discover specific nouns that it gets wrong,
  -- we can add a lookup table.
  
  -- | Either pluralize a noun or not, depending on the value of the
  --   number.
  number :: Int -> Text -> Text
  number 1 = id
  number _ = plural
  
  -- | Surround some text in single quotes.
  squote :: Text -> Text
  squote t = T.concat ["'", t, "'"]
  
  -- | Surround some text in double quotes.
  quote :: Text -> Text
  quote t = T.concat ["\"", t, "\""]
  
  -- | Surround some text in backticks.
  bquote :: Text -> Text
  bquote t = T.concat ["`", t, "`"]
  
  -- | Surround some text in parentheses.
  parens :: Text -> Text
  parens t = T.concat ["(", t, ")"]
  
  -- | Surround some text in square brackets.
  brackets :: Text -> Text
  brackets t = T.concat ["[", t, "]"]
  
  -- | Make a list of things with commas and the word "and".
  commaList :: [Text] -> Text
  commaList [] = ""
  commaList [t] = t
  commaList [s, t] = T.unwords [s, "and", t]
  commaList ts = T.unwords $ map (`T.append` ",") (init ts) ++ ["and", last ts]
  
  ------------------------------------------------------------
  -- Some orphan instances
  
  deriving instance FromJSON TimeSpec
  deriving instance ToJSON TimeSpec
  
  ------------------------------------------------------------
  -- Validation utilities
  
  -- | Require that a Boolean value is @True@, or throw an exception.
  holdsOr :: Has (Throw e) sig m => Bool -> e -> m ()
  holdsOr b e = unless b $ throwError e
  
  -- | Require that a 'Maybe' value is 'Just', or throw an exception.
  isJustOr :: Has (Throw e) sig m => Maybe a -> e -> m a
  Just a `isJustOr` _ = return a
  Nothing `isJustOr` e = throwError e
  
  -- | Require that an 'Either' value is 'Right', or throw an exception
  --   based on the value in the 'Left'.
  isRightOr :: Has (Throw e) sig m => Either b a -> (b -> e) -> m a
  Right a `isRightOr` _ = return a
  Left b `isRightOr` f = throwError (f b)
  
  -- | Require that a 'Validation' value is 'Success', or throw an exception
  --   based on the value in the 'Failure'.
  isSuccessOr :: Has (Throw e) sig m => Validation b a -> (b -> e) -> m a
  Success a `isSuccessOr` _ = return a
  Failure b `isSuccessOr` f = throwError (f b)
  
  ------------------------------------------------------------
  -- Template Haskell utilities
  
  -- See https://stackoverflow.com/questions/38143464/cant-find-inerface-file-declaration-for-variable
  liftText :: T.Text -> Q Exp
  liftText txt = AppE (VarE 'T.pack) <$> lift (T.unpack txt)
  
  ------------------------------------------------------------
  -- Fused-Effects Lens utilities
  
  (<+=) :: (Has (State s) sig m, Num a) => LensLike' ((,) a) s a -> a -> m a
  l <+= a = l <%= (+ a)
  {-# INLINE (<+=) #-}
  
  (<%=) :: (Has (State s) sig m) => LensLike' ((,) a) s a -> (a -> a) -> m a
  l <%= f = l %%= (\b -> (b, b)) . f
  {-# INLINE (<%=) #-}
  
  (%%=) :: (Has (State s) sig m) => Over p ((,) r) s s a b -> p a (r, b) -> m r
  l %%= f = state (swap . l f)
  {-# INLINE (%%=) #-}
  
  (<<.=) :: (Has (State s) sig m) => LensLike ((,) a) s s a b -> b -> m a
  l <<.= b = l %%= (,b)
  {-# INLINE (<<.=) #-}
  
  (<>=) :: (Has (State s) sig m, Semigroup a) => ASetter' s a -> a -> m ()
  l <>= a = modify (l <>~ a)
  {-# INLINE (<>=) #-}
  
  ------------------------------------------------------------
  -- Other lens utilities
  
  _NonEmpty :: Lens' (NonEmpty a) (a, [a])
  _NonEmpty = lens (\(x :| xs) -> (x, xs)) (const (uncurry (:|)))
  
  ------------------------------------------------------------
  -- Some set utilities
  
  -- | Remove any sets which are supersets of other sets.  In other words,
  --   (1) no two sets in the output are in a subset relationship
  --   (2) every element in the input is a superset of some element in the output.
  --
  --   Note this can also be seen as minimizing a boolean expression in positive
  --   conjunctive normal form.
  --
  -- >>> import qualified Data.Set as S
  -- >>> rss = map S.toList . S.toList . removeSupersets . S.fromList . map S.fromList
  --
  -- >>> rss [[1,2,3], [1]]
  -- [[1]]
  --
  -- >>> rss [[1,2,3], [2,4], [2,3]]
  -- [[2,3],[2,4]]
  --
  -- >>> rss [[], [1], [2,3]]
  -- [[]]
  --
  -- >>> rss [[1,2], [1,3], [2,3]]
  -- [[1,2],[1,3],[2,3]]
  removeSupersets :: Ord a => Set (Set a) -> Set (Set a)
  removeSupersets ss = S.filter (not . isSuperset) ss
   where
    isSuperset s = any (`S.isSubsetOf` s) (S.delete s ss)
  
  -- | Given a list of /nonempty/ sets, find a hitting set, that is, a
  --   set which has at least one element in common with each set in the
  --   list.  It is not guaranteed to be the /smallest possible/ such
  --   set, because that is NP-hard.  Instead, we use a greedy algorithm
  --   that will give us a reasonably small hitting set: first, choose
  --   all elements in singleton sets, since those must necessarily be
  --   chosen.  Now take any sets which are still not hit, and find an
  --   element which occurs in the largest possible number of remaining
  --   sets. Add this element to the set of chosen elements, and filter
  --   out all the sets it hits.  Repeat, choosing a new element to hit
  --   the largest number of unhit sets at each step, until all sets are
  --   hit.  This algorithm produces a hitting set which might be larger
  --   than optimal by a factor of lg(m), where m is the number of sets
  --   in the input.
  --
  -- >>> import qualified Data.Set as S
  -- >>> shs = smallHittingSet . map S.fromList
  --
  -- >>> shs ["a"]
  -- fromList "a"
  --
  -- >>> shs ["ab", "b"]
  -- fromList "b"
  --
  -- >>> shs ["ab", "bc"]
  -- fromList "b"
  --
  -- >>> shs ["acd", "c", "aef", "a"]
  -- fromList "ac"
  --
  -- >>> shs ["abc", "abd", "acd", "bcd"]
  -- fromList "cd"
  --
  -- Here is an example of an input for which @smallHittingSet@ does
  -- /not/ produce a minimal hitting set. "bc" is also a hitting set and
  -- is smaller.  b, c, and d all occur in exactly two sets, but d is
  -- unluckily chosen first, leaving "be" and "ac" unhit and
  -- necessitating choosing one more element from each.
  --
  -- >>> shs ["bd", "be", "ac", "cd"]
  -- fromList "cde"
  smallHittingSet :: Ord a => [Set a] -> Set a
  smallHittingSet ss = go fixed (filter (S.null . S.intersection fixed) choices)
   where
    (fixed, choices) = first S.unions . partition ((== 1) . S.size) . filter (not . S.null) $ ss
  
    go !soFar [] = soFar
    go !soFar cs = go (S.insert best soFar) (filter (not . (best `S.member`)) cs)
     where
      best = mostCommon cs
  
    -- Given a nonempty collection of sets, find an element which is shared among
    -- as many of them as possible.
    mostCommon :: Ord a => [Set a] -> a
    mostCommon = fst . maximumBy (comparing snd) . M.assocs . histogram . concatMap S.toList
  