content: -- | A wrapper for the fuzzyset library that makes it a Fuzzy map
  {-# LANGUAGE TemplateHaskell, DerivingVia, TupleSections #-}
  module Lamdu.Fuzzy
      ( Fuzzy, make, matches
      , Distance(..), isFuzzy, distanceInts
      , memoableMake
      ) where
  
  import qualified Control.Lens as Lens
  import           Control.Lens.Extended ((~~>))
  import           Data.Char (toLower)
  import           Data.List (sortOn, groupBy)
  import           Data.MMap (MMap)
  import qualified Data.Text as Text
  import           Data.Tuple (swap)
  import           Data.Vector ((!))
  import qualified Data.Vector as Vector
  import qualified Text.EditDistance as EditDistance
  import           Text.PrettyPrint (($+$), (<+>))
  import qualified Text.PrettyPrint as Pretty
  import           Text.PrettyPrint.HughesPJClass (Pretty(..))
  
  import           Lamdu.Prelude
  
  data Fuzzy a =
      Fuzzy (MMap Text a) (MMap Char (Fuzzy a))
      deriving stock Generic
      deriving (Semigroup, Monoid) via Generically (Fuzzy a)
  
  pPrintMMap :: (k -> Pretty.Doc) -> (v -> Pretty.Doc) -> MMap k v -> Pretty.Doc
  pPrintMMap pk pv m =
      m ^@.. Lens.ifolded
      <&> (\(k, v) -> pk k <+> " => " <+> pv v)
      & Pretty.vcat
  
  instance Pretty a => Pretty (Fuzzy a) where
      pPrint (Fuzzy xs m) =
          payload $+$ Pretty.nest 4 (pPrintMMap pPrint pPrint m)
          where
              payload = pPrintMMap (Pretty.text . Text.unpack) pPrint xs
  
  trieOf :: Semigroup a => Text -> String -> a -> Fuzzy a
  trieOf s [] x = Fuzzy (s ~~> x) mempty
  trieOf s (c:cs) x =
      Fuzzy mempty (toLower c ~~> cont) <> cont
      where
          cont = trieOf s cs x
  
  make :: Ord a => [(Text, a)] -> Fuzzy (Set a)
  make items =
      items
      <&> (\(text, x) -> trieOf text (Text.unpack text) (mempty & Lens.contains x .~ True))
      & mconcat
  
  -- | Takes the number of skips in given string that are allowed
  trieMatch :: Semigroup a => Int -> String -> Fuzzy a -> MMap Text a
  trieMatch _ [] (Fuzzy xs _) = xs
  trieMatch n (c:cs) t@(Fuzzy _ m) =
      (m ^.. Lens.ix (toLower c) <&> trieMatch n cs)
      ++ [trieMatch (n-1) cs t | n > 0] -- <-- try skipping a char
      & mconcat
  
  data Distance = Distance
      { _isFuzzy :: Bool
      , _distanceInts :: [Int]
      } deriving (Eq, Ord, Show)
  
  Lens.makeLenses ''Distance
  
  distance :: Text -> Text -> Distance
  distance rawX rawY =
      Distance
      { _isFuzzy = all (== 1) nonFuzzyScores
      , _distanceInts =
          nonFuzzyScores ++
          [ EditDistance.restrictedDamerauLevenshteinDistance
            EditDistance.defaultEditCosts (Text.unpack x) (Text.unpack y)
          ]
      }
      where
          nonFuzzyScores =
              concat
              [ check (==)
              , check Text.isPrefixOf
              , check Text.isInfixOf
              ]
          score True = 0
          score False = 1
          check f = map score [f rawX rawY, f x y]
          x = Text.toLower rawX
          y = Text.toLower rawY
  
  allowedSkips :: Int
  allowedSkips = 0
  
  matches :: Semigroup a => Text -> Fuzzy a -> MMap Text a
  matches = trieMatch allowedSkips . Text.unpack
  
  memoableMake ::
      ([(Text, Int)] -> Fuzzy (Set Int)) -> [([Text], a)] -> Text -> [(Distance, a)]
  memoableMake memoMake pairs =
      -- Keep as explicit lambda syntax (not in LHS) for
      -- performance/sharing reasons (where clauses do not depend on
      -- last Text param)
      \text ->
      matches text f ^@.. Lens.ifolded
      <&> _1 %~ distance text
      & sortOn swap & groupBy (\(_, x) (_, y) -> x == y) <&> head
      & sortOn fst
      >>= flatten
      <&> _2 %~ (v !)
      where
          -- Use const cause lower scores are better, use those for dups
          flatten (dist, indices) = indices ^.. Lens.folded <&> (,) dist
          f =
              pairs ^@.. Lens.ifolded
              >>= (\(idx, (texts, _)) -> texts <&> (, idx))
              & memoMake
          v = pairs <&> snd & Vector.fromList
  