{-# OPTIONS_GHC -Wall         #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  HarmTrace.Audio.ChromaChord
-- Copyright   :  (c) 2010-2012 Universiteit Utrecht, 2012 University of Oxford
-- License     :  GPL3
--
-- Maintainer  :  bash@cs.uu.nl, jpm@cs.ox.ac.uk
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Summary: Recognise audio chroma vectors into textual chord descriptions.
--------------------------------------------------------------------------------

module HarmTrace.Audio.ChromaChord ( createChordRanks, beatSync
                                   , mergeByBeat, addBeatTimeStamp
                                   , mergeAndTimeStamp, meanBeatSyncVectors
                                   , module Numeric.LinearAlgebra
                                   , module Numeric.GSL.Statistics 
                                   ) where

import Constants (maxProbChordListLength, cutOffProbability)
                                   
import HarmTrace.Audio.DataParser (shift)
import HarmTrace.Audio.ChordTypes
import HarmTrace.Base.MusicRep

-- import Text.Printf (printf)
import Data.List (sortBy,find) -- , elemIndices, minimumBy)
import Data.Maybe (isJust, fromJust, mapMaybe)
import Data.Ord (comparing)
-- import Data.Function (on)

-- http://hackage.haskell.org/package/hmatrix
import Numeric.LinearAlgebra hiding (find)
-- import qualified Numeric.GSL.Statistics as GSL hiding (mean)
import Numeric.GSL.Statistics  hiding (mean)

--------------------------------------------------------------------------------
-- Matching chords and chroma
--------------------------------------------------------------------------------  

-- | Wraps Chord candidate lists into a 'BeatTimedData' structure
addBeatTimeStamp :: [BeatBar] -> [[a]] -> [BeatTimedData [a]]
addBeatTimeStamp = beatTime (BeatBar (0,Four)) where

  beatTime :: BeatBar ->  [BeatBar] -> [[a]] -> [BeatTimedData [a]]
  beatTime _  [] [] = []
  beatTime (BeatBar (on, onbt)) (next@(BeatBar (off, _)) : bs) (x : xs) =
    BeatTimedData x onbt on off : beatTime next bs xs
  beatTime _  _  _  = error "addBeatTimeStamp:: asynchronous beats and data"

-- | Given a list of beats, a list of grouped data items, and a merging function
-- 'mergeAndTimeStamp' returns a list of 'BeatTimedData'. Before wrapping the
-- the grouped data items, e.g. chord candidates, the list is reduced by the 
-- provided merging function.
mergeAndTimeStamp  ::([a] -> a)-> [BeatBar] -> [[a]] -> [BeatTimedData a]
mergeAndTimeStamp f = merge (BeatBar (0,Four)) where
  -- merge :: BeatBar -> [BeatBar] -> [[a]] -> [BeatTimedData a]
  merge  _                   []    []       = []
  merge (BeatBar (on, bt)) beats (x : xs) = 
    let (off : rest) = drop (length x -1) beats
    in  BeatTimedData (f x) bt on (timeStamp off) : merge off rest xs
  merge _ _ _ = error "mergeAndTimeStamp: asynchronous beats and data"

-- | Merges chord segments, adding a bias toward merging at the first
-- and the third 'Beat' (specified by 'canMerge').
mergeByBeat :: [BeatTimedData [ProbChord]] -> [BeatTimedData [ProbChord]]
mergeByBeat [] = []
mergeByBeat [a] = [a]
mergeByBeat (x:y:xs) 
  | canMerge (getBeat x) (getBeat y) && isJust xy = mergeByBeat (fromJust xy:xs)
  | otherwise                                     = x : mergeByBeat (y:xs) where
      xy = merge x y  
      -- merges two Timed 'ProbChord's using intersectPC (currently in Utils.hs)
      merge :: BeatTimedData [ProbChord] -> BeatTimedData [ProbChord] 
            -> Maybe (BeatTimedData [ProbChord])
      merge a b 
        | not $ null m = Just (BeatTimedData m (getBeat a) (onset a) (offset b))
        | otherwise    = Nothing
          where m = intersectPC (getData a) (getData b)
      -- specifies which combinations of beats are allowed to merge
      canMerge :: Beat -> Beat -> Bool
      canMerge One   _ = True
      canMerge Three _ = True
      canMerge _     _ = False


