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

-- $Id: Matching.hs 1260 2011-06-14 15:18:21Z bash $
module HarmTrace.Matching.Matching (getMatch, printBPM
                                   , getDownRight, wbMatch, wbMatchF
                                   , collectMatch, align, getWeightMatch
                                   ) where

import HarmTrace.Base.MusicRep
import HarmTrace.Matching.Sim
import HarmTrace.HAnTree.HAn 

import Data.Array
import Debug.Trace

--------------------------------------------------------------------------------
-- Parameters
--------------------------------------------------------------------------------   

inDel, matchW :: Float
inDel = -1
matchW = 4

max3 :: (Ord a) => (t, a) -> (t, a) -> (t, a) -> (t, a)    
max3 = lazyMax3
--------------------------------------------------------------------------------
-- Matching
--------------------------------------------------------------------------------    

-- prints a match
printBPM ::  [ChordLabel] -> [ChordLabel] -> IO()
printBPM t1' t2'  = putStrLn ("score: " ++ show simVal ++ '\n' :
                              "self sim 1: "++ show (maxSim t1) ++ '\n' :
                              "self sim 2: "++ show (maxSim t2) ++ '\n' :
                    algn t1 t2 (reverse match)) where
  -- hardcode C major for now ...
  t1 = map (toChordDegree (Key (Note Nothing C) MajMode)) t1'
  t2 = map (toChordDegree (Key (Note Nothing C) MajMode)) t2'
  -- (match, simVal) = getDownRight $ wbMatch t1 t2   
  tab    = trace (show t1 ++"\n" ++ show t2) (wbMatchF t1 t2)
  simVal = getDownRight $ tab
  match  = collectMatch tab
  -- algn :: Sim a => [a] -> [a] -> [(Int, Int)] -> [Char]
  algn a@(ha:ta) b@(hb:tb) m@((ma,mb):ms) 
    | matcha && matchb = show ha ++ "\t** "  ++ (show $ sim ha hb) ++ 
                                    " **\t"  ++ show hb ++ '\n':(algn ta tb ms)
    | matcha           = "             \t\t" ++ show hb ++ '\n':(algn a tb m)
    | matchb           = show ha ++ '\n'                       :(algn ta b m)
    | otherwise        = show ha ++ "\t\t" ++  show hb ++ '\n' :(algn ta tb m)
    where matcha = (getLoc ha) == ma
          matchb = (getLoc hb) == mb
  algn _         _         _             = ""  


-- returns a similarity value                              
getMatch :: Key -> [ChordLabel] -> [ChordLabel] -> Float
getMatch key ta tb = (weight * weight) / 
                     (maxSim ta' * maxSim tb' * matchW * matchW) where
  -- (_match,weight) = getWeightMatch ta' tb' 
  (_match,weight) = align ta' tb' 
  ta' = map (toChordDegree key) ta
  tb' = map (toChordDegree key) tb
 
-- selects the most lower right cell in the wbMatch' matrix
getWeightMatch :: (Sim a, GetDur a) => [a] -> [a] -> ([a], Float)
getWeightMatch  _ []    = ([],0)
getWeightMatch  [] _    = ([],0)
getWeightMatch  a   b   = (result,simVal) where
  (match, simVal) = getDownRight $ wbMatch a b 
  result          = snd . unzip $ filter (\x -> fst x `elem` mfst) (zip [0..] a)
  mfst            = reverse $ map fst match

align :: (Sim a, GetDur a) => [a] -> [a] -> ([a], Float)
align _  [] = ([],0.0)
align [] _  = ([],0.0)
align a  b  = (m, getDownRight t) where
  t  = wbMatchF a b
  cm = (map fst $ collectMatch t)
  m  = fst . unzip $ filter (\(_,x) -> x `elem` cm) (zip a [0..]) 
    
wbMatchF :: (Sim a, GetDur a) => [a] -> [a] -> Array (Int, Int) (Float)
wbMatchF _ []  = listArray ((0,0),(0,0)) (repeat 0.0)
wbMatchF [] _  = listArray ((0,0),(0,0)) (repeat 0.0)
wbMatchF 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        
  dura = listArray (0,la) (map (fromIntegral . getDur) a')
  durb = listArray (0,lb) (map (fromIntegral . getDur) b')
  match :: Int -> Int -> Float
  match i j = let s' = 2 * matchW * sim (a!i) (b!j) in 
              if  s' > 0 then s' else inDel * (min (dura!i) (durb!j)) -- durWeight
  -- inDelj = inDel * (fromIntegral $ getDur (b!j))
  -- this is the actual core recursive definintion of the algorithm
  recur i j = max3'(m!(i-1,j)   + inDel * (dura!i)) 
                   (m!(i-1,j-1) + match i j) 
                   (m!(i,j-1)   + inDel * (durb!j))
  m = array ((0,0),(la,lb)) 
    (((0,0), max (match 0 0) 0) :
    [((0,j), max0 (m!(0,j-1) + inDel * (durb!j)) (match 0 j)) | j <- [1..lb]] ++
    [((i,0), max0 (m!(i-1,0) + inDel * (dura!i)) (match i 0)) | i <- [1..la]] ++
    [((i,j), recur i j) | i <- [1..la], j <- [1..lb]])      

max3' :: (Ord a, Num a) => a -> a -> a -> a    
max3' a b c = max a (max0 b c)
-- max3' w nw n = if n > nw then n else max0 nw w -- not correct yet

max0  :: (Ord a, Num a) => a -> a -> a 
max0 a b = max a (max b 0)


collectMatch :: Array (Int, Int) Float -> [(Int,Int)]
collectMatch a = collect a (snd $ bounds a) []
collect :: Array (Int, Int) Float -> (Int,Int) -> [(Int,Int)] -> [(Int,Int)]
collect a c@(0,0) m = if a!c > 0 then c : m else m
collect a c@(i,0) m = if a!(i,0) > a!(i-1,0) then c : m else collect a (i-1,0) m
collect a c@(0,j) m = if a!(0,j) > a!(0,j-1) then c : m else collect a (0,j-1) m
collect a c@(i,j) m 
  | a!c > snd o = collect a (fst o) (c : m)
  | otherwise   = collect a (fst o) m where 
      o = realMax3 ((i-1,j)  , a!(i-1,j))
                   ((i-1,j-1), a!(i-1,j-1))
                   ((i,j-1)  , a!(i,j-1))
    
wbMatch :: (Sim a, GetDur 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 =  l where
    l = if   s > 0 
        then max3 (merge i j di (m!(i-1,j)))
                  (merge i j s  (m!(i-1,j-1)))
                  (merge i j dj (m!(i,j-1)))
        else max3 (m!(i-1,j)) (m!(i,j-1)) (m!(i-1,j-1))
    s = sim (a!i) (b!j) 
    di = inDel * getDurWeight (a!(i-1)) (b!j)
    dj = inDel * getDurWeight (a!i)     (b!(j-1))
  m =  array ((0,0),(la,lb)) 
    (((0,0), match 0 0) :
    [((0,j), maxByWeight (m!(0,j-1)) (match 0 j)) | j <- [1..lb]] ++
    [((i,0), maxByWeight (m!(i-1,0)) (match i 0)) | i <- [1..la]] ++
    [((i,j), concatMatch i j) | i <- [1..la], j <- [1..lb]])      

lazyMax3 :: (Ord a) => (t, a) -> (t, a) -> (t, a) -> (t, a)    
lazyMax3 w@(_,w') nw@(_,nw') n@(_,n') =  if n' > nw' then n else 
                                        (if w' > nw' then w else nw)    

realMax3 :: (Ord a) => (t, a) -> (t, a) -> (t, a) -> (t, a)    
realMax3 w nw n = maxByWeight nw (maxByWeight w n) where

maxByWeight :: Ord a => (t,a) -> (t,a) -> (t,a)                                        
maxByWeight 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.  
merge :: Int -> Int -> Float -> ([(Int,Int)], Float) -> ([(Int,Int)], Float)
merge 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