Layers.hs

module Data.Interval.Layers (
  Layers (Layers),
  Data.Interval.Layers.fromList,
  Data.Interval.Layers.toList,
  empty,
  singleton,
  insert,
  pile,
  squash,
  squashing,
  isquashing,
  land,
  landAbove,
  thickness,
  thickest,
  dig,
  remove,
  (\-),
  baseline,
  difference,
  truncate,
  (\=),
  toStepFunction,
  integrate,

  -- ** Helper functions
  nestings,
) where

import Algebra.Lattice.Levitated (Levitated (Top))
import Data.Data (Data)
import Data.Foldable qualified as Foldable
import Data.Foldable.WithIndex (FoldableWithIndex, ifoldMap)
import Data.Functor.WithIndex
import Data.Group (Group (..))
import Data.Heap (Heap)
import Data.Heap qualified as Heap
import Data.Interval (
  Adjacency (..),
  Interval,
  OneOrTwo (..),
  pattern Whole,
  pattern (:---:),
  pattern (:|-|:),
 )
import Data.Interval qualified as Interval
import Data.Interval.Borel (Borel)
import Data.Interval.Borel qualified as Borel
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Traversable.WithIndex (TraversableWithIndex (itraverse))
import GHC.Generics (Generic)
import Prelude hiding (truncate)

-- The 'Layers' of an ordered type @x@ are like the 'Borel' sets,
-- but that keeps track of how far each point has been "raised" in @y@.
newtype Layers x y = Layers (Map (Interval x) y)
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable, Generic, Data)
deriving newtype instance FunctorWithIndex (Interval x) (Layers x)
deriving newtype instance FoldableWithIndex (Interval x) (Layers x)
instance TraversableWithIndex (Interval x) (Layers x) where
  itraverse ::
    (Applicative f) =>
    (Interval x -> a -> f b) -> Layers x a -> f (Layers x b)
  itraverse f (Layers s) = Layers <$> itraverse f s

instance (Ord x, Ord y, Semigroup y) => Semigroup (Layers x y) where
  (<>) :: (Ord x, Ord y, Semigroup y) => Layers x y -> Layers x y -> Layers x y
  Layers s1 <> Layers s2 =
    Layers
      . Map.fromAscList
      . nestingsAsc
      . Heap.fromList
      $ Map.toAscList (Map.unionWith (<>) s1 s2)

instance (Ord x, Ord y, Semigroup y) => Monoid (Layers x y) where
  mempty :: (Ord x, Ord y, Semigroup y) => Layers x y
  mempty = Layers mempty

instance (Ord x, Ord y, Group y) => Group (Layers x y) where
  invert :: (Ord x, Ord y, Group y) => Layers x y -> Layers x y
  invert (Layers s) = Layers (fmap invert s)

-- | A blank canvas.
empty :: Layers x y
empty = Layers Map.empty

-- | @singleton ix y@ is the rectangle with base @ix@ of thickness @y@.
singleton :: (Ord x) => Interval x -> y -> Layers x y
singleton ix y = Layers (Map.singleton ix y)

-- | Draw the 'Layers' of specified bases and thicknesses.
fromList :: (Ord x, Ord y, Semigroup y) => [(Interval x, y)] -> Layers x y
fromList = Layers . Map.fromList . nestings

-- | Get all of the bases and thicknesses in the 'Layers'.
toList :: (Ord x) => Layers x y -> [(Interval x, y)]
toList (Layers s) = Map.toList s

-- | Ignore the 'Layers' and focus only on whether points are 'Data.Interval.within'
-- any contained 'Interval' or not.
squash :: (Ord x) => Layers x y -> Borel x
squash (Layers s) = foldMap Borel.singleton (Map.keys s)

-- | 'squash' together the intervals satisfying a predicate.
squashing :: (Ord x) => (y -> Bool) -> Layers x y -> Borel x
squashing = isquashing . const