intersectPC :: [ProbChord] -> [ProbChord] -> [ProbChord]
intersectPC a b = let (a',b') = order a b 
                  in mapMaybe (findAndMergePC a') b' where
  -- N.B. the probabilities are not divided by their length so > 1,
  -- due to the addition
  findAndMergePC :: [ProbChord] -> ProbChord -> Maybe ProbChord
  findAndMergePC pcs pc = case find (== pc) pcs of
    (Just pc') -> Just (ProbChord (chordLab pc) (prob pc + prob pc'))  
    Nothing    -> Nothing

  -- takes two lists and returns a tuple where the first element is the smallest
  -- and the second element is the largest of the two lists
  order  :: [a] -> [a] -> ([a],[a])
  order x y
    | length x <= length y = (x,y)
    | otherwise            = (y,x)

-- | Synchronises the 'ChordinoData' with the list of beats 
-- by grouping the 'ChordinoLines' of the 'ChordinoData' in separate lists.
beatSync :: BeatBarTrackData -> ChordinoData -> BeatChroma
beatSync _   []     = []
beatSync bt (cd:cs) = beatSync' (getBeatTrack bt) [cd] (cd:cs) where
  beatSync' :: [NumData] -> [ChordinoLine] -> [ChordinoLine] -> [[ChordinoLine]]
  beatSync' _  _   [] = []
  beatSync' [] _   c  = [c]
  beatSync' (b:bs) prv c -- we also store the previous group in case beat < time
    | null x    = prv : beatSync' bs prv xs 
    | otherwise = x   : beatSync' bs x   xs 
        where (x, xs) = span ((>=) b . time) c

--------------------------------------------------------------------------------
-- Matrix Functions for matching chords
-------------------------------------------------------------------------------- 

-- | Having a matrix of beat-synchronised bass and treble chromagrams and a 
-- chord dictionary, the probability of a chord sounding at a particular beat is
-- estimated by calculating the Euclidean distance between the chord structures 
-- and the chroma feature. These distances are calculated for every chord 
-- candidate at every beat. Next, we sort the chord candidates by descending 
-- Euclidean distance. To obtain a relative measure of the fit 
-- between a chord candidate and the chroma vector in the range [0,1],
-- the distances are normalised by dividing them by distance of the best 
-- matching chord candidate. 
createChordRanks :: BeatChroma -> [[ProbChord]]
createChordRanks = 
  map (selectTop . normalize . sortTake . matchCDictionary). meanBeatSyncVectors 
    where
    sortTake, normalize :: [ProbChord] -> [ProbChord]
    sortTake          = take maxProbChordListLength . sortBy (comparing prob)

    normalize l@(h:_) = let ph = prob h in map (\p -> p{prob = ph / prob p }) l 
    normalize []      = []
    
    selectTop l -- selects the everything with a probability > x
      | null s       = none 
      -- so far, this had not positive effect
      --   length s > maxProbChordListLength = none 
      | otherwise    = s
          where s    = takeWhile ((> cutOffProbability) . prob) l 
                none = [ProbChord noneLabel 0.0]

-- takes the mean of every "beat block" and these Vectors as one Matrix
-- Each row of this matrix corresponds to the chroma within one beat
meanBeatSyncVectors :: BeatChroma -> [Vector NumData] -- [[ChordinoLine]]
meanBeatSyncVectors = map mean . beatSyncMatrix

-- takes the median of every "beat block" and these Vectors as one Matrix
-- Each row of this matrix corresponds to the chroma within one beat
-- N.B. does not perform as well as meanBeatSyncVectors
-- medianBeatSyncVectors :: BeatChroma -> [Vector NumData]
-- medianBeatSyncVectors = 
   -- map (fromList . map GSL.median . toColumns) . beatSyncMatrix
 
-- creates a list of matrices, in which each matrix corresponds to the
-- collection of chroma vectors within one beat (drop the time stamp)
beatSyncMatrix :: BeatChroma -> [Matrix NumData]
beatSyncMatrix = map (dropColumns 1 . toChromaMatrix)

-- converts a ChordinoData into a Matrix
toChromaMatrix :: ChordinoData -> Matrix NumData
toChromaMatrix = fromLists . map mergeLine where
  mergeLine :: ChordinoLine -> [NumData]
  mergeLine (ChordinoLine tm bs tb) = tm : bs ++ tb    
  
-- matches all transpositions of a chord structure with a chroma vector
matchCDictionary :: Vector NumData -> [ProbChord]
matchCDictionary v = map (matchStruct v) chordDictionary

-- calculate an Euclidean (PNorm2) norm 
-- I also tried using the maximum norm, but this gave inferior results
matchStruct :: Vector NumData -> ChordCand -> ProbChord
matchStruct chroma (ChordCand r _ir None cs) = 
  ProbChord (Chord r None [] 0 1) (pnorm PNorm2 (chroma - fromList (cs ++ cs)))
