-----------------------------------------------------------------------------
-- Copyright 2018, Ideas 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)
--
-- A type class for expressing choice, preference, and left-biased choice.
-- The 'Menu' datatype implements the type class by keeping all the
-- alternatives.
--
-----------------------------------------------------------------------------

module Ideas.Common.Strategy.Choice
   ( -- * Choice type class
     Choice(..)
     -- * Menu data type
   , Menu, (|->), doneMenu, eqMenuBy
     -- * Queries
   , elems, bests, bestsOrdered
   , isEmpty, hasDone, getByIndex, cut
     -- * Generalized functions
   , onMenu, onMenuWithIndex
   ) where

import Data.Maybe

infixr 3 .|., ./., |>, :|:, :/:, :|>
infixr 5 |->, :->

------------------------------------------------------------------------
-- Choice type class

-- | Laws: '.|.', './.' '|>' are all associative, and have 'empty' as their
-- unit element.
class Choice a where
   -- | Nothing to choose from.
   empty :: a
   -- | Normal (unbiased) choice.
   (.|.) :: a -> a -> a
   -- | Left-preference.
   (./.) :: a -> a -> a
   -- | Left-biased choice.
   (|>) :: a -> a -> a
   -- | One of the alternatives in a list (unbiased).
   choice     :: [a] -> a
   preference :: [a] -> a
   orelse     :: [a] -> a
   -- default implementation
   choice     xs = if null xs then empty else foldr1 (.|.) xs
   preference xs = if null xs then empty else foldr1 (./.) xs
   orelse     xs = if null xs then empty else foldr1 (|>)  xs

instance Choice [a] where
   empty    = []
   (.|.)    = (++)
   (./.)    = (++)
   xs |> ys = if null xs then ys else xs
   choice   = concat

instance Choice b => Choice (a -> b) where
   empty       = const empty
   (f .|. g) a = f a .|. g a
   (f ./. g) a = f a ./. g a
   (f |> g)  a = f a  |> g a

------------------------------------------------------------------------
-- Menu data type

-- Invariants for the Menu datatype:
-- (Unit) The left-hand side of :|: and :|> cannot be Empty
-- (Asso) :|: and :|> are balanced to the right

-- | A menu offers choices and preferences. It stores singleton bindings (thus
-- acting as a finite map) and one special element ('doneMenu'). It is an
-- instance of the 'Functor' and 'Monad' type classes.
data Menu k a = k :-> a
              | Done
              | Empty
              | Menu k a :|: Menu k a
              | Menu k a :/: Menu k a -- left-preference
              | Menu k a :|> Menu k a -- left-biased

instance (Eq k, Eq a) => Eq (Menu k a) where
   (==) = eqMenuBy (==) (==)

instance Choice (Menu k a) where
   empty  = Empty

   p0 .|. rest = rec p0 -- maintain invariant
    where
     rec Empty     = rest
     rec (p :|: q) = p :|: rec q
     rec p         = case rest of -- strict: also check rhs
                        Empty -> p
                        _     -> p :|: rest

   p0 ./. rest = rec p0 -- maintain invariant
    where
     rec Empty     = rest
     rec (p :/: q) = p :/: rec q
     rec p         = p :/: rest

   p0 |> rest = rec p0 -- maintain invariant
    where
     rec Empty     = rest
     rec (p :|> q) = p :|> rec q
     rec p         = p :|> rest

instance Functor (Menu k) where
   fmap f = rec
    where
      rec (p :|: q) = rec p :|: rec q
      rec (p :/: q) = rec p :/: rec q
      rec (p :|> q) = rec p :|> rec q
      rec (k :-> a) = k :-> f a
      rec Done      = Done
      rec Empty     = Empty

-- | Singleton binding
(|->) :: a -> s -> Menu a s
(|->) = (:->)

-- | Special element for denoting success
doneMenu :: Menu k a
doneMenu = Done

hasDone :: Menu k a -> Bool
hasDone (p :|: q) = hasDone p || hasDone q
hasDone (p :/: q) = hasDone p || hasDone q
hasDone (p :|> _) = hasDone p
hasDone (_ :-> _) = False
hasDone Done      = True
hasDone Empty     = False