-- | Perform 'squashing' with a test that accepts the 'Interval' as an argument.
isquashing :: (Ord x) => (Interval x -> y -> Bool) -> Layers x y -> Borel x
isquashing f s = flip ifoldMap s \ix y -> if f ix y then Borel.singleton ix else Borel.empty

-- | Treating 'mempty' as sea level, consider the 'Borel' set of a provided
-- 'Layers' that is "land".
--
-- An improvement over 'squash' in that it will not return 'I.Whole'
-- if 'baseline' or some involved interval calculations have been used.
land :: (Ord x, Monoid y, Ord y) => Layers x y -> Borel x
land = landAbove mempty

-- | Given a "sea level", consider the 'Borel' set of a provided 'Layers'
-- that is "land".
--
-- An improvement over 'squash' in that it will not return 'I.Whole'
-- if 'baseline' or some involved interval calculations have been used.
landAbove :: (Ord x, Ord y) => y -> Layers x y -> Borel x
landAbove sea (Layers s) = flip Map.foldMapWithKey s \i y ->
  if y > sea then Borel.singleton i else Borel.empty

-- | @insert ix y l@ draws over @l@ a rectangle with base @ix@ of thickness @y@.
insert ::
  (Ord x, Ord y, Semigroup y) =>
  Interval x ->
  y ->
  Layers x y ->
  Layers x y
insert ix y = (<>) (singleton ix y)

-- | Flipped synonym for 'insert'.
-- Mnemonic: "pile" this much onto the existing 'Layers'
-- over the given 'Interval'.
pile ::
  (Ord x, Ord y, Semigroup y) =>
  y ->
  Interval x ->
  Layers x y ->
  Layers x y
pile = flip insert

-- | Take away a thickness over a given base from the 'Layers'.
dig :: (Ord x, Ord y, Group y) => y -> Interval x -> Layers x y -> Layers x y
dig y ix = insert ix (invert y)

-- | Completely remove an 'Interval' from the 'Layers'.
remove :: (Ord x, Ord y, Semigroup y) => Interval x -> Layers x y -> Layers x y
remove ix (Layers s) = flip (`Map.foldlWithKey'` empty) s \acc jx y ->
  acc <> case jx Interval.\\ ix of
    Nothing -> mempty
    Just (One kx) -> singleton kx y
    Just (Two kx lx) -> fromList [(kx, y), (lx, y)]

-- | Fliped infix version of 'remove'.
(\-) :: (Ord x, Ord y, Semigroup y) => Layers x y -> Interval x -> Layers x y
(\-) = flip remove

-- | Add the given thickness to every point.
baseline :: (Ord x, Ord y, Semigroup y) => y -> Layers x y -> Layers x y
baseline = insert Whole

-- | "Excavate" the second argument from the first.
difference :: (Ord x, Ord y, Group y) => Layers x y -> Layers x y -> Layers x y
difference layers (Layers s) =
  foldr (uncurry (flip dig)) layers (Map.toAscList s)

-- | Restrict the range of the 'Layers' to the given 'Interval'.
truncate ::
  (Ord x, Ord y, Semigroup y) => Interval x -> Layers x y -> Layers x y
truncate ix (Layers s) =
  flip (`Map.foldlWithKey'` empty) s \acc jx y ->
    maybe id (`insert` y) (Interval.intersect ix jx) acc

-- | Flipped infix version of 'truncate'.
(\=) :: (Ord x, Ord y, Semigroup y) => Layers x y -> Interval x -> Layers x y
(\=) = flip truncate

-- |
-- @'integrate' diff hgt ix l@ calculates the area under the 'Interval' @ix@
-- using the measure @diff@ of the interval multiplied by the height @hgt@
-- of the layers over each sub-interval in the layers.
integrate ::
  (Ord x, Ord y, Semigroup y, Num z) =>
  (x -> x -> z) ->
  (y -> z) ->
  Interval x ->
  Layers x y ->
  Maybe z
integrate diff hgt ix layers =
  let Layers (Map.assocs -> s) = layers \= ix
      f (jx, y) maccum = do
        acc <- maccum
        d <- Interval.measuring diff jx
        pure $ acc + d * hgt y
   in foldr f (Just 0) s