matchStruct chroma (ChordCand r _ir sh cs) = 
  -- Chord root shorthand degrees description_str repetitions
  ProbChord (Chord r sh [] 0 1) (sqrt (bss * bss + treble * treble))
    where treble = pnorm PNorm2 (subVector 12 12 chroma - fromList cs)
          -- (_ir,bss) = minimumBy (compare `on` snd) . map matchInv $ rootInvs cs   
          (_ir,bss) = matchInv r   

          -- calculates the euclidean distance between the bass chromagram
          -- and all bass note inversions.
          matchInv :: Root -> (Root, NumData)
          matchInv ir = (ir, pnorm PNorm2 (subVector 0 12 chroma - bcs)) where
            bcs         = fromList (pre ++ [1] ++ tail post)
            (pre, post) = splitAt (toSemitone ir) (shortHandToCS None)

          -- For a given chord structure, compute all possible 
          -- roots for inversions
          -- rootInvs :: ChordStruct -> [Root]
          -- rootInvs cs = map iThRoot (elemIndices 1 cs)
          
--------------------------------------------------------------------------------
-- The Chord Dictionary
--------------------------------------------------------------------------------

-- the chord dictionary of all chords that are matched
chordDictionary :: [ChordCand]
chordDictionary = concatMap transpose12 [minBound..] -- all shorthands :-)
  
-- takes a ChordStruct and returns a list containing the 12 transposed versions
-- (including the original) of the ChordStruct, times the number of possible
-- inversions (3 or 4, depending on the number of notes in the shorthand).
transpose12 :: Shorthand -> [ChordCand]
transpose12 sh 
  | null cstruct = []
  | sh == None   = [ChordCand (Note Nothing N) (Note Nothing N) None cstruct]
  | otherwise    = concatMap mkChordCands [0..11]

  where
      cstruct :: ChordStruct
      cstruct = shortHandToCS sh

      mkChordCands :: Int -> [ChordCand]
      mkChordCands n = [ ChordCand r r sh shiftedCS]
                       where shiftedCS = shift (12-n) cstruct
                             r         = iThRoot n
      
-- Get the i-th chord root, for 0 <= i < 12
-- For i >= 12, the result is given modulo 12
-- JPM: I actually think this is not necessary because we know that
-- 0 <= i < 12, but it makes things safer anyway
iThRoot :: Int -> Root
iThRoot n | n < 0            = error "iThRoot: negative index"
          | n >= 0 && n < 12 = chromaPC !! n
          | otherwise        = iThRoot (n `mod` 12)

shortHandToCS :: Shorthand -> ChordStruct
shortHandToCS sh = case sh of 
  --        0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10,  11
  --        C,  Db, D,  Eb, E,  F,  F#, G,  Ab, A,  Bb,  B
  Maj   -> [1,  0,  0,  0,  1,  0,  0,  1,  0,  0,  0,   0 ] 
  Min   -> [1,  0,  0,  1,  0,  0,  0,  1,  0,  0,  0,   0 ]
  -- Dim   -> [1,  0,  0,  1,  0,  0,  1,  0,  0,  0,  0,   0 ]   
  -- HDim7 -> [1,  0,  0,  1,  0,  0,  1,  0,  0,  0,  1,   0 ]   
  -- Dim   -> [1,  0,  0,  1,  0,  0,  1,  0,  0,  0,  0,   0 ]   
  -- Maj6  -> [1,  0,  0,  0,  1,  0,  0,  1,  0,  1,  0,   0 ] 
  Sev   -> [1,  0,  0,  0,  1,  0,  0,  1,  0,  0,  1,   0 ] 
  -- Dim7  -> [1,  0,  0,  1,  0,  0,  1,  0,  0,  1,  0,   0 ] 
  -- Maj7  -> [1,  0,  0,  0,  1,  0,  0,  1,  0,  0,  0,   1 ] 
  -- Min7  -> [1,  0,  0,  1,  0,  0,  0,  1,  0,  0,  1,   0 ]
  -- Min6  -> [1,  0,  0,  1,  0,  0,  0,  1,  0,  1,  0,   0 ]
  -- Sus4  -> [1,  0,  0,  0,  0,  1,  0,  1,  0,  0,  0,   0 ] 
  
  None  -> [0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,   0 ] 
  -- None -> [1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,   1 ] 
  _     -> []
  -- none
  -- all
  -- ....

--------------------------------------------------------------------------------
-- general Matrix stuff
--------------------------------------------------------------------------------

-- given a matrix, calculates the mean vector 
mean :: (Product t, Fractional t) => Matrix t -> Vector t
mean a = constant (recip . fromIntegral . rows $ a) (rows a) <> a