{-# LANGUAGE UnicodeSyntax, FlexibleInstances #-}
module Data.MarkovAlgo
  (Algo, Rule (..), Var (..),
   antecedent, consequent, constructor,
   expand,
   parseRule,
   buildAlgo,
   runMarkov)
  where

import Data.List

-- | Markov's algorithm itself
type Algo c = [Rule c]

-- | One rule in algorithm
data Rule c = [c] :-> [c]   -- ^ Non-terminating rule
            | [c] :->. [c]  -- ^ Terminating rule
  deriving (Eq)

instance Show (Rule Char) where
  show (a :-> s) = a ++ " → " ++ s
  show (a :->. s) = a ++ " →. " ++ s

instance Show (Rule (Var Char)) where
  show (a :-> s) = show a ++ " → " ++ show s
  show (a :->. s) = show a ++ " →. " ++ show s

-- | Get antecedent of rule
antecedent ∷ Rule c → [c]
antecedent (xs :-> _) = xs
antecedent (xs :->. _) = xs

-- | Get consequent of rule
consequent ∷ Rule c → [c]
consequent (_ :-> ys) = ys
consequent (_ :->. ys) = ys

-- | Get data constructor of Rule
constructor ∷ Rule c → ([d] → [d] → Rule d)
constructor (_ :-> _) = (:->)
constructor (_ :->. _) = (:->.)

-- | Variable for rules
data Var α = L α    -- ^ Literal char
           | V Int  -- ^ Variable with given number

instance Show (Var Char) where
  show (L c) = [c]
  show (V c) = show c

isChar ∷ Var α → Bool
isChar (L _) = True
isChar _     = False

isVar ∷ Var α → Bool
isVar = not . isChar

-- | Expand algorithm with variables into algorithm without variables
expand ∷ Eq α
       ⇒ [α]               -- ^ Alphabet
       → Algo (Var α)      -- ^ Algorithm with variables
       → Algo α
expand 𝔞 = concatMap (expandRule 𝔞)

expandRule ∷ Eq α ⇒ [α] → Rule (Var α) → [Rule α]
expandRule 𝔞 rule = nub $ map toChar (expandRule' 𝔞 rule)

toChar ∷ Rule (Var α) → Rule α
toChar rule = (constructor rule) (toChar' $ antecedent rule) (toChar' $ consequent rule)
  where
    toChar' ∷ [Var α] → [α]
    toChar' = map toChar''

    toChar'' (L c) = c
    toChar'' (V _) = error "Internal error"

expandRule' ∷ [α] → Rule (Var α) → [Rule (Var α)]
expandRule' 𝔞 rule = [(constructor rule) a c | (a,c) ← expandString 𝔞 (antecedent rule) (consequent rule)]

expandString ∷ [α] → [Var α] → [Var α] → [([Var α], [Var α])]
expandString 𝔞 as cs
  = map (subst as cs) $ mapM (const 𝔞) [1..n]
      where
        n = length (filter isVar (as ++ cs))

subst ∷ [Var α] → [Var α] → [α] → ([Var α], [Var α])
subst as cs xs = (subst' as xs, subst' cs xs)
  where
    subst' []        _      = []
    subst' (L c: vs) cs = L c: subst' vs cs
    subst' (V n: vs) cs = L (cs !! n): subst' vs cs
    subst' (V _: _)  [] = error "Internal error: too few arguments in `subst'!"

parseString ∷ Eq α ⇒ [α] → [α] → [Var α]
parseString xs s = map toVar s
  where
    toVar c = case elemIndex c xs of
                Nothing → L c
                Just n  → V n

-- | Create generic Rule from concrete Rule
parseRule ∷ Eq α 
          ⇒ [α]    -- ^ Names of variables
          → Rule α -- ^ Concrete rule
          → Rule (Var α)
parseRule xs rule = (constructor rule) (parseString xs $ antecedent rule) (parseString xs $ consequent rule)

replace ∷ Eq α ⇒ [α] → [α] → [α] → [α]
replace _ _ [] = []
replace old new s@(c:cs) 
  | old `isPrefixOf` s = new ++ (drop (length old) s)
  | otherwise          = c: replace old new cs

applyRule ∷ Eq α ⇒ Rule α → [α] → (Maybe [α], Bool)
applyRule rule s | a `isInfixOf` s = (Just $ replace a c s, shouldStop rule)
                 | otherwise       = (Nothing, shouldStop rule)
                     where
                      a = antecedent rule
                      c = consequent rule

shouldStop ∷ Rule c → Bool
shouldStop (_ :-> _) = False
shouldStop (_ :->. _) = True

-- | Run concrete (without variables) Markov's algorithm
runMarkov ∷ Eq α
          ⇒ Algo α       -- ^ Algorithm itself
          → [α]          -- ^ Start string
          → [α]
runMarkov algo s = runMarkov' algo algo s
  where
    runMarkov' :: Eq α ⇒ Algo α → Algo α → [α] → [α]
    runMarkov' _ [] s = s
    runMarkov' algo (r:rs) s = 
      case applyRule r s of
        (Just res, False) → runMarkov' algo algo res
        (Just res, True)  → res
        (Nothing,  _)     → runMarkov' algo rs s

-- | Build concrete algo from simple text description with variables
buildAlgo ∷ Eq α
          ⇒ [α]     -- ^ Alphabet
          → [α]     -- ^ Names of variables
          → Algo α  -- ^ Description of algorithm
          → Algo α
buildAlgo 𝔞 vs algo = expand 𝔞 $ map (parseRule vs) algo