-- | Get the thickness of the 'Layers' at a point.
thickness :: (Ord x, Semigroup y) => Levitated x -> Layers x y -> Maybe y
thickness x (Layers s) = case Map.lookupLE (x :|-|: x) s of
  Just (ix, y) | x `Interval.within` ix -> Just y
  _ -> Nothing

-- | Where and how thick is the thickest 'Interval'?
thickest :: (Ord x, Ord y) => Layers x y -> Maybe (Interval x, y)
thickest (Layers s) =
  flip (`Map.foldlWithKey'` Nothing) s \acc ix y -> Just case acc of
    Nothing -> (ix, y)
    Just (_, y') | y > y' -> (ix, y)
    Just (ix', y') -> (ix', y')

-- | Convert the 'Layers' into a list of beginning-points and heights,
-- that define a step function piecewise.
toStepFunction :: (Ord x, Ord y, Monoid y) => Layers x y -> [(Levitated x, y)]
toStepFunction = go . Data.Interval.Layers.toList
 where
  go = \case
    [(il :---: iu, iy), (j@(jl :---: Top), jy)]
      | iu == jl -> (il, iy) : go [(j, jy)]
      | otherwise -> (il, iy) : (iu, mempty) : go [(j, jy)]
    (il :---: iu, iy) : (j@(jl :---: _), jy) : is
      | iu == jl -> (il, iy) : go ((j, jy) : is)
      | otherwise -> (il, iy) : (iu, mempty) : go ((j, jy) : is)
    [(il :---: iu, iy)] -> [(il, iy), (iu, mempty)]
    [] -> []

nestings ::
  (Ord x, Ord y, Semigroup y) =>
  [(Interval x, y)] ->
  [(Interval x, y)]
nestings = nestingsAsc . Heap.fromList

nestingsAsc ::
  (Ord x, Ord y, Semigroup y) =>
  Heap (Interval x, y) ->
  [(Interval x, y)]
nestingsAsc heap = case firstTwo of
  Nothing -> Foldable.toList heap
  Just ((i', iy), (j', jy), js) -> case Interval.adjacency i' j' of
    Before i j -> (i, iy) : nestingsAsc (Heap.insert (j, jy) js)
    Meets i j k ->
      (i, iy) : nestingsAsc (Heap.fromList [(j, iy <> jy), (k, jy)] <> js)
    Overlaps i j k ->
      nestingsAsc do
        Heap.fromList [(i, iy), (j, iy <> jy), (k, jy)] <> js
    Starts i j ->
      nestingsAsc do
        Heap.fromList [(i, iy <> jy), (j, jy)] <> js
    During i j k ->
      nestingsAsc do
        Heap.fromList [(i, jy), (j, iy <> jy), (k, jy)] <> js
    Finishes i j ->
      nestingsAsc do
        Heap.fromList [(i, iy), (j, iy <> jy)] <> js
    Identical i -> nestingsAsc (Heap.insert (i, iy <> jy) js)
    FinishedBy i j ->
      nestingsAsc do
        Heap.fromList [(i, iy), (j, iy <> jy)] <> js
    Contains i j k ->
      nestingsAsc do
        Heap.fromList [(i, iy), (j, iy <> jy), (k, iy)] <> js
    StartedBy i j ->
      nestingsAsc do
        Heap.fromList [(i, iy <> jy), (j, iy)] <> js
    OverlappedBy i j k ->
      nestingsAsc do
        Heap.fromList [(i, jy), (j, iy <> jy), (k, iy)] <> js
    MetBy i j k ->
      (i, jy) : nestingsAsc (Heap.fromList [(j, iy <> jy), (k, iy)] <> js)
    After i j -> (i, jy) : nestingsAsc (Heap.insert (j, iy) js)
 where
  firstTwo = do
    (min1, heap') <- Heap.uncons heap
    (min2, heap'') <- Heap.uncons heap'
    pure (min1, min2, heap'')