    {-# LANGUAGE DeriveFunctor #-}
    {-# LANGUAGE MultiParamTypeClasses #-}
    {-# LANGUAGE TemplateHaskell #-}
    {-# LANGUAGE DeriveGeneric #-}
    {-# LANGUAGE DeriveAnyClass #-}
    module Brick.BorderMap
        ( BorderMap
        , Edges(..)
        , eTopL, eBottomL, eRightL, eLeftL
        , empty, emptyCoordinates, singleton
        , insertH, insertV, insert
        , unsafeUnion
        , coordinates, bounds
        , values
        , lookupRow, lookupCol, lookupH, lookupV, lookup
        , setCoordinates, crop, expand
        , translate
        ) where

    import Brick.Types.Common (Edges(..), Location(..), eTopL, eBottomL, eRightL, eLeftL, origin)
    import Control.Applicative (liftA2)
    import Data.IMap (IMap, Run(Run))
    import GHC.Generics
    import Control.DeepSeq
    import Prelude hiding (lookup)
    import qualified Data.IMap as IM

    -- | Internal use only.
    neighbors :: Edges a -> Edges (a, a)
    neighbors (Edges vt vb vl vr) = Edges horiz horiz vert vert where
        horiz = (vl, vr)
        vert  = (vt, vb)

    -- Invariant: corner values are present on all the edges incident on that
    -- corner. Widthless or heightless rectangles replicate the IMaps exactly on
    -- the two coincident edges.
    --
    -- Practically speaking, this means for lookup you can look on any edge that
    -- could contain the key you care about, while for insertion you must insert on
    -- every edge that could contain the keys being inserted.

    -- | A @BorderMap a@ is like a @Map Location a@, except that there is a
    -- rectangle, and only 'Location's on the border of this rectangle are
    -- retained. The 'BorderMap' can be queried for the position and size of the
    -- rectangle. There are also efficient bulk query and bulk update operations
    -- for adjacent positions on the border.
    data BorderMap a = BorderMap
        { _coordinates :: Edges Int
        , _values :: Edges (IMap a)
        } deriving (Eq, Ord, Show, Functor, Read, Generic, NFData)

    -- | Given a rectangle (specified as the coordinates of the top, left, bottom,
    -- and right sides), initialize an empty 'BorderMap'.
    emptyCoordinates :: Edges Int -> BorderMap a
    emptyCoordinates cs = BorderMap { _coordinates = cs, _values = pure IM.empty }

    -- | An empty 'BorderMap' that only tracks the point (0,0).
    empty :: BorderMap a
    empty = emptyCoordinates (pure 0)

    -- | A 'BorderMap' that tracks only the given the point (and initially maps it
    -- to the given value).
    singleton :: Location -> a -> BorderMap a
    singleton l v = translate l . insert origin v $ empty

    {-# INLINE coordinates #-}
    -- | The positions of the edges of the rectangle whose border is retained in a
    -- 'BorderMap'. For example, if @coordinates m = e@, then the top border
    -- contains the 'Location's on row @eTop e@ and between columns @eLeft e@ to
    -- @eRight e@ inclusive.
    coordinates :: BorderMap a -> Edges Int
    coordinates = _coordinates

    -- | A complementary way to query the edges of the rectangle whose border is
    -- retained in a 'BorderMap'. For example, if @bounds m = b@, then a
    -- 'Location'\'s column must be between @fst (eTop b)@ and @snd (eTop b)@ to be
    -- retained. See also 'coordinates', which is in most cases a more natural
    -- border query.
    bounds :: BorderMap a -> Edges (Int, Int)
    bounds = neighbors . coordinates

    {-# INLINE values #-}
    -- | Maps giving the values along each edge. Corner values are replicated in
    -- all relevant edges.
    values :: BorderMap a -> Edges (IMap a)
    values = _values

    -- | Bulk insertion of horizontally-adjacent values. The 'Location' gives the
    -- start point, and the 'Run' extends in the "larger columns" direction.
    insertH :: Location -> Run a -> BorderMap a -> BorderMap a
    insertH = insertDirAgnostic (Edges insertPar insertPar insertPerp insertPerp) . swapLoc
        where
        swapLoc (Location (col, row)) = Location (row, col)

    -- | Bulk insertion of vertically-adjacent values. The 'Location' gives the
    -- start point, and the 'Run' extends in the "larger rows" direction.
    insertV :: Location -> Run a -> BorderMap a -> BorderMap a
    insertV = insertDirAgnostic (Edges insertPerp insertPerp insertPar insertPar)

    insertDirAgnostic
        :: Edges (Location -> Run a -> Int -> (Int, Int) -> IMap a -> IMap a)
        -> Location -> Run a -> BorderMap a -> BorderMap a
    insertDirAgnostic insertions l r m =
        m { _values = insertions <*> pure l <*> pure r <*> coordinates m <*> bounds m <*> _values m }

    insertPar, insertPerp :: Location -> Run a -> Int -> (Int, Int) -> IMap a -> IMap a
    insertPar (Location (kPar, kPerp)) r herePar (loPerp, hiPerp)
        | kPar == herePar && loPerp <= kPerp + IM.len r - 1 && kPerp <= hiPerp
            = IM.insert beg r { IM.len = end - beg + 1 }
        | otherwise = id
        where
        beg = max kPerp loPerp
        end = min (kPerp + IM.len r - 1) hiPerp
    insertPerp (Location (kPar, kPerp)) r herePerp (loPar, hiPar)
        | loPar <= kPar && kPar <= hiPar && kPerp <= herePerp && herePerp <= kPerp + IM.len r - 1
            = IM.insert kPar r { IM.len = 1 }
        | otherwise = id

    -- | Insert a single value at the given location.
    insert :: Location -> a -> BorderMap a -> BorderMap a
    insert l = insertV l . Run 1

    -- | Look up all values on a given row. The 'IMap' returned maps columns to
    -- values.
    lookupRow :: Int -> BorderMap a -> IMap a
    lookupRow row m
        | row == eTop    (coordinates m) = eTop    (_values m)
        | row == eBottom (coordinates m) = eBottom (_values m)
        | otherwise = IM.fromList
            $  [(eLeft   (coordinates m), Run 1 a) | Just a <- [IM.lookup row (eLeft   (_values m))]]
            ++ [(eRight  (coordinates m), Run 1 a) | Just a <- [IM.lookup row (eRight  (_values m))]]

    -- | Look up all values on a given column. The 'IMap' returned maps rows to
    -- values.
    lookupCol :: Int -> BorderMap a -> IMap a
    lookupCol col m
        | col == eLeft   (coordinates m) = eLeft   (_values m)
        | col == eRight  (coordinates m) = eRight  (_values m)
        | otherwise = IM.fromList
            $  [(eTop    (coordinates m), Run 1 a) | Just a <- [IM.lookup col (eTop    (_values m))]]
            ++ [(eBottom (coordinates m), Run 1 a) | Just a <- [IM.lookup col (eBottom (_values m))]]

    -- | Bulk lookup of horizontally-adjacent values. The 'Location' gives the
    -- starting point, and the 'Run' extends in the "larger columns" direction. The
    -- 'IMap' returned maps columns to values.
    lookupH :: Location -> Run ignored -> BorderMap a -> IMap a
    lookupH (Location (col, row)) r = IM.restrict col r . lookupRow row

    -- | Bulk lookup of vertically-adjacent values. The 'Location' gives the
    -- starting point, and the 'Run' extends in the "larger rows" direction. The
    -- 'IMap' returned maps rows to values.
    lookupV :: Location -> Run ignored -> BorderMap a -> IMap a
    lookupV (Location (col, row)) r = IM.restrict row r . lookupCol col

    -- | Look up a single position.
    lookup :: Location -> BorderMap a -> Maybe a
    lookup (Location (col, row)) = IM.lookup row . lookupCol col

    -- | Set the rectangle being tracked by this 'BorderMap', throwing away any
    -- values that do not lie on this new rectangle.
    setCoordinates :: Edges Int -> BorderMap a -> BorderMap a
    setCoordinates coordinates' m = BorderMap
        { _values = values'
        , _coordinates = coordinates'
        }
        where
        bounds' = neighbors coordinates'
        values' = pure gc
            <*> _coordinates m
            <*> coordinates'
            <*> bounds'
            <*> _values m
            <*> Edges { eTop = lookupRow, eBottom = lookupRow, eLeft = lookupCol, eRight = lookupCol }
        gc oldPar newPar (loPerp, hiPerp) imPar lookupPerp
            | oldPar == newPar = IM.restrict loPerp (Run (hiPerp-loPerp+1) ()) imPar
            | otherwise = lookupPerp newPar m

    -- | Ensure that the rectangle being tracked by this 'BorderMap' extends no
    -- farther than the given one.
    crop :: Edges Int -> BorderMap a -> BorderMap a
    crop cs m = setCoordinates (shrink <*> cs <*> coordinates m) m where
        shrink = Edges
            { eTop    = max
            , eBottom = min
            , eLeft   = max
            , eRight  = min
            }

    -- | Ensure that the rectangle being tracked by this 'BorderMap' extends at
    -- least as far as the given one.
    expand :: Edges Int -> BorderMap a -> BorderMap a
    expand cs m = setCoordinates (grow <*> cs <*> coordinates m) m where
        grow = Edges
            { eTop    = min
            , eBottom = max
            , eLeft   = min
            , eRight  = max
            }

    -- | Move a 'BorderMap' by adding the given 'Location' to all keys in the map.
    translate :: Location -> BorderMap a -> BorderMap a
    -- fast path: do nothing for (0,0)
    translate (Location (0, 0)) m = m
    translate (Location (c, r)) m = BorderMap
        { _coordinates = liftA2 (+)          cOffsets (_coordinates m)
        , _values      = liftA2 IM.addToKeys vOffsets (_values      m)
        }
        where
        cOffsets = Edges { eTop = r, eBottom = r, eLeft = c, eRight = c }
        vOffsets = Edges { eTop = c, eBottom = c, eLeft = r, eRight = r }

    -- | Assumes the two 'BorderMap's are tracking the same rectangles, but have
    -- disjoint keys. This property is not checked.
    unsafeUnion :: BorderMap a -> BorderMap a -> BorderMap a
    unsafeUnion m m' = m { _values = liftA2 IM.unsafeUnion (_values m) (_values m') }