-- | Second assignment for IB016 Seminar on Functional Programming, spring 2016
--
-- Your task is to implement a priority queue based on heap data structure. A
-- priority queue is a data structure which holds values sorted by priority
-- (so that the value with the highest priority is at the top of the heap). The
-- interface is given by this file, together with suggested time complexities
-- for each of the functions. Your implementation should meet these time
-- complexities. 
--
-- It is suggested to use binary heap. You can also use more advanced versions
-- of heaps, for example binomial heap or Fibonacci heap. It is possible to
-- meet all the requirements with binary heap which is reasonably simple to
-- implement. You can get extra points for more advanced heap implementations.
--
--   * <https://en.wikipedia.org/wiki/Binary_heap Binary heap>
--   * <https://en.wikipedia.org/wiki/Binomial_heap Binomial heap>
--   * <https://en.wikipedia.org/wiki/Fibonacci_heap Fibonacci heap>
--
-- You may define your own (internal) functions that the module does not export
-- (e.g. helper functions for tree rotations and such).
-- You are not required (but still strongly encouraged) to document such functions
-- using the <https://www.haskell.org/haddock/doc/html/ch03s08.html Haddock markup>.
-- You can use any modules from the @Base@ package if you wish.
-- However, you may not use any other packages.
--
-- Name: Name Surname
-- UID: 123456

module PriorityQueue
    ( -- * Priority Queue Type
      PriorityQueue
      -- * Construction
    , empty
    , singleton
    , fromList
    , insert
    , union
      -- * Querying
    , isEmpty
    , size
    , getTop
    , toDescList
      -- * Modification
    , extractTop
    , extractGetTop
    , replaceTop
    , modifyTop
      -- * Debugging
    , valid
    ) where

import Data.List ( unfoldr )
import Control.Applicative ( (<$>) )
import Data.Function ( on )

-- | Defines a 'PriorityQueue' type. It should implement a (binary) heap. Note
-- that the data constructors of this type are not exported from the module to
-- disallow direct manipulations. (If they should have been exported the module
-- header entry would read @PriorityQueue (..)@).
--
-- Note: You should define 'Show', 'Eq', and 'Functor' instances. 'Show' and
-- 'Eq' should be defined by hand, not derived. Nevertheless, it is useful to
-- derive 'Show' instance during development and change to a custom-defined
-- instance once you trust your implementation.
data PriorityQueue p v
      = Node { nsize  :: Int
             , nprio  :: p
             , nval   :: v
             , nleft  :: PriorityQueue p v
             , nright :: PriorityQueue p v
             }
      | Empty

-- | /O(n log n)/. Define instance of 'Show' similar to the one for
-- 'Data.Map.Map'.
--
-- >>>  show empty
-- "fromList []"
--
-- >>> show $ singleton 42 ()
-- "fromList [(42,())]
--
instance (Show p, Show v, Ord p) => Show (PriorityQueue p v) where
  show = ("fromList " ++) . show . toDescList

-- These functions are not exported from the module, however, they can be
-- useful for directing descend into the appropriate branch of a binary heap in
-- insert.
maxPathForSize :: Int -> Int
maxPathForSize 0 = 0
maxPathForSize n = 1 + maxPathForSize (n `div` 2)

minPathForSize :: Int -> Int
minPathForSize n
    | 2 ^ mh - 1 == n = mh
    | otherwise       = mh - 1
  where
    mh = maxPathForSize n

-- | /O(1)/. An empty 'PriorityQueue'.
--
-- >>> empty
-- fromList []
empty :: PriorityQueue p v
empty = Empty

-- | /O(1)/. A 'PriorityQueue' with a single element.
--
-- >>> singleton 42 "hi"
-- fromList [(42,"hi")]
singleton :: Ord p => p -> v -> PriorityQueue p v
singleton p v = Node 1 p v Empty Empty

-- | /O(1)/. Is the 'PriorityQueue' empty?
--
-- >>> isEmpty empty
-- True
-- >>> isEmpty $ singleton 42 "hi"
-- False
isEmpty :: PriorityQueue p v -> Bool
isEmpty Empty = True
isEmpty _     = False

-- | /O(1)/. The number of elements in the 'PriorityQueue'.
--
-- >>> size $ singleton 42 "hi"
-- 1
size :: PriorityQueue p v -> Int
size Node { nsize = sz } = sz
size Empty = 0

