{-# LANGUAGE NamedFieldPuns #-}
module Data.OrderedMap (
      OrderedMap
    , build
    , elems
    , fromList
    , get
    , keys
    , lookup
    , mapi
    , set
    , toList
  ) where

import Data.Map (Map, (!), mapWithKey)
import qualified Data.Map as Map (fromList, insert, lookup, member)
import Prelude hiding (lookup)

data OrderedMap k a = OrderedMap {
      assoc :: Map k a
    , keys :: [k]
  } 

instance (Ord k, Show k, Show a) => Show (OrderedMap k a) where
  show = show . toList

instance Functor (OrderedMap k) where
  fmap f orderedMap = orderedMap {assoc = fmap f (assoc orderedMap)}

instance Ord k => Foldable (OrderedMap k) where
  foldMap f (OrderedMap {assoc, keys}) = foldMap f $ (assoc !) <$> keys

instance Ord k => Traversable (OrderedMap k) where
  sequenceA (OrderedMap {assoc, keys}) =
    (flip OrderedMap keys) <$> sequenceA assoc 

elems :: Ord k => OrderedMap k a -> [a]
elems (OrderedMap {assoc, keys}) = (assoc !) <$> keys

toList :: Ord k => OrderedMap k a -> [(k, a)]
toList (OrderedMap {assoc, keys}) = (\k -> (k, assoc ! k)) <$> keys

fromList :: Ord k => [(k, a)] -> OrderedMap k a
fromList keyValueList = OrderedMap {
      assoc = Map.fromList keyValueList
    , keys = fst <$> keyValueList
  }

build :: Ord k => (k -> a) -> [k] -> OrderedMap k a
build f keys = OrderedMap {
      assoc = Map.fromList $ (\k -> (k, f k)) <$> keys
    , keys
  }

get :: Ord k => k -> OrderedMap k a -> a
get k = (! k) . assoc

lookup :: Ord k => k -> OrderedMap k a -> Maybe a
lookup k = (Map.lookup k) . assoc

set :: Ord k => k -> a -> OrderedMap k a -> OrderedMap k a
set k v orderedMap@(OrderedMap {assoc})
  | Map.member k assoc = orderedMap {assoc = Map.insert k v assoc}
  | otherwise = orderedMap

mapi :: Ord k => (k -> a -> b) -> OrderedMap k a -> OrderedMap k b
mapi f orderedMap = orderedMap {
    assoc = mapWithKey f $ assoc orderedMap
  }

{-
cons :: Ord k => k -> a -> OrderedMap k a -> OrderedMap k a
cons k a orderedMap =

snoc :: Ord k => k -> a -> OrderedMap k a -> OrderedMap k a

alter :: Ord k => (Maybe a -> Maybe a) -> k -> OrderedMap k a -> OrderedMap k a
alter 
-}