{-# LANGUAGE FlexibleContexts #-}
{-# OPTIONS_GHC -Wall         #-}

module HarmTrace.Audio.Key where

import HarmTrace.Audio.ChordTypes
import HarmTrace.Audio.Utils
import HarmTrace.Audio.BeatChroma ( beatSyncKeyStrenth, keyMap)
import HarmTrace.Base.MusicRep

import Prelude as P hiding ( map, length, head, last, mapM, mapM_, max
                           , maximum, reverse, tail, null, concatMap )

-- N.B. Vector inside this module refers to a different type than Vector
--      in the HarmTrace.Audio.BeatChroma module
import Data.Vector as V 
import qualified Data.List as L

import Data.Ord (comparing)
import Text.Printf (printf)

--------------------------------------------------------------------------------
-- Chroma key estimation
--------------------------------------------------------------------------------
modulPenalty :: NumData
modulPenalty = 1.0

getBeatSyncKeyFromChroma :: [NumData] -> [ChordinoLine] -> [[Key]]
getBeatSyncKeyFromChroma bts key = 
  groupKeys . getKeyFromTable $ selectKey bts key

selectKey :: [NumData] -> [ChordinoLine] ->  Vector (Vector (Int, NumData))
selectKey _bts  []  = empty 
selectKey  []  _key = empty
selectKey  bts  key = k where
  -- start by calculating the beat synchronised key strenght for all 
  -- 24 keys (ordered by HarmTrace.Audio.BeatChroma.keyMap)
  m :: Vector (Vector NumData)
  m = fromList . L.map fromList $ beatSyncKeyStrenth bts key
  
  -- calculate for every beat the maximum key (the index) and the 
  -- profile correlation (snd)
  maxima :: Vector (Int, NumData)
  {-# INLINE maxima  #-}
  maxima = map (\x -> (maxIndex x, maximum x)) m
  
  -- we fill a beat x 24 table and store the cumulative key strength.
  -- we can chose to stay in the current key or we can modulate which is
  -- penalised by modulPenalty, we also store the index so we can follow
  -- the path back to the first beat
  fill :: Int -> Int -> (Int, NumData)
  {-# INLINE fill  #-}
  fill 0 j = (j, (m!0)!j)
  fill i j = let (mj, mv) = maxima!i -- current max
                 noModul  = (j , (snd ((k!(i-1))!j)) + ((m!i)!j))
                 modul    = (mj, (snd ((k!(i-1))!j)) + mv - modulPenalty)
             in max2 modul noModul 

  k = generate (length m) (generate 24 . fill)

max2 :: (Int, NumData) -> (Int, NumData) -> (Int, NumData)
{-# INLINE max2  #-}
max2 t1@(_, s1) t2@(_, s2) = if s1 > s2 then t1 else t2

getKeyFromTable :: Vector (Vector (Int, NumData)) -> [Key]
getKeyFromTable k = L.map ((!!) keyMap) (L.reverse yek) where
  yek   = collectMax (fst $ maximumBy (comparing snd) (last k)) (reverse k)
  -- given the table calulated with selectKey, this function calculates 
  -- the actual key assignment for every beat
  collectMax :: Int -> Vector (Vector (Int, NumData)) -> [Int]
  collectMax startj l 
    | null l     = []
    | otherwise  = fst ((head l) ! startj) : collectMax m (tail l) where
        m = fst $ maximumBy (comparing snd) (head l)

printKeyTable :: [NumData] -> [ChordinoLine] -> IO ()
printKeyTable bts chrm = 
  let showLn :: Vector (Int, NumData) -> IO ()
      showLn x = do mapM_ (\(i,f) -> putStr (printf "(%d, %.2f)" i f)) x
                    putStr "\n"
  in mapM_ showLn $ selectKey bts chrm 

naiveBeatSyncKey :: BeatTrackerData -> [ChordinoLine] -> [Key]
naiveBeatSyncKey bts key = 
  L.map (((!!) keyMap) . maxListIndex) $ beatSyncKeyStrenth bts key

--------------------------------------------------------------------------------
-- key strengthpParsing
--------------------------------------------------------------------------------

groupKeys :: [Key] -> [[Key]]
groupKeys ks = L.group . L.concat $ groupMinSize 16 (getGlobalKey ks) ks 

getGlobalKey :: [Key] -> Key
getGlobalKey = mode