module Data.NonEmpty.Map (
   T,
   insert,
   singleton,
   member,
   lookup,
   minViewWithKey,
   maxViewWithKey,
   fromList,
   toAscList,
   flatten,
   union,
   unionLeft,
   unionRight,
   ) where

import qualified Data.NonEmpty as NonEmpty

import qualified Data.Map as Map
import Data.Map (Map, )

import Control.Monad (mzero, )
import Data.Maybe (fromMaybe, )
import Data.Tuple.HT (forcePair, )
import Data.Ord.HT (comparing, )

import Prelude hiding (lookup, )


{-
The first field will always contain the smallest element.
-}
data T k a = Cons (k, a) (Map k a)
   deriving (Eq, Ord)

instance (Show k, Show a) => Show (T k a) where
   showsPrec p xs =
      showParen (p>10) $
         showString "NonEmptyMap.fromList " .
         showsPrec 11 (toAscList xs)


insert :: Ord k => k -> a -> Map k a -> T k a
insert = curry $ insertGen fst

insertGen :: Ord k => (((k,a),(k,a)) -> (k,a)) -> (k,a) -> Map k a -> T k a
insertGen select y xt =
   uncurry Cons $
   fromMaybe (y, xt) $ do
      (x,xs) <- Map.minViewWithKey xt
      case comparing fst y x of
         GT -> return (x, uncurry Map.insert y xs)
         EQ -> return (select (y,x), xs)
         LT -> mzero

singleton :: k -> a -> T k a
singleton k a = Cons (k,a) Map.empty

member :: (Ord k) => k -> T k a -> Bool
member y (Cons x xs) =
   y == fst x || Map.member y xs

lookup :: (Ord k) => k -> T k a -> Maybe a
lookup y (Cons x xs) =
   if y == fst x
     then Just $ snd x
     else Map.lookup y xs

minViewWithKey :: T k a -> ((k,a), Map k a)
minViewWithKey (Cons x xs) = (x,xs)

maxViewWithKey :: (Ord k) => T k a -> ((k,a), Map k a)
maxViewWithKey (Cons x xs) =
   forcePair $
   case Map.maxViewWithKey xs of
      Nothing -> (x,xs)
      Just (y,ys) -> (y, uncurry Map.insert x ys)

fromList :: (Ord k) => NonEmpty.T [] (k,a) -> T k a
fromList (NonEmpty.Cons x xs) = uncurry insert x $ Map.fromList xs

toAscList :: T k a -> NonEmpty.T [] (k,a)
toAscList (Cons x xs) = NonEmpty.cons x $ Map.toAscList xs

flatten :: (Ord k) => T k a -> Map k a
flatten (Cons x xs) = uncurry Map.insert x xs

union :: (Ord k) => T k a -> T k a -> T k a
union (Cons x xs) (Cons y ys) =
   uncurry Cons $
   case Map.union xs ys of
      zs ->
         case comparing fst x y of
            LT -> (x, Map.union zs $ uncurry Map.singleton y)
            GT -> (y, uncurry Map.insert x zs)
            EQ -> (x, zs)

unionLeft :: (Ord k) => Map k a -> T k a -> T k a
unionLeft xs (Cons y ys) =
   insertGen snd y $ Map.union xs ys

unionRight :: (Ord k) => T k a -> Map k a -> T k a
unionRight (Cons x xs) ys =
   insertGen fst x $ Map.union xs ys