content: module Markov where
  
  import           MusicData
  import           Utility
  
  -- import           Numeric.LinearAlgebra (Matrix, R, (><))
  
  import           Data.Map              (Map)
  import qualified Data.Map              as Map (empty, insertWith, lookup)
  import qualified Data.Map.Merge.Strict as Map' (merge, preserveMissing,
                                                  zipWithMatched)
  import qualified Data.Map.Strict       as Strict (Map)
  import qualified Data.Map.Strict       as Map' (difference, elems, empty,
                                                  fromList, insert, insertWith,
                                                  keys, lookup, member, toList)
  import qualified Data.Maybe            as Maybe (fromMaybe)
  import qualified Data.Set              as Set (fromList, size, toList)
  
  import Control.Parallel.Strategies (parMap, rseq)
  import qualified Data.Map.Strict as StrictMap
  import qualified Data.Set as Set
  
  -- |representation of bigrams and containing deterministic cadence sequences
  type Bigram  = (Cadence, Cadence)
  
  -- |representation of counts for each trigram
  type TransitionCounts = Strict.Map Bigram Double
  
  -- -- |representations of the Markov transition matrix
  -- type TransitionMatrix = Matrix R
  
  -- |representation of markov transition matrix as key-value pairs
  type MarkovMap = Map Cadence [(Cadence, Double)]
  
  -- |mapping from list of events into list of existing preceding bigrams
  bigrams :: [a] -> [(a, a)]
  bigrams (x:xs)
    | length (x:xs) < 2 = []
    | otherwise = bigram (x:xs) : bigrams xs
    where bigram (x:y:ys) = (\a b -> (a, b)) x y
  
  -- |lifted 'shortcut' `toCadence` which operates on a list of integer lists
  toCadences    :: (Integral a, Num a) => [[a]] -> [Cadence]
  toCadences xs = toCadence <$> bigrams (flatTriad <$> xs)
  
  -- |mapping from input data into all theoretically possible bigrams
  pairs   :: [Cadence] -> [Bigram]
  pairs xs =
    let cs  = unique xs
     in [ (x, y) | x <- cs, y <- cs ]
  
  -- |mapping from input data into all possible trigrams with counts of zero
  zeroCounts             :: [Cadence] -> TransitionCounts
  zeroCounts xs           =
    let mInsert acc key = Map'.insert key 0 acc
     in foldl mInsert Map'.empty $ pairs xs
  
  -- |mapping from input data to counts of all occurring transitions
  cadenceCounts        :: [Cadence] -> TransitionCounts
  cadenceCounts xs      =
    let mInsert acc key = Map'.insertWith (+) key 1 acc
     in foldl mInsert Map'.empty $ bigrams xs
  
  ---- |mapping from input to counts of cadences, including 'stationary' movements
  --transitionCounts       :: [Cadence] -> TransitionCounts
  --transitionCounts xs     = mergeMaps (foldl mInsert cadences (Map'.keys diff)) zeros
  --  where diff            = Map'.difference zeros cadences
  --        zeros           = zeroCounts xs
  --        cadences        = cadenceCounts xs
  --        mergeMaps m1 m2 = Map'.merge Map'.preserveMissing Map'.preserveMissing
  --                          (Map'.zipWithMatched (\k x y -> x)) m1 m2
  --        keys k          = [ (fst k, nxt) | nxt <- unique xs ]
  --        newKey k        = (fst k, fst k)
  --        member k        = sequenceA [ f ks | f <- [Map'.member], ks <- keys k ] (cadences)
  --        mInsert acc key
  --          | all (\x -> x == False) (member key) == True =
  --              Map'.insert (newKey key) 1 acc
  --          | otherwise   =
  --              Map'.insert key (Maybe.fromMaybe 0 $ Map'.lookup key $ cadences) acc
  
  transitionCounts :: [Cadence] -> TransitionCounts
  transitionCounts xs =
    let bigrams = parMap rseq (\x -> (x, 1)) $ Markov.bigrams xs
        bigramMap = StrictMap.fromListWith (+) bigrams
        zeroCounts = StrictMap.fromList $ zip (pairs xs) (repeat 0)
    in StrictMap.unionWith (+) bigramMap zeroCounts
  
  -- |helper function for probabilityList which generates probability sublists
  transitionProbs      :: [Cadence] -> [Double] -> [[Double]]
  transitionProbs _ []  = []
  transitionProbs xs ys =
    let j               = Set.size $ Set.fromList xs
        i               = j^2
        xx              = take j ys
        recurse         = drop j ys
     in fmap (/ sum xx) xx : transitionProbs xs recurse
  
  -- |mapping from list of Cadences into list of transitions with probabilities
  probabilityMap   :: [Cadence] -> Map Bigram Double
  probabilityMap xs = Map'.fromList $ zip (Map'.keys $ zeroCounts xs) $ concat
                       . transitionProbs xs $ Map'.elems $ transitionCounts xs
  
  -- -- |mapping from list of cadences into transition matrix
  -- transitionMatrix   :: [Cadence] -> TransitionMatrix
  -- transitionMatrix xs =
  --   let n             = Set.size $ Set.fromList xs
  --    in (n><n) $ Map'.elems $ probabilityMap xs :: Matrix R
  
  -- |mapping from list of cadences into map with possible next probabilities
  markovMap              :: [Cadence] -> MarkovMap
  markovMap xs            = foldl mInsert Map.empty $ pairs xs
    where mInsert acc key = Map.insertWith (++) (fst key) (pList key) acc
          pList key       = [(snd key, Maybe.fromMaybe 0 $ Map.lookup key pMap)]
          pMap            = probabilityMap xs