-- | /O(log n)/. @insert prio val pq@ inserts the element @val@ with the
-- priority @prio@ into the 'PriorityQueue' @pq@.
--
-- >>> insert 42 "hi" empty
-- fromList [(42,"hi")]
--
-- >>> insert 0 "" $ singleton 42 "hi"
-- fromList [(42,"hi"),(0,"")]
insert :: Ord p => p -> v -> PriorityQueue p v -> PriorityQueue p v
insert p v node@Node { nleft = l, nright = r, nprio = rp, nsize = sz }
    | lsz == size r || minPathForSize lsz < maxPathForSize lsz =
        case insert p v l of
            nl@Node { nprio = dp } -> if rp >= dp then node { nleft = nl, nsize = sz + 1 }
                                                  else swap (\r d -> r { nleft = d }) node nl
    | otherwise =
        case insert p v r of
            nr@Node { nprio = dp } -> if rp >= dp then node { nright = nr, nsize = sz + 1 }
                                                  else swap (\r d -> r { nright = d }) node nr
  where
    lsz = size l
    swap bind root descendant = bind nroot ndesc
      where
        ndesc = descendant { nprio = nprio root, nval = nval root }
        nroot = root { nprio = nprio descendant, nval = nval descendant, nsize = sz + 1 }
insert p v Empty = singleton p v

-- | /O(1)/. Get the element with highest priority from the given 'PriorityQueue'.
-- Returns 'Nothing' if and only if the 'PriorityQueue' is 'empty'.
--
-- >>> getTop $ insert 0 "" $ singleton 42 "hi"
-- Just (42,"hi")
--
-- >>> getTop empty
-- Nothing
getTop :: Ord p => PriorityQueue p v -> Maybe (p, v)
getTop Empty = Nothing
getTop Node { nprio = p, nval = v } = Just (p, v)

