content: -----------------------------------------------------------------------
  -- |
  -- Module           : Data.IntervalsMap
  -- Description      : Nested intervals
  -- Copyright        : (c) Galois, Inc 2020
  -- License          : BSD3
  -- Maintainer       : Daniel Matichuk <dmatichuk@galois.com>
  -- Stability        : provisional
  ------------------------------------------------------------------------
  
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE FunctionalDependencies #-}
  
  module Data.Parameterized.IntervalsMap
    ( IntervalF(..)
    , mkIntervalF
    , Intervals(..)
    , IntervalsMap
    , intersecting
    , unionWith
    , unionWithM
    , singleton
    , insertWith
    , insertWithM
    , intersectionWith
    , mapMIntersecting
    , fromList
    , toList
    , empty
    , IM.Interval(..)
    , mergeIntervalsF
    , mergeWithM
    , AsOrd(..)
    ) where
  
  
  import           Data.Kind ( Type )
  import           Data.Maybe (catMaybes)
  
  import           Data.IntervalMap.Strict ( IntervalMap )
  import qualified Data.IntervalMap.Strict as IM
  import qualified Data.IntervalMap.Interval as IM
  import qualified Data.IntervalMap.Generic.Strict as IMG
  
  import           Data.Parameterized.Classes
  import qualified Data.Parameterized.Context as Ctx
  
  newtype AsOrd f tp where
    AsOrd :: { unAsOrd :: f tp } -> AsOrd f tp
  
  instance TestEquality f => Eq (AsOrd f tp) where
    (AsOrd a) == (AsOrd b) = case testEquality a b of
      Just Refl -> True
      _ -> False
  
  instance OrdF f => Ord (AsOrd f tp) where
    compare (AsOrd a) (AsOrd b) = toOrdering $ compareF a b
  
  newtype IntervalF f tp where
    IntervalF :: IM.Interval (AsOrd f tp) -> IntervalF f tp
  
  mkIntervalF ::
    IM.Interval (f tp) -> IntervalF f tp
  mkIntervalF ival = IntervalF $ fmap AsOrd ival
  
  instance TestEquality f => TestEquality (IntervalF f) where
    testEquality (IntervalF i1) (IntervalF i2) = case testEquality (unAsOrd (IM.lowerBound i1)) (unAsOrd (IM.lowerBound i2)) of
      Just Refl | i1 == i2 -> Just Refl
      _ -> Nothing
  
  deriving instance TestEquality f => Eq (IntervalF f tp)
  
  deriving instance OrdF f => Ord (IntervalF f tp)
  
  newtype Intervals f ctx = Intervals (Ctx.Assignment (IntervalF f) ctx)
  
  deriving instance TestEquality f => Eq (Intervals f ctx)
  
  instance OrdF f => Ord (Intervals f ctx) where
    compare (Intervals (rest1 Ctx.:> a1)) (Intervals (rest2 Ctx.:> a2)) =
      compare a1 a2 <> compare (Intervals rest1) (Intervals rest2)
    compare (Intervals Ctx.Empty) (Intervals Ctx.Empty) = EQ
  
  data IntervalsMap (f :: k -> Type) (ctx :: Ctx.Ctx k) tp where
    IntervalsMapCons ::
      IntervalMap (AsOrd f idx) (IntervalsMap f ctx tp) ->
      IntervalsMap f (ctx Ctx.::> idx) tp
    IntervalsMapHead :: tp -> IntervalsMap f Ctx.EmptyCtx tp
  
  instance Functor (IntervalsMap f ctx) where
    fmap f ims = case ims of
      IntervalsMapCons ims' -> IntervalsMapCons (fmap (fmap f) ims')
      IntervalsMapHead v -> IntervalsMapHead $ f v
  
  instance Foldable (IntervalsMap f ctx) where
    foldMap f (IntervalsMapCons ims') = foldMap (foldMap f) ims'
    foldMap f (IntervalsMapHead v) = f v
  
  instance Traversable (IntervalsMap f ctx) where
    traverse f (IntervalsMapCons ims') = IntervalsMapCons <$> traverse (traverse f) ims'
    traverse f (IntervalsMapHead v) = IntervalsMapHead <$> f v
  
  intersecting ::
    OrdF f =>
    IntervalsMap f ctx tp ->
    Intervals f ctx ->
    IntervalsMap f ctx tp
  intersecting (IntervalsMapCons ims) (Intervals (rest Ctx.:> IntervalF k)) =
    let
      top = IM.intersecting ims k
    in IntervalsMapCons $ fmap (\ims' -> intersecting ims' (Intervals rest)) top
  intersecting v (Intervals Ctx.Empty) = v
  
  fromList ::
    OrdF f =>
    [(Intervals f (ctx Ctx.::> a), tp)] ->
    IntervalsMap f (ctx Ctx.::> a) tp
  fromList es = foldr (unionWith (\l _ -> l)) empty (map (uncurry singleton) es)
  
  toList ::
    IntervalsMap f ctx tp ->
    [(Intervals f ctx, tp)]
  toList (IntervalsMapCons ims) =
    concat $ map (\(k, es) -> addTo k (toList es)) $ (IM.toList ims)
    where
      addTo :: IM.Interval (AsOrd f a) -> [(Intervals f ctx, tp)] -> [(Intervals f (ctx Ctx.::> a), tp)]
      addTo ival = map (\(Intervals ivalf, a) -> (Intervals $ ivalf Ctx.:> IntervalF ival, a))
  toList (IntervalsMapHead v) = [(Intervals Ctx.empty, v)]
  
  unionWith ::
    OrdF f =>
    (a -> a -> a) ->
    IntervalsMap f ctx a ->
    IntervalsMap f ctx a ->
    IntervalsMap f ctx a
  unionWith f (IntervalsMapCons ims1) (IntervalsMapCons ims2) =
    IntervalsMapCons $ IM.unionWith (unionWith f) ims1 ims2
  unionWith f (IntervalsMapHead v1) (IntervalsMapHead v2) = IntervalsMapHead $ f v1 v2
  
  unionWithM ::
    forall f m a ctx.
    OrdF f =>
    Monad m =>
    (a -> a -> m a) ->
    IntervalsMap f ctx a ->
    IntervalsMap f ctx a ->
    m (IntervalsMap f ctx a)
  unionWithM f ims1 ims2 = sequenceA $ unionWith go (fmap return ims1) (fmap return ims2)
    where
      go :: m a -> m a -> m a
      go f1 f2 = do
        v1 <- f1
        v2 <- f2
        f v1 v2
  
  data MergeResult a b =
      MergeLeft a
    | MergeRight b
    | MergeCombined a b
  
  mergeWithM ::
    forall f m a b c ctx.
    OrdF f =>
    Monad m =>
    (a -> m c) ->
    (b -> m c) ->
    (a -> b -> m c) ->
    IntervalsMap f ctx a ->
    IntervalsMap f ctx b ->
    m (IntervalsMap f ctx c)
  mergeWithM inLeft inRight combine ims1 ims2 = do
    traverse eval $ unionWith go (fmap MergeLeft ims1) (fmap MergeRight ims2)
    where
      eval :: MergeResult a b -> m c
      eval (MergeLeft a) = inLeft a
      eval (MergeRight b) = inRight b
      eval (MergeCombined a b) = combine a b
  
      go :: MergeResult a b -> MergeResult a b -> MergeResult a b
      go (MergeLeft f1) (MergeRight f2) = MergeCombined f1 f2
      go _ _ = error "mergeWithM: unexpected MergeResult"
  
  
  singleton ::
    Intervals f ctx ->
    tp ->
    IntervalsMap f ctx tp
  singleton (Intervals (rest Ctx.:> IntervalF k)) v = IntervalsMapCons $ IM.singleton k (singleton (Intervals rest) v)
  singleton (Intervals Ctx.Empty) v = IntervalsMapHead v
  
  empty :: IntervalsMap f (ctx Ctx.::> a) tp
  empty = IntervalsMapCons IM.empty
  
  insertWith ::
    OrdF f =>
    (tp -> tp -> tp) ->
    Intervals f ctx ->
    tp ->
    IntervalsMap f ctx tp ->
    IntervalsMap f ctx tp
  insertWith f k v = unionWith f (singleton k v)
  
  
  insertWithM ::
    forall m f ctx tp.
    Monad m =>
    OrdF f =>
    (tp -> tp -> m tp) ->
    Intervals f ctx ->
    tp ->
    IntervalsMap f ctx tp ->
    m (IntervalsMap f ctx tp)
  insertWithM f k v ims = sequenceA $ insertWith go k (return v) (fmap return ims)
    where
      go :: m tp -> m tp -> m tp
      go f1 f2 = do
        v1 <- f1
        v2 <- f2
        f v1 v2
  
  intersectionWith ::
    OrdF f =>
    (a -> b -> c) ->
    IntervalsMap f ctx a ->
    IntervalsMap f ctx b ->
    IntervalsMap f ctx c
  intersectionWith f (IntervalsMapCons ims1) (IntervalsMapCons ims2) =
    IntervalsMapCons $ IM.intersectionWith (intersectionWith f) ims1 ims2
  intersectionWith f (IntervalsMapHead v1) (IntervalsMapHead v2) = IntervalsMapHead $ f v1 v2
  
  mapMIntersecting' ::
    Monad m =>
    OrdF f =>
    Intervals f ctx ->
    (tp -> m (Maybe tp)) ->
    IntervalsMap f ctx tp ->
    m (Maybe (IntervalsMap f ctx tp))
  mapMIntersecting' (Intervals (rest Ctx.:> IntervalF k)) f (IntervalsMapCons ims) = do
    ims' <- mapMIntersectingBase k (\_ -> mapMIntersecting' (Intervals rest) f) ims
    case IM.size ims' of
      0 -> return Nothing
      _ -> return $ Just (IntervalsMapCons ims')
  mapMIntersecting' (Intervals Ctx.Empty) f (IntervalsMapHead v) = fmap IntervalsMapHead <$> f v
  
  -- | Adjust entries which intersect the given interval
  mapMIntersecting ::
    Monad m =>
    OrdF f =>
    Intervals f (ctx Ctx.::> a) ->
    (tp -> m (Maybe tp)) ->
    IntervalsMap f (ctx Ctx.::> a) tp ->
    m (IntervalsMap f (ctx Ctx.::> a) tp)
  mapMIntersecting i f ims = mapMIntersecting' i f ims >>= \case
    Just ims' -> return ims'
    Nothing -> return $ IntervalsMapCons IM.empty
  
  
  mapMIntersectingBase ::
    forall k v e m.
    Monad m =>
    IMG.Interval k e =>
    Ord k =>
    k ->
    (k -> v -> m (Maybe v)) ->
    IMG.IntervalMap k v ->
    m (IMG.IntervalMap k v)
  mapMIntersectingBase k f im = do
    let (pref, inter, suf) = IM.splitIntersecting im k
    case IM.size inter of
      0 -> return im
      _ -> do
        im' <- catMaybes <$> mapM go (IM.toAscList inter)
        return $ IM.fromDistinctAscList (IM.toAscList pref ++ im' ++ IM.toAscList suf)
    where
      go :: (k, v) -> m (Maybe (k, v))
      go (k', v) = f k' v >>= \case
        Just v' -> return $ Just (k', v')
        Nothing -> return Nothing
  
  
  mergeIntervals ::
    Ord a =>
    IM.Interval a ->
    IM.Interval a ->
    IM.Interval a
  mergeIntervals i1 i2 = case (leftClosed, rightClosed) of
    (True, True) -> IM.ClosedInterval lower upper
    (False, True) -> IM.IntervalOC lower upper
    (True, False) -> IM.IntervalCO lower upper
    (False, False) -> IM.OpenInterval lower upper
    where
      leftClosed = (IM.leftClosed i1 && lo1 <= lo2) || (IM.leftClosed i2 && lo2 <= lo1)
      rightClosed = (IM.rightClosed i1 && hi2 <= hi1) || (IM.rightClosed i2 && hi1 <= hi2)
      lo1 = IM.lowerBound i1
      lo2 = IM.lowerBound i2
      hi1 = IM.upperBound i1
      hi2 = IM.upperBound i2
      lower = min lo1 lo2
      upper = max hi1 hi2
  
  mergeIntervalsF ::
    OrdF f =>
    IntervalF f a ->
    IntervalF f a ->
    IntervalF f a
  mergeIntervalsF (IntervalF i1) (IntervalF i2) = IntervalF (mergeIntervals i1 i2)
  