module Heap
	(Heap, heapCreate, heapSize, heapLength, heapExchange,
	heapSort, heapMaximum, heapIncreaseKey, heapInsert,
	heapFind, heapRemove, heapRemoveAt, heapMerge,
	heapLastElement, heapIntersect, heapDifference, heapEquals,
	heapShow, heapMap, heapFilter, heapContains,
	heapLesserOrEq, heapFold, heapTrimToLength, heapDoubleLength) where

import Array

newtype Heap a = HeapC (Int, Int, (Array Int a))

instance (Eq a, Ord a) => Eq (Heap a) where
	(==) = heapEquals
instance Ord a => Ord (Heap a) where
	(<=) = heapLesserOrEq

-- A Heap starts with room for 10 elements, but it will dynamically
-- expand to fit any number of elements, courtesy heapInsert calling heapDoubleLength
heapCreate :: (Heap a)
heapCreate = HeapC (0, 10, (array (1,10) []))

heapSize :: (Heap a) -> Int
heapSize (HeapC (b,q,a)) = b

heapLength :: (Heap a) -> Int
heapLength (HeapC (a,q,b)) = q

heapArray :: (Heap a) -> (Array Int a)
heapArray (HeapC (a,q,b)) = b

heapExchange :: (Heap a) -> Int -> Int -> (Heap a)
heapExchange (HeapC (d,q,a)) b c = HeapC (d, q, a // [(b,a!c),(c,a!b)])

heapSort :: Ord a => (Heap a) -> Int -> (Heap a)
heapSort (HeapC (b,q,a)) i
	= let l = i*2 ; r = i*2+1 in
		let large = heapSortHelper l b (a!l) (a!i) i in
			let largest = heapSortHelper r b (a!r) (a!large) large in
				if largest /= i
					then heapSort (heapExchange (HeapC (b,q,a)) i largest) largest
					else (HeapC (b,q,a))

heapSortHelper :: Ord a => Int -> Int -> a -> a -> Int -> Int
heapSortHelper a b c d e
	| a <= b && c > d	= a
	| otherwise			= e

{-
heapFromArray :: Ord a => Int -> (Array Int a) -> (Heap a)
heapFromArray a b
	| a==0		= heapCreate
	| otherwise	= 
-}

heapMaximum :: (Heap a) -> a
heapMaximum (HeapC (b,q,a)) = a!1

heapLastElement :: (Heap a) -> a
heapLastElement (HeapC (b,q,a)) = a!b

heapIncreaseKey :: Ord a => (Heap a) -> Int -> a -> (Heap a)
heapIncreaseKey (HeapC (b,q,a)) i k
	| k < a!i		= error "new key is smaller than current key"
	| otherwise		= heapIncreaseKeyHelper (HeapC (b,q,(a // [(i,k)]))) i

heapIncreaseKeyHelper :: Ord a => (Heap a) -> Int -> (Heap a)
heapIncreaseKeyHelper (HeapC (b,q,a)) i
	| i > 1 && a!(i `div` 2) < a!i
		= heapIncreaseKeyHelper (heapExchange (HeapC (b,q,a)) i (i `div` 2)) (i `div` 2)
	| otherwise						= (HeapC (b,q,a))

-- This will set the array size to anything
heapSetLength :: (Heap a) -> Int -> (Heap a)
heapSetLength (HeapC (a,q,b)) z
	| z<1		= error "heap length must be >= 1"
	| z>=a		= (HeapC (a, z, array (1,z) (assocs b)))
	| otherwise = (HeapC (z, z, array (1,z) (assocs b)))

-- This will reduce the array size to the heap size
-- Unused, but could be useful for saving a Heap to I/O
heapTrimToLength :: (Heap a) -> (Heap a)
heapTrimToLength (HeapC (a,q,b)) = (HeapC (a, a+1, array (1,a+1) (assocs b)))

-- If we run out of space, double the size of the array.
-- This ensures that we don't have to increase the size that often, since
-- creating a new Array is probably not the cheapest operation.
heapDoubleLength :: (Heap a) -> (Heap a)
heapDoubleLength (HeapC (a,q,b)) = (HeapC (a, q*2, array (1,q*2) (assocs b)))

-- It's not an error to try to insert an existing element
heapInsert :: Ord a => (Heap a) -> a -> (Heap a)
heapInsert (HeapC (b,q,a)) k
	| b+1 > q					= heapInsert (heapDoubleLength (HeapC (b,q,a))) k
	| (heapFind (HeapC (b,q,a)) k) == 0	= heapIncreaseKey (HeapC (b+1, q, (a // [(b+1, k)]))) (b+1) k
	| otherwise					= (HeapC (b,q,a))

heapFind :: Ord a => (Heap a) -> a -> Int
heapFind (HeapC (b,q,a)) k = heapFindHelper (HeapC (b,q,a)) k 1

-- This performs lots of checks to make sure it doesn't ask for out of bounds elements
-- It is safe to search both child trees and + their result since an element can only occur once
-- in a heap, courtesy heapInsert using heapFind.
heapFindHelper :: Ord a => (Heap a) -> a -> Int -> Int
heapFindHelper (HeapC (b,q,a)) k i
	| (i <= b) && (k == a!i)				= i
	| (i*2 <= b) && (i*2+1 <= b) && (k <= a!(i*2)) && (k <= a!(i*2+1))
		= (heapFindHelper (HeapC (b,q,a)) k (i*2)) + (heapFindHelper (HeapC (b,q,a)) k (i*2+1))
	| (i*2 <= b) && (k <= a!(i*2))		= heapFindHelper (HeapC (b,q,a)) k (i*2)
	| (i*2+1 <= b) && (k <= a!(i*2+1))	= heapFindHelper (HeapC (b,q,a)) k (i*2+1)
	| otherwise							= 0 -- dirty hack, but it works

heapRemove :: Ord a => (Heap a) -> a -> (Heap a)
heapRemove (HeapC (b,q,a)) k
	| heapFind (HeapC (b,q,a)) k == 0	= (HeapC (b,q,a))
	| otherwise		= heapRemoveAt (HeapC (b,q,a)) (heapFind (HeapC (b,q,a)) k)

heapRemoveAt :: Ord a => (Heap a) -> Int -> (Heap a)
heapRemoveAt (HeapC (b,q,a)) i
	| i==0		= error "cannot remove the 0th element"
	| i>b		= error "index to remove is out of bounds"
	| i==b		= (HeapC (b-1,q,a))
	| otherwise = heapSort (HeapC (b-1,q,(a // [(i,a!b)]))) i

heapMerge :: Ord a => (Heap a) -> (Heap a) -> (Heap a)
heapMerge a (HeapC (c,q,b))
	| c==0		= a
	| otherwise	= let x = heapLastElement (HeapC (c,q,b)) in
					let y = heapRemoveAt (HeapC (c,q,b)) c in
						heapMerge (heapInsert a x) y

heapIntersect :: Ord a => (Heap a) -> (Heap a) -> (Heap a)
heapIntersect (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0 || c==0	= heapCreate
	| otherwise		= heapIntersectHelper heapCreate (HeapC (a,q,b)) (HeapC (c,w,d))

heapIntersectHelper :: Ord a => (Heap a) -> (Heap a) -> (Heap a) -> (Heap a)
heapIntersectHelper a (HeapC (b,q,c)) (HeapC (d,w,e))
	| b==0 || d==0	= a
	| otherwise		= let x = heapLastElement (HeapC (b,q,c)) in
						let y = heapRemoveAt (HeapC (b,q,c)) b in
							let z = heapFind (HeapC (d,w,e)) x in
								if z == 0
									then heapIntersectHelper a y (HeapC (d,w,e))
									else heapIntersectHelper (heapInsert a x) y (heapRemoveAt (HeapC (d,w,e)) z)

heapDifference :: Ord a => (Heap a) -> (Heap a) -> (Heap a)
heapDifference (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0 || c==0	= (HeapC (a,q,b))
	| otherwise		= let x = heapLastElement (HeapC (c,w,d)) in
						let y = heapRemoveAt (HeapC (c,w,d)) c in
							let z = heapFind (HeapC (a,q,b)) x in
								if z==0
									then heapDifference (HeapC (a,q,b)) y
									else heapDifference (heapRemoveAt (HeapC (a,q,b)) z) y

heapEquals :: Ord a => (Heap a) -> (Heap a) -> Bool
heapEquals (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0 && c==0	= True
	| a/=c			= False
	| otherwise		= let x = heapLastElement (HeapC (c,w,d)) in
						let y = heapRemoveAt (HeapC (c,w,d)) c in
							let z = heapFind (HeapC (a,q,b)) x in
								if z==0
									then False
									else heapEquals (heapRemoveAt (HeapC (a,q,b)) z) y

-- This function is lexiographically correct
heapLesserOrEq :: Ord a => (Heap a) -> (Heap a) -> Bool
heapLesserOrEq (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0 && c==0	= True
	| a==0 && c/=0	= True -- empty heaps are always lesser than non-empty heaps
	| a/=0 && c==0	= False
	| b!1 > d!1		= False
	| b!1 < d!1		= True
	| otherwise		= heapLesserOrEq (heapRemoveAt (HeapC (a,q,b)) 1) (heapRemoveAt (HeapC (c,w,d)) 1)

heapContains :: Ord a => (Heap a) -> (Heap a) -> Bool
heapContains (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0 && c==0	= True
	| a/=0 && c==0	= True -- This will be the usual end-result
	| a==0 && c/=0	= False
	| otherwise		= let x = heapLastElement (HeapC (c,w,d)) in
						let y = heapRemoveAt (HeapC (c,w,d)) c in
							let z = heapFind (HeapC (a,q,b)) x in
								if z==0
									then False
									else heapContains (HeapC (a,q,b)) y

heapMap :: (Ord a, Ord b) => (a -> b) -> (Heap a) -> (Heap b)
heapMap f (HeapC (a,q,b)) = heapMapHelper f (HeapC (a,q,b)) heapCreate

heapMapHelper :: (Ord a, Ord b) => (a -> b) -> (Heap a) -> (Heap b) -> (Heap b)
heapMapHelper f (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0		= (HeapC (c,w,d))
	| otherwise	= let x = heapLastElement (HeapC (a,q,b)) in
					let y = heapRemoveAt (HeapC (a,q,b)) a in
						heapMapHelper f y (heapInsert (HeapC (c,w,d)) (f x))

heapFilter :: Ord a => (a -> Bool) -> (Heap a) -> (Heap a)
heapFilter f (HeapC (a,q,b)) = heapFilterHelper f (HeapC (a,q,b)) heapCreate

heapFilterHelper :: Ord a => (a -> Bool) -> (Heap a) -> (Heap a) -> (Heap a)
heapFilterHelper f (HeapC (a,q,b)) (HeapC (c,w,d))
	| a==0		= (HeapC (c,w,d))
	| otherwise	= let x = heapLastElement (HeapC (a,q,b)) in
					let y = heapRemoveAt (HeapC (a,q,b)) a in
						if (f x)
							then heapFilterHelper f y (heapInsert (HeapC (c,w,d)) x)
							else heapFilterHelper f y (HeapC (c,w,d))

heapFold :: Ord b => (a -> b -> a) -> a -> (Heap b) -> a
heapFold f a (HeapC (c,q,d))
	| c==0		= a
	| otherwise	= let x = heapLastElement (HeapC (c,q,d)) in
					let y = heapRemoveAt (HeapC (c,q,d)) c in
						heapFold f (f a x) y

heapShow :: Ord a => (a -> String) -> (Heap a) -> String
heapShow f (HeapC (a,q,b)) = "{" ++ (heapShowHelper f (HeapC (a,q,b)) "") ++ "}"

heapShowHelper :: Ord a => (a -> String) -> (Heap a) -> String -> String
heapShowHelper f (HeapC (a,q,b)) s
	| a==0		= s
	| otherwise	= let x = heapMaximum (HeapC (a,q,b)) in
					let y = heapRemoveAt (HeapC (a,q,b)) 1 in
						heapShowHelper f y (s ++ (f x) ++ ",")