{-# OPTIONS_GHC -Wall -fno-warn-orphans #-}

-- $Id: Matching.hs 1260 2011-06-14 15:18:21Z bash $
module HarmTrace.Matching.Matching (getSimMatch, getMatch, filterTrans
                             , collectMChord, printBPM, MatchChord (..)) where

import HarmTrace.Base.MusicRep
import HarmTrace.Tokenizer.Tokens
import HarmTrace.HAnTree.Tree
import HarmTrace.HAnTree.HAn
import HarmTrace.Matching.Sim

import Data.Maybe
import Data.Array
import Data.Ord
import Data.List (maximumBy)

--------------------------------------------------------------------------------
-- Matching
--------------------------------------------------------------------------------    

-- Top level functions:

-- prints a match
printBPM :: Int -> Tree HAn -> Tree HAn -> IO()
printBPM sampRate t1' t2'  = putStrLn ("score: " ++ show simVal ++ '\n' :
                              "self sim 1: "++ show (maxSim t1) ++ '\n' :
                              "self sim 2: "++ show (maxSim t2) ++ '\n' :
                    align t1 t2 (reverse match)) where
  t1              = collectMChord sampRate t1'
  t2              = collectMChord sampRate t2'
  (match, simVal) = getDownRight $ wbMatch t1 t2      
  align :: [MatchChord] -> [MatchChord] -> [(Int, Int)] -> [Char]
  align a@(ha:ta) b@(hb:tb) m@((ma,mb):ms) 
    | matcha && matchb = show ha ++ "\t** "  ++ (show $ sim ha hb) ++ 
                                    " **\t"  ++ show hb ++ '\n':(align ta tb ms)
    | matcha           = "             \t\t" ++ show hb ++ '\n':(align a tb m)
    | matchb           = show ha ++ '\n'                       :(align ta b m)
    | otherwise        = show ha ++ "\t\t" ++  show hb ++ '\n' :(align ta tb m)
    where matcha = (mgetLoc ha) == ma
          matchb = (mgetLoc hb) == mb
  align _         _         _             = ""  

-- returns a list of matched chords
getSimMatch :: Int -> Tree HAn -> Tree HAn -> [MatchChord]
getSimMatch sr ta tb = fst $ getWeightMatch 
                              (collectMChord sr ta) (collectMChord sr tb)

-- returns a similarity value                              
getMatch :: Int -> Tree HAn -> Tree HAn -> Float
getMatch sampleRate ta' tb' = (weight * weight) / (maxSim ta * maxSim tb) where
  (_match,weight) = getWeightMatch ta tb 
  ta = (collectMChord sampleRate ta') 
  tb = (collectMChord sampleRate tb')                               
                              
-- first argument is the sample rate, second argument is the parse tree
collectMChord :: Int -> Tree HAn -> [MatchChord]  
collectMChord sr ts = number $ collectMChord' 
                                (undefined :: HFunc) NoTrans ts where
  collectMChord' :: HFunc -> Trans -> Tree HAn -> [MatchChord]
  collectMChord'  f t (Node (HAnChord ct)  _     _) = 
    convChord sr f (filterTrans t) ct -- sample rate s
  collectMChord' _f t (Node (HAnFunc  hf) cs     _) = 
    concatMap (collectMChord' hf t) cs    
  collectMChord'  f t (Node (HAnTrans ht) (c:cs) _) = 
    collectMChord' f ht c ++ concatMap (collectMChord' f  t) cs    
  collectMChord'  f t (Node _             cs     _) = 
    concatMap (collectMChord' f  t) cs 
  -- the first argument r (Int) is the sample rate in beats. If it is set to one
  -- every beat is expanded to a chord. 
  convChord :: Int -> HFunc -> Trans -> ChordToken -> [MatchChord]
  convChord r f t  (ChordToken rt ct chds s _loc _dur) = 
    concatMap convert chds where
      convert :: ChordLabel -> [MatchChord]
      convert (Chord _ sh ad l d)   
        |d == 1    = [MChord rt sh ct ad s l 1 f t]
        |otherwise = replicate (round (fromIntegral d / fromIntegral r :: Float)) 
                  (MChord rt sh ct ad s l r f t)
  
  number :: [MatchChord] -> [MatchChord]
  number t  = number' t 0
  number' :: [MatchChord] -> Int -> [MatchChord]
  number' [] _ = []
  number' (MChord rt sh ct ad s _loc d f t: ns) ix = 
    (MChord rt sh ct ad s ix d f t) : number' ns (ix+1)

filterTrans :: Trans -> Trans
filterTrans s@(SecDom  _ _) = s
filterTrans s@(SecMin  _ _) = s
filterTrans s@(DiatDom _ _) = s
filterTrans s@(DimTrit _ _) = s
filterTrans _               = NoTrans
  
data MatchChord = MChord { _mchordRoot        :: ScaleDegree
                         , _mchordShorthand   :: Shorthand
                         , _mclassType        :: ClassType
                         , _mchordAdditions   :: [Addition]
                         , _mstatus           :: ParseStatus
                         , mgetLoc            :: Int -- the index of the chord  
                         , mduration          :: Int -- in the list of tokens
                         , _hfunc             :: HFunc
                         , _trans             :: Trans
                         } 
                         
instance Show MatchChord where
  show (MChord rt  sh  _clss  _add  _stat  lc  dr  f t ) = 
     show f ++ ':' : show t  ++ ':' : show rt ++ ':' : show sh 
            ++ '-' : show lc ++ ':' : show dr
               
instance Sim MatchChord where
  sim (MChord rt  _sh  clss  _add  _stat  _loc  dur1 _fnc  _trns )
      (MChord rt2 _sh2 clss2 _add2 _stat2 _loc2 dur2 _fnc2 _trns2) 
    =  (mChord) * (fromIntegral $ min dur1 dur2)  where
       {-  mFnc   = if fnc  == fnc2  then 0.15 else 0.0
         mTrns  = if stat /= Deleted && stat2 /= Deleted 
                  && trns == trns2   && trns  /= NoTrans then 0.25 else 0.0 -}
         mChord = if rt == rt2 && clss == clss2 then 1.0 else 0.0
        -- mDel   = if isNotDel  && mChord > 0    then 0.1 else 0.0
         -- mSh    = if mChord > 0 &&  sh == sh2  then 0.1 else 0.0
         --mStat  = if stat == stat2 && then 0.05 else 0.0
         --isNotDel  = stat /= Deleted && stat2 /= Deleted
         -- d      = fromIntegral $ min dur1 dur2

maxSim :: [MatchChord] -> Float
maxSim =  foldr (\a b -> sim a a + b) 0
  

  
-- selects the most lower right cell in the wbMatch' matrix
getWeightMatch :: [MatchChord] -> [MatchChord] -> ([MatchChord], Float)
getWeightMatch  _ []    = ([],0)
getWeightMatch  [] _    = ([],0)
getWeightMatch  a   b   = (result,simVal) where
  (match, simVal) = getDownRight $ wbMatch a b 
  mfst            = reverse $ map fst match
  result          = catMaybes $ map f a
  f x  
    | (mgetLoc x) `elem` mfst = Just x
    | otherwise   = Nothing


wbMatch :: Sim a => [a] -> [a] -> Array (Int, Int) ([(Int, Int)], Float)
wbMatch _ []  = listArray ((0,0),(0,0)) (repeat ([],0.0))
wbMatch [] _  = listArray ((0,0),(0,0)) (repeat ([],0.0))
wbMatch a' b' = m where
  la = length a'-1
  lb = length b'-1
  a  = listArray (0,la) a'  -- we need random access and therefore
  b  = listArray (0,lb) b'  -- convert the lists to arrays        
  match :: Int -> Int -> ([(Int,Int)],Float)
  match i j = if s > 0 then ([(i,j)],s) else ([],0) where s = sim (a!i) (b!j)
  -- this is the actual core recursive definintion of the algorithm
  concatMatch i j = maximumBy (comparing getWeight) l where
    l = if   s > 0 -- put the diagonal at the back to prefer symmetry
        then [ merge2 i j s (m!(i-1,j))  
             , merge2 i j s (m!(i,j-1)) 
             , merge2 i j s (m!(i-1,j-1))]
        else [m!(i-1,j), m!(i,j-1), m!(i-1,j-1)]
    s = sim (a!i) (b!j)    
  m = array ((0,0),(la,lb)) 
    (((0,0), match 0 0) :
    [((0,j), maxWeight (m!(0,j-1)) (match 0 j)) | j <- [1..lb]] ++
    [((i,0), maxWeight (m!(i-1,0)) (match i 0)) | i <- [1..la]] ++
    [((i,j), concatMatch i j) | i <- [1..la], j <- [1..lb]])      
   
maxWeight :: (a, Float) -> (a, Float) -> (a, Float)
maxWeight a@(_,wa) b@(_,wb) = if wa >= wb then a else b       
 
-- merges two tuples contianing the matchings, weight and cumulative depth of both
-- matched trees.  
merge2 :: Int -> Int -> Float -> ([(Int,Int)], Float) -> ([(Int,Int)], Float)
merge2 i j s p@(prv, w)
  | isFree prv i fst && isFree prv j snd = ((i,j) : prv, w + s) 
  | otherwise = p where
      isFree :: [a] -> Int -> (a -> Int) -> Bool 
      isFree prv' a f = null prv' || a > f (head prv')                  

--------------------------------------------------------------------------------
-- Some LCES helper functions
-------------------------------------------------------------------------------- 
getDownRight :: (Ix i) => Array i e -> e
getDownRight n = n ! snd (bounds n)

-- returns the weight of a match and is synonymous to snd
getWeight :: (a, b) -> b
getWeight (_,w) = w