-- | Equality with a comparison function for the elements
eqMenuBy :: (k -> k -> Bool) -> (a -> a -> Bool) -> Menu k a -> Menu k a -> Bool
eqMenuBy eqK eqA = test
 where
   test (p1 :|: p2) (q1 :|: q2) = test p1 q1 && test p2 q2
   test (p1 :/: p2) (q1 :/: q2) = test p1 q1 && test p2 q2
   test (p1 :|> p2) (q1 :|> q2) = test p1 q1 && test p2 q2
   test (k1 :-> a1) (k2 :-> a2) = eqK k1 k2 && eqA a1 a2
   test Done        Done        = True
   test Empty       Empty       = True
   test (p :/: Empty) q = test p q
   test (p :|> Empty) q = test p q
   test p (q :/: Empty) = test p q
   test p (q :|> Empty) = test p q
   test _ _ = False

------------------------------------------------------------------------
-- Queries

-- | Returns all elements that are in the menu.
elems :: Menu k a -> [(k, a)]
elems = ($ []) . rec
 where
   rec (p :|: q) = rec p . rec q
   rec (p :/: q) = rec p . rec q
   rec (p :|> q) = rec p . rec q
   rec (k :-> a) = ((k, a):)
   rec Done      = id
   rec Empty     = id

-- | Returns only the best elements that are in the menu with respect to
-- left-biased choices.
bests :: Menu k a -> [(k, a)]
bests = bestsWith (++)

-- | Returns only the best elements that are in the menu, with a given ordering.
bestsOrdered :: (k -> k -> Ordering) -> Menu k a -> [(k, a)]
bestsOrdered cmp = bestsWith merge
 where
   -- merge two lists with comparison function
   merge lx@(x:xs) ly@(y:ys) =
      case cmp (fst x) (fst y) of
         GT -> y : merge lx ys
         _  -> x : merge xs ly
   merge [] ys = ys
   merge xs [] = xs

-- helper: takes combinator for (:|:)
bestsWith:: ([(k, a)] -> [(k, a)] -> [(k, a)]) -> Menu k a -> [(k, a)]
bestsWith f = rec
 where
   rec (p :|: q) = f (rec p) (rec q)
   rec (p :/: q) = rec p ++ rec q
   rec (p :|> _) = rec p
   rec (k :-> a) = [(k, a)]
   rec Done      = []
   rec Empty     = []

-- | Is the menu empty?
isEmpty :: Menu k a -> Bool
isEmpty Empty = True
isEmpty _     = False -- because of invariant

-- | Get an element from the menu by its index.
getByIndex :: Int -> Menu k a -> Maybe (k, a)
getByIndex n = listToMaybe . drop n . elems

-- | Only keep the best elements in the menu.
cut :: Menu k a -> Menu k a
cut (p :|: q) = cut p .|. cut q
cut (p :/: q) = cut p ./. cut q
cut (p :|> _) = cut p
cut (k :-> a) = k |-> a
cut Done      = doneMenu
cut Empty     = empty

------------------------------------------------------------------------
-- Generalized functions

-- | Generalized monadic bind, with the arguments flipped.
{-# INLINE onMenu #-}
onMenu :: Choice b => (k -> a -> b) -> b -> Menu k a -> b
onMenu f e = rec
 where
   rec (p :|: q) = rec p .|. rec q
   rec (p :/: q) = rec p ./. rec q
   rec (p :|> q) = rec p  |> rec q
   rec (k :-> a) = f k a
   rec Done      = e
   rec Empty     = empty

-- | Maps a function over a menu that also takes the index of an element.
{-# INLINE onMenuWithIndex #-}
onMenuWithIndex :: Choice b => (Int -> k -> a -> b) -> b -> Menu k a -> b
onMenuWithIndex f e = snd . rec 0
 where
   rec n (p :|: q) = let (n1, pn) = rec n p
                         (n2, qn) = rec n1 q
                     in (n2, pn .|. qn)
   rec n (p :/: q) = let (n1, pn) = rec n p
                         (n2, qn) = rec n1 q
                     in (n2, pn ./. qn)
   rec n (p :|> q) = let (n1, pn) = rec n p
                         (n2, qn) = rec n1 q
                     in (n2, pn |> qn)
   rec n (k :-> a) = (n+1, f n k a)
   rec n Done      = (n, e)
   rec n Empty     = (n, empty)