module Heap (Heap(Empty, Node), getValue, insert, pop, fromList) where

import Prelude hiding (foldl, foldr)
import Data.Foldable
import Data.Monoid

data Heap a = Empty | Node (Heap a) a (Heap a)

instance (Show a) => Show (Heap a) where
  show Empty = "E"
  show (Node lh a rh) = "( " ++ show a ++ " ( L = " ++ show lh ++ " ) " ++ " ( R = " ++ show rh ++ " ) )"

instance Foldable Heap where
  foldMap _ Empty = mempty
  foldMap f (Node l a r) = f a `mappend` foldMap f l `mappend` foldMap f r
  foldr _ z Empty = z
  foldr f z (Node Empty a Empty) = f a z
  foldr f z (Node l a r) = f a (foldr f (foldr f z r) l)

getValue :: Heap a -> a
getValue Empty = undefined
getValue (Node _ a _) = a

insert :: Ord a => a -> Heap a -> Heap a
insert x h = arrange $ _insert x h

_insert :: Ord a => a -> Heap a -> Heap a
_insert x Empty = Node Empty x Empty
_insert x (Node Empty a Empty) = Node (Node Empty x Empty) a Empty
_insert x (Node Empty a r) = Node (Node Empty x Empty) a r
_insert x (Node l a Empty) = Node l a (Node Empty x Empty)
_insert x (Node l a r) = Node l a (_insert x r)

arrange :: Ord a => Heap a -> Heap a
arrange Empty = Empty
arrange (Node Empty a Empty) = Node Empty a Empty
arrange (Node Empty a r) = let (Node rl rv rr) = arrange r
                               in if a <= rv
                                    then Node Empty a (Node rl rv rr)
                                    else Node Empty rv (arrange $ Node rl a rr)
arrange (Node l a Empty) = let (Node ll lv lr) = arrange l
                               in if a <= lv
                                    then Node (Node ll lv lr) a Empty
                                    else Node (arrange $ Node ll a lr) lv Empty
arrange (Node l a r) = let ar = arrange r
                           al = arrange l
                           in replace al a ar
  where
    replace (Node ll lv lr) a' (Node rl rv rr)
      | a <= lv && a <= rv = Node (arrange $ Node ll lv lr) a' (arrange $ Node rl rv rr)
      | a > lv && a <= rv = Node (arrange $ Node ll a' lr) lv (Node rl rv rr)
      | a <= lv && a > rv = Node (Node ll lv lr) rv (arrange $ Node rl a' rr)
      | a > lv && a > rv && lv >= rv = Node (Node ll lv lr) rv (arrange $ Node rl a' rr)
      | otherwise = Node (arrange $ Node ll a' lr) lv (Node rl rv rr)
    replace _ _ _ = undefined

fromList :: Ord a => [a] -> Heap a
fromList = foldl (flip insert) Empty

maxValue :: Ord a => Heap a -> a
maxValue (Node Empty a Empty) = a
maxValue (Node l _ Empty) = maxValue l
maxValue (Node Empty _ r) = maxValue r
maxValue (Node l _ r) = max (maxValue l) (maxValue r)
maxValue _ = undefined

removeMaxValue :: Ord a => Heap a -> Heap a
removeMaxValue Empty = undefined
removeMaxValue (Node Empty _ Empty) = Empty
removeMaxValue (Node l@(Node _ _ _) a Empty) = Node (removeMaxValue l) a Empty
removeMaxValue (Node Empty a r@(Node _ _ _)) = Node Empty a (removeMaxValue r)
removeMaxValue (Node l@(Node _ _ _) a r@(Node _ _ _))
  | maxValue l <= maxValue r = Node l a (removeMaxValue r)
  | otherwise = Node (removeMaxValue l) a r

pop :: Ord a => Heap a -> (Maybe a, Heap a)
pop Empty = (Nothing, Empty)
pop (Node Empty a Empty) = (Just a, Empty)
pop (Node l a Empty) = (Just a, l)
pop (Node Empty a r) = (Just a, r)
pop (Node l a r) = (Just a, arrange heap)
  where
    heap
      | maxValue l <= maxValue r = let rv = maxValue r
                                       in Node l rv (removeMaxValue r)
      | otherwise = let lv = maxValue l
                                       in Node (removeMaxValue l) lv r