-- | /O(log n)/. Extract (remove) the element with highest priority from the
-- given 'PriorityQueue'. Returns 'empty' if the 'PriorityQueue' is 'empty'.
--
-- >>> extractTop $ insert 0 "" $ singleton 42 "hi"
-- fromList [(0,"")]
--
-- >>> extractTop empty
-- fromList []
extractTop :: Ord p => PriorityQueue p v -> PriorityQueue p v
extractTop Empty = Empty
extractTop node = case root' of
                    Empty -> Empty
                    root  -> bubbleDown (root { nprio = lastp, nval = lastv })
  where
    (lastp, lastv, root') = extractLast node

-- | Adjust heap which can have broken root.
bubbleDown Empty = Empty
bubbleDown node@Node { nleft = Empty, nright = Empty } = node
bubbleDown node@Node { nleft = l, nright = Empty }
    | nprio l <= nprio node = node
    | otherwise = node { nprio = nprio l, nval = nval l
                       , nleft = bubbleDown (l { nprio = nprio node, nval = nval node }) }
bubbleDown node@Node { nleft = l, nright = r }
    | maxdesc <= nprio node = node
    | nprio l >= nprio r = node { nprio = nprio l, nval = nval l
                              , nleft = bubbleDown (l { nprio = nprio node, nval = nval node }) }
    | otherwise = node { nprio = nprio r, nval = nval r
                       , nright = bubbleDown (r { nprio = nprio node, nval = nval node }) }
  where
    maxdesc = max (nprio l) (nprio r)

-- | /O(log n)/. Get and extract the element with highest priority from the
-- given 'PriorityQueue'. Returns 'Nothing' if and only if the 'PriorityQueue'
-- is 'empty'.
--
-- >>> extractGetTop $ insert 0 "" $ singleton 42 "hi"
-- Just ((42,"hi"),fromList [(0,"")])
--
-- >>> extractGetTop empty
-- Nothing
extractGetTop :: Ord p => PriorityQueue p v -> Maybe ((p, v), PriorityQueue p v)
extractGetTop pq = flip (,) (extractTop pq) <$> getTop pq

extractLast :: PriorityQueue p v -> (p, v, PriorityQueue p v)
extractLast Empty = error "extractLast: Empty"
extractLast node@Node { nleft = Empty, nright = Empty } = (nprio node, nval node, Empty)
extractLast node@Node { nsize = sz, nleft = l, nright = r }
    | maxPathForSize lsz > maxPathForSize rsz
        = case extractLast l of
            (p, v, nl) -> (p, v, node { nleft = nl, nsize = sz - 1 } )
    | otherwise = case extractLast r of
            (p, v, nr) -> (p, v, node { nright = nr, nsize = sz - 1 } )
  where
    lsz = size l
    rsz = size r

-- | /O(log n)/. Replace the highest priority element with new priority and
-- value.
--
-- >>> replaceTop 2 "a" $ fromList [(1, ""), (3, "")]
-- fromList [(2,"a"),(1,"")]
--
-- >>> replaceTop 0 "a" $ fromList [(1, ""), (3, "")]
-- fromList [(1,""),(0,"a")]
replaceTop :: Ord p => p -> v -> PriorityQueue p v -> PriorityQueue p v
replaceTop p v = modifyTop (\_ _ -> Just (p, v))

-- | /O(log n)/. Replace the highest priority element using a function callback
-- which is called with the current priority and value of the top element. If
-- the callback returns 'Nothing' the root is removed, otherwise new values of
-- priority and value are determined from the returned Just value.
--
-- >>> modifyTop (\k v -> Just (k + 2, v)) $ fromList [(1, ""), (2, "")]
-- fromList [(4,""),(1,"")]
--
-- >>> modifyTop (\k v -> Just (k - 2, v)) $ fromList [(1, ""), (2, "")]
-- fromList [(1,""),(0,"")]
modifyTop :: Ord p => (p -> v -> Maybe (p, v)) -> PriorityQueue p v -> PriorityQueue p v
modifyTop _ Empty = Empty
modifyTop f node@Node { nprio = p, nval = v } = case f p v of
    Just (p', v') -> bubbleDown $ node { nprio = p', nval = v' }
    Nothing       -> extractTop node

-- | /O(n log n)/. Build a 'PriorityQueue' from a list of priority – value pairs.
--
-- Note: this can be done in /O(n)/, you can be awarded bonus points if you
-- implement this version.
-- 
-- >>> fromList [(0, ()), (4, ()), (17, ()), (3, ())]
-- fromList [(17,()),(4,()),(3,()),(0,())]
fromList :: Ord p => [(p, v)] -> PriorityQueue p v
fromList = foldl (\pq (p, v) -> insert p v pq) empty

-- | /O(n log n)/. Convert the given 'PriorityQueue' to a list of its priority
-- – value pairs ordered by priority in descending order.
--
-- >>> toDescList $ fromList [(0, ()), (4, ()), (17, ()), (3, ())]
-- [(17,()),(4,()),(3,()),(0,())]
toDescList :: Ord p => PriorityQueue p v -> [(p, v)]
toDescList = unfoldr extractGetTop

-- | /O(n log n)/. Compute union of two 'PriorityQueue's.
--
-- Note: this can be done in /O(n)/, you can be awarded bonus points if you
-- implement this version.
--
-- >>> fromList [(0, ()), (4, ()), (17, ()), (3, ())] `union` singleton 42 ()
-- fromList [(42,()),(17,()),(4,()),(3,()),(0,())]
union :: Ord p => PriorityQueue p v -> PriorityQueue p v -> PriorityQueue p v
union a b = fromList (toDescList a ++ toDescList b)

-- | Define @'PriorityQueue' p@ to be an instance of 'Functor', keep in mind the
-- standard behaviour of 'Functor'.
instance Functor (PriorityQueue p) where
    fmap _ Empty = Empty
    fmap f node@Node { nval = v, nleft = l, nright = r }
        = node { nval = f v, nleft = fmap f l, nright = fmap f r }

-- | Check validity of a given 'PriorityQueue', this should validate that 'size'
-- returns the right value for the heap and all its subheaps and that all the
-- requirements given by the data structure are met (e.g. for binary heap,
-- it should check that the tree is left aligned, lengths of paths from root to
-- leaf differ at most by one, and that the priority is nonascending on every
-- path from root to leaf). You should use this function for testing.
valid :: Show p => Ord p => PriorityQueue p v -> Bool
valid Empty = True
valid pq = fst (checkSize pq) && all sorted paths && leftAligned paths
  where
    checkSize Empty = (True, 0)
    checkSize n@Node { nleft = l, nright = r } = (vl && vr && sz == size n, sz)
      where
        (vl, sl) = checkSize l
        (vr, sr) = checkSize r
        sz = 1 + sl + sr

    paths = pa pq
    pa Empty = [[]]
    pa Node { nprio = p, nleft = l, nright = r } = map (p:) $ pa l ++ pa r

    sorted xs = and $ zipWith (>=) xs (tail xs)

    leftAligned [] = True
    leftAligned (x:xs) = all (\p -> length p == lx - 1) shorter
      where
        lx = length x
        shorter = dropWhile (\p -> length p == lx) xs