-----------------------------------------------------------------------------
-- Copyright 2019, Advise-Me project team. This file is distributed under 
-- the terms of the Apache License 2.0. For more information, see the files
-- "LICENSE.txt" and "NOTICE.txt", which are included in the distribution.
-----------------------------------------------------------------------------
-- |
-- Maintainer  :  bastiaan.heeren@ou.nl
-- Stability   :  provisional
-- Portability :  portable (depends on ghc)
--
-----------------------------------------------------------------------------

module Bayes.EliminationOrdering 
   ( EliminationOrdering, nextVariable
   , simpleOrder, staticOrder, minFactorOrder, minFactorVars
   ) where

import Bayes.Factor
import Data.List

-- depends on 'factors' and 'keeps'
newtype EliminationOrdering = Elim { fromElim :: [Dimensions] -> [String] -> Maybe (String, EliminationOrdering) }

nextVariable :: EliminationOrdering -> [Factor] -> [String] -> Maybe (String, EliminationOrdering)
nextVariable (Elim f) fs = f (map dimensions fs)

simpleOrder :: EliminationOrdering
simpleOrder = Elim $ \ds keeps -> 
   case filter (`notElem` keeps) (vars ds) of
      []  -> Nothing
      y:_ -> Just (y, simpleOrder)

staticOrder :: [String] -> EliminationOrdering
staticOrder xs = Elim $ \ds keeps -> 
   let rec [] = fromElim simpleOrder ds keeps
       rec (y:ys) 
          | y `elem` keeps = rec ys
          | otherwise      = Just (y, staticOrder ys)
   in rec xs

minFactorOrder :: EliminationOrdering
minFactorOrder = Elim $ \ds keeps -> 
   case minFactorNext (`notElem` keeps) ds of
      Just v  -> Just (v, minFactorOrder)
      Nothing -> Nothing

minFactorNext :: (String -> Bool) -> [Dimensions] -> Maybe String
minFactorNext p ds =
   case sortOn fst (map f vs) of
      [] -> Nothing
      (_, v):_ -> Just v
 where
   vs  = filter p (vars ds)
   f v = let xs  = filter (hasVarD v) ds
             res = mergesD xs
         in (size res - size xs, v)

-- |elimination order of variables corresponding to the min-factor heuristic
minFactorVars :: [Factor] -> [String]
minFactorVars fs = 
   case minFactorNext (const True) (map dimensions fs) of 
      Just v  -> v : minFactorVars (eliminate fs v)
      Nothing -> []