--------------------------------------------------------------------------------
-- |
-- Module      :  HarmTrace.Audio.Evaluation
-- 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: Basic parsers for parsing VAMP csv files.
--------------------------------------------------------------------------------

module HarmTrace.Audio.Evaluation (
    relCorrectOverlap, printRelCorrectOverlap
  ) where

import Constants
import HarmTrace.Audio.ChordTypes
import HarmTrace.Audio.Annotate (preProcessData, preProcessKeyData)
import HarmTrace.Base.MusicRep 

import Data.List (genericLength, zipWith5)
import Text.Printf(printf)
import System.IO (stdout,hFlush)

-- TODO this is a parameter and should some how be integrated into Constants.hs
-- this functions determines when two chords are considered the same
eqFunc :: ChordLabel -> ChordLabel -> Bool
eqFunc = chordTriadEq

--------------------------------------------------------------------------------
-- Evaluation functions
--------------------------------------------------------------------------------

-- chordClassEq :: ChordLabel -> ChordLabel -> Bool
-- chordClassEq (Chord (Note Nothing N) None _ _ _)
             -- (Chord (Note Nothing N) None _ _ _)       = True
-- chordClassEq (Chord (Note Nothing N) None _ _ _) _     = False
-- chordClassEq _     (Chord (Note Nothing N) None _ _ _) = False
-- chordClassEq (Chord r1 sh1 _ _ _) (Chord r2 sh2 _ _ _) =
  -- (toSemitone r1) == (toSemitone r2) && (toClassType sh1 == toClassType sh2)

chordTriadEq :: ChordLabel -> ChordLabel -> Bool
chordTriadEq (Chord (Note Nothing N) None _ _ _)
             (Chord (Note Nothing N) None _ _ _)       = True
chordTriadEq (Chord (Note Nothing N) None _ _ _) _     = False
chordTriadEq _     (Chord (Note Nothing N) None _ _ _) = False
chordTriadEq (Chord r1 sh1 _ _ _) (Chord r2 sh2 _ _ _) =
  toSemitone r1 == toSemitone r2 && toMode sh1 == toMode sh2

-- | Calculates the relative correct overlap, which is the recall
-- of matching frames, and defined as the nr of matching frames (sampled at
-- an 10 milisecond interval) divided by all frames.
relCorrectOverlap :: ChordAnnotation -> ChordAnnotation -> Double
relCorrectOverlap a b =
    foldl countMatch 0 (zipWith eqFunc sama samb) / tot where
    sama = sample a
    samb = sample b
    tot  = max (genericLength sama) (genericLength samb)

-- | does the same thing as relCorrectOverlap, but it also prints the
-- chords and uses a lower sample rate. N.B. the number output by 
-- 'printRelCorrectOverlap' might differ from the output of 
-- 'relCorrectOverlap', because a different sample rate might be used (see
-- 'Constants').
printRelCorrectOverlap :: (AudioFeat -> ChordBeatAnnotation) -> AudioFeat
                       -> ChordAnnotation -> IO Double
printRelCorrectOverlap annotator af@(AudioFeat chrm btbar afk) gt = do
  let keys = snd $ preProcessKeyData chrm btbar afk
      -- BUG:  now alswo when we are evaluating a simple annotator grouping is 
      --       is displayed, this is wrong. printRelCorrectOverlap should
      --       be independend of the kind of annotator.
      blks :: [BeatTimedData [ProbChord]]
      blks  = concatMap segChords $ preProcessData Nothing af

      -- sample the info for printing and evaluation
      samaf = sampleWith displaySampleRate (annotator af)
      samgt = sampleWith displaySampleRate gt
      sambk = sampleWith displaySampleRate blks
      samk  = sampleWith displaySampleRate keys

      tot   = max (genericLength samaf) (genericLength samgt)
      showEq m = if m then "==" else "/=" 
      printEval :: NumData -> ChordLabel -> ChordLabel -> Key -> [ProbChord] 
                -> IO Bool
      printEval t g a b c = 
         do putStrLn (printf "%.2f" t ++ '\t' : showEq equal ++ '\t' : show g
                      ++ '\t' : show a ++ '\t' : show b ++ '\t' : show c)
                      >> hFlush stdout
            return equal where equal = g `eqFunc` a
  putStrLn "time\tmatch\tGT\t\tMPTREE\tkey\toptional chords"
  m <- sequence (zipWith5 printEval [0.0,displaySampleRate ..] 
                                    samgt samaf samk sambk)
  return (foldl countMatch 0 m / tot)

countMatch :: Double -> Bool -> Double 
countMatch x y | y         = succ x -- count the number of matching frames
               | otherwise = x

-- given a chord annotation sample the chord label at every 10 ms
sample :: Timed t => [t a] -> [a]
sample = sampleWith evaluationSampleRate

-- like sample, but takes a sample rate (seconds :: Float) as argument
sampleWith :: Timed t => NumData -> [t a] -> [a]
sampleWith rate =  sampleAt [0.00, rate .. ] 

        
-- samples at specific points in time, specified in a list
sampleAt :: Timed t => [NumData] -> [t a] -> [a]
sampleAt  _  [] = [] -- below, will never occur
sampleAt []  _  = error "Harmtrace.Audio.Evaluation: No sampling grid specified" 
sampleAt (t:ts) (c:cs)
  | t <= offset c = getData c : sampleAt ts (c:cs)
  | otherwise     = sampleAt (t:ts) cs