{-# OPTIONS_GHC -Wall #-}

module HarmTrace.Matching.FlatMatch where 

import Data.Array
import Data.List

-- import HarmTrace.Base.MusicRep
import HarmTrace.HAnTree.Tree
import HarmTrace.Matching.GuptaNishimura (getDownRight)

-- Toplevel:
-- Returns a Maximum Included Subtree (MIS). This is the three that maximizes 
-- the number of matching nodes in a tree respecting the order and 
-- ancestorship relations. The size of a MIS is always greater then or equal
-- to the size of the LCES of Gupta and Nishimura. 
getFlatMatch  :: (Eq t) => Tree t -> Tree t -> [Tree t]
getFlatMatch  ta tb = buildLCES (matchToTree ta ma) (matchToTree tb mb) where
  (ma,mb) = getFlatMatch' ta tb
getFlatMatch' :: (Eq a) => Tree a -> Tree a -> ([Int], [Int])
getFlatMatch' ta tb = unzip $ reverse $ fst3 $ getDownRight $ wbMatch (potPret ta) (potPret tb) simLab
      

buildLCES :: [Tree t] -> [Tree t] -> [Tree t]
buildLCES [] [] = []
buildLCES [] _  = error "buildLCES error"
buildLCES _  [] = error "buildLCES error"
buildLCES a@(ta:tas) b@(tb:tbs)
  | size ta > size tb = buildLCES (flatten a) b
  | size ta < size tb = buildLCES a (flatten b)
  | otherwise  
  = Node (getLabel ta) (buildLCES (getChild ta) (getChild tb)) Nothing 
  : buildLCES tas tbs  
  
flatten :: [Tree t] -> [Tree t]  
flatten  []    = [] 
flatten (t:[]) = flatRight t
flatten (t:ts) = flatLeft  t ++ ts   

flatRight :: Tree t -> [Tree t]
flatRight t@(Node _ []     _ ) = [t]
flatRight   (Node a (c:cn) pn) = [lf, Node a (c' ++ cn) pn] where
  (lf, c') = getFirstLeaf c

getFirstLeaf :: Tree t -> (Tree t, [Tree t])
getFirstLeaf ta@(Node _ [] _ ) = (ta, [])
getFirstLeaf    (Node a (c:cs) pn ) = (lf, [Node a (cn' ++ cs) pn]) where
  (lf, cn') = getFirstLeaf c
  
flatLeft :: Tree t -> [Tree t]  
flatLeft t@(Node _ []     _ ) = [t] 
flatLeft   (Node l  c     pn ) =  c ++ [Node l [] pn] 

-- returns the actual matching
-- TODO ensure potPret?
wbMatch :: [Tree a] -> [Tree a] -> (Tree a -> Tree a -> Bool) -> 
            Array (Int, Int) ([(Int, Int)], Int, Int)
wbMatch _ []  _ = listArray ((0,0),(0,0)) (repeat ([],0,0))
wbMatch [] _  _ = listArray ((0,0),(0,0)) (repeat ([],0,0))
wbMatch a' b' simf = 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 i j = if simf (a!i) (b!j) then ([(i,j)],1,1) else ([],0,0)
  -- this is the actual core recursive definintion of the algorithm
  concatMatch i j = maximumBy depthWeight l where
    l = if simf (a!i) (b!j) -- put the diagonal at the back to prefer symmetry
        then [merge (i,j) (m!(i-1,j)) a b, merge (i,j) (m!(i,j-1)) a b, merge (i,j) (m!(i-1,j-1)) a b]
        else [m!(i-1,j), m!(i,j-1), m!(i-1,j-1)]
  m = array ((0,0),(la,lb)) 
    (((0,0), match 0 0) :
    [((0,j), maximumBy depthWeight [m!(0,j-1), match 0 j]) | j <- [1..lb]] ++
    [((i,0), maximumBy depthWeight [m!(i-1,0), match i 0]) | i <- [1..la]] ++
    [((i,j), concatMatch i j) | i <- [1..la], j <- [1..lb]])      
    
-- returns the weight of a matching    
getWeight :: ([(Int,Int)], Int, Int) -> Int
getWeight (_,w,_) = w

-- compares two matching on the basis of their weight 
-- and in case of equal weight on the basis of the cummulative
-- depth of both compared trees
depthWeight :: (a, Int, Int) -> (a, Int, Int) -> Ordering
depthWeight (_,w,d) (_,w',d')
  |  w < w'      = LT
  |  w > w'      = GT
  | (d - d') < 0 = LT
  | (d - d') > 0 = GT
  | otherwise    = EQ 
  
-- merges two tuples contianing the matchings, weight and cumulative depth of both
-- matched trees.  
merge :: (Int,Int) -> ([(Int,Int)], Int, Int) -> Array Int (Tree b) -> Array Int (Tree b) ->
         ([(Int,Int)], Int, Int)
merge e@(i,j) p@(prv, w, d) a b  -- trace ((show e) ++ ":"++show prv++" d: " ++ show d ++ " update: " ++ show (d + (levelUp prv i fst a) + (levelUp prv j snd b)))
  | isFree prv i fst && isFree prv j snd = (e : prv, w + 1, d + depthInc) 
  | otherwise = p where
  depthInc = if levelUp prv i fst a && levelUp prv j snd b then 1 else 0

isFree :: [a] -> Int -> (a -> Int) -> Bool 
isFree prv i f = null prv || i > f (head prv) 
  
-- returns True if given a previous matching prv the the node i of tree a 
-- moves up a level and False otherwise
-- N.B. this should retrieve the preorder number, i.e. use potPret
-- If the preorder number of the previous match (in postorder!) is larger than
-- the current preorder number we move up in the three and increase the depth
levelUp :: [a] -> Int -> (a -> Int) -> Array Int (Tree b) -> Bool
levelUp prv i f a = null prv || getPn (a!f (head prv)) > getPn (a!i)
  
-- similarity measure for comparing tree labels
simLab :: (Eq a) => Tree a -> Tree a -> Bool    
simLab ta tb = getLabel ta == getLabel tb     

-- similarity measure for comparing anything
simEq :: (Eq a) => a -> a -> Bool
simEq a b = a == b    
  
fst3 :: (a, Int, Int) -> a    
fst3 (f, _,_) = f