{-# LANGUAGE TemplateHaskell, ViewPatterns, RankNTypes, GeneralizedNewtypeDeriving, DeriveTraversable, ScopedTypeVariables, FlexibleInstances, MultiParamTypeClasses #-}

module Postdelay.PrioMap
  ( PrioMap, prioMap, prioMap', _Endo
  , prio, prio', prios
  , squash

  , PrioEndo, prioEndo, prioEndo'
  ) where

import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap

import Control.Lens

import Data.Foldable
import Data.Semigroup
  

makePrisms ''Endo
  
newtype PrioMap p a = PrioMap (IntMap a)
  deriving (Functor, Foldable, Traversable)
makePrisms ''PrioMap

instance Semigroup a => Semigroup (PrioMap p a) where
  (PrioMap a) <> (PrioMap b) = PrioMap $ IntMap.unionWith (<>) a b

instance Semigroup a => Monoid (PrioMap p a) where
  mempty = PrioMap mempty
  mappend = (<>)

instance Enum p => FunctorWithIndex p (PrioMap p) where
  imap f (PrioMap intMap) = PrioMap $ imap (f . toEnum) intMap

instance Enum p => FoldableWithIndex p (PrioMap p) where
  ifoldMap f (PrioMap intMap) = ifoldMap (f . toEnum) intMap

instance Enum p => TraversableWithIndex p (PrioMap p) where
  itraverse f (PrioMap intMap) = PrioMap <$> itraverse (f . toEnum) intMap


prioMap :: (Enum p, Monoid a) => Iso' (PrioMap p a) a
-- ^ `prioMap` squashes priority information into `0`
prioMap = prioMap' $ toEnum 0

prioMap' :: (Enum p, Monoid a) => p -> Iso' (PrioMap p a) a
-- ^ `prioMap' p` squashes priority information into `p`
prioMap' (fromEnum -> p) = _PrioMap . iso fold (IntMap.singleton p)

prio :: Enum p => p -> Lens' (PrioMap p a) (Maybe a)
prio (fromEnum -> p) = _PrioMap . at p

squash :: Semigroup a => PrioMap p (PrioMap p a) -> PrioMap p a
squash = unwrapMonoid . foldMap WrapMonoid

prios :: forall f p arr a. (Bounded p, Enum p, Semigroup a, Applicative f, Indexable p arr) => arr (Maybe a) (f (Maybe a)) -> (PrioMap p a -> f (PrioMap p a))
prios (indexed -> alter) pMap = foldr cons_f (pure $ PrioMap IntMap.empty) [minBound .. maxBound]
  where
    cons_f :: p -> f (PrioMap p a) -> f (PrioMap p a)
    cons_f p x = (<>) <$> (maybe mempty (PrioMap . IntMap.singleton (fromEnum p)) <$> alter p (pMap ^. prio p)) <*> x


type PrioEndo p a = PrioMap p (Endo a)

prioEndo :: Enum p => Iso' (PrioEndo p a) (Endo a)
prioEndo = prioMap

prioEndo' :: Enum p => p -> Iso' (PrioEndo p a) (Endo a)
prioEndo' = prioMap'

prio' :: Enum p => p -> Setter' (PrioEndo p a) a
prio' p = sets $ \(Endo -> mod) -> mappend $ review (prioMap' p) mod