module Core where

import           Data.List (maximumBy)
import qualified Data.Set  as Set
import qualified Data.Text as T

type Base = T.Text
type Suffix = T.Text
type Segment = T.Text

type SWord = T.Text

data Stemmer = Stemmer
  { wordList :: !(Set.Set SWord)
  , bases    :: !(Set.Set Base)
  , suffixes :: !(Set.Set Suffix)
  } deriving Show

defaultStemmer :: Stemmer
defaultStemmer = Stemmer
  { wordList = Set.empty
  , bases = Set.empty
  , suffixes = Set.empty
  }

suffixFrequency :: Stemmer -> Suffix -> Int
suffixFrequency stemmer suffix = Set.size . Set.filter (`T.isSuffixOf` suffix) . wordList $ stemmer

segmentFrequency :: Stemmer -> Segment -> Int
segmentFrequency stemmer segment = Set.size
                                 . Set.filter (T.isInfixOf segment)
                                 . wordList $ stemmer

nonFinalFrequency :: Stemmer -> Segment -> Int
nonFinalFrequency stemmer segment = segFreq - sufFreq
  where segFreq = segmentFrequency stemmer segment
        sufFreq = suffixFrequency stemmer segment

curveDrop :: Stemmer -> Suffix -> Double
curveDrop stemmer suffix
  | actualSuffixFrequency == 0 = 0
  | otherwise = maximum . map curveDropFunc . prependHelper $ suffix
  where prependHelper suf      = [ T.concat [T.singleton c, suf] | c <- charecterRange]
        charecterRange         = ['a' .. 'z']
        curveDropFunc suf      = (1 - (newSuffixFrequency suf / actualSuffixFrequency)) / bestPercentageDrop
        newSuffixFrequency suf = fromIntegral . suffixFrequency stemmer $ suf
        actualSuffixFrequency  = fromIntegral $ suffixFrequency stemmer suffix
        bestPercentageDrop     = 1 - (1 / fromIntegral (length charecterRange))

randomAdjustment :: Stemmer -> Suffix -> Double
randomAdjustment stemmer suffix
  | nonFinalProbabilty stemmer suffix > 0 = suffixProbabilty stemmer suffix / nonFinalProbabilty stemmer suffix
  | otherwise = 1

suffixProbabilty :: Stemmer -> Suffix -> Double
suffixProbabilty stemmer suffix = ps / psum
  where ps   = fromIntegral (suffixFrequency stemmer suffix)
        psum = fromIntegral (sum . Set.toList . Set.map (suffixFrequency stemmer) . suffixes $ stemmer)

nonFinalProbabilty :: Stemmer -> Segment -> Double
nonFinalProbabilty stemmer suffix = ps / psum
  where ps   = fromIntegral (nonFinalFrequency stemmer suffix)
        psum = fromIntegral (sum . Set.toList . Set.map (nonFinalFrequency stemmer) . suffixes $ stemmer)

score :: Stemmer -> Segment -> Double
score stemmer segment = c * ra * f
  where c  = curveDrop stemmer segment
        ra = randomAdjustment stemmer segment
        f  = fromIntegral $ suffixFrequency stemmer segment

mostLikelySWord :: T.Text -> Stemmer -> (SWord, Base, Suffix)
mostLikelySWord word stemmer = toSWord . maximumBy maxFunc . map (\(b, s) -> (b, s, score stemmer s)) $ sWords word
  where maxFunc (_, _, f1) (_, _, f2) = compare f1 f2
        toSWord (b, s, _) = (T.concat [b, s], b, s)

sWords :: T.Text -> [(Base, Suffix)]
sWords word = map (`T.splitAt` word) [0 .. T.length word]

addSWord :: Stemmer -> (SWord, Base, Suffix) -> Stemmer
addSWord stemmer@(Stemmer wl bs ss) (sWord, b, s) = stemmer
  { wordList = Set.insert sWord wl
  , bases = Set.insert b bs
  , suffixes = Set.insert s ss
  }
