{-# LANGUAGE NamedFieldPuns #-}
module Data.RangeMap (
      HalfLine(..)
    , RangeMap
    , empty
    , get
    , line
    , segment
    , set
  ) where

type Color b = Maybe b
data Branch a b = Branch {
      lower :: RangeMap a b
    , bound :: a
    , weight :: Int
    , higher :: RangeMap a b
  } deriving (Show, Read)

data RangeMap a b =
    Color (Color b)
  | Tree (Branch a b)
  deriving (Show, Read)

empty :: RangeMap a b
empty = Color Nothing

line :: b -> RangeMap a b
line = Color . Just

get :: Ord a => RangeMap a b -> a -> Maybe b
get (Color color) _ = color
get (Tree (Branch {lower, bound, higher})) key =
  get (if key < bound then lower else higher) key

data HalfLine a = Above a | Below a

origin :: HalfLine a -> a
origin (Above a) = a
origin (Below a) = a

set :: Ord a => HalfLine a -> Color b -> RangeMap a b -> RangeMap a b
set halfLine color (Tree branch@(Branch {lower, bound, higher})) =
  if origin halfLine < bound
    then
      let newLower = set halfLine color lower in
      Tree $ balance (branch {
          lower = newLower
        , weight = weight branch + 1
      })
    else
      let newHigher = set halfLine color higher in
      Tree $ balance (branch {
          higher = newHigher
        , weight = weight branch + 1
      })

set (Above bound) color rangeMap = Tree (Branch {
      lower = rangeMap
    , bound
    , weight = 0
    , higher = Color color
  })

set (Below bound) color rangeMap = Tree (Branch {
      lower = Color color
    , bound
    , weight = 0
    , higher = rangeMap
  })

segment :: (a, a) -> b -> RangeMap a b
segment (from, to) b = Tree (Branch {
      lower = empty
    , bound = from
    , weight = 1
    , higher = Tree (Branch {
        lower = line b
      , bound = to
      , weight = 0
      , higher = empty
    })
  })

popMax :: Branch a b -> ((a, Color b), RangeMap a b)
popMax (Branch {lower, bound, higher = Color color}) = ((bound, color), lower)
popMax branch@(Branch {higher = Tree higherBranch, weight}) =
  let (pair, subTree) = popMax higherBranch in
  (pair, Tree (branch {higher = subTree, weight = weight - 1}))

popMin :: Branch a b -> ((a, Color b), RangeMap a b)
popMin (Branch {lower = Color color, bound, higher}) = ((bound, color), higher)
popMin branch@(Branch {lower = Tree lowerBranch, weight}) =
  let (pair, subTree) = popMin lowerBranch in
  (pair, Tree (branch {lower = subTree, weight = weight - 1}))

rotateLeft :: Ord a => Branch a b -> Branch a b
rotateLeft branch@(Branch {higher = Color _}) = branch
rotateLeft branch@(Branch {lower, bound, higher = Tree higherBranch}) =
  let ((minBound, color), newHigher) = popMin higherBranch in
  branch {
      lower = set (Above bound) color lower
    , bound = minBound
    , higher = newHigher
  }

rotateRight :: Ord a => Branch a b -> Branch a b
rotateRight branch@(Branch {lower = Color _}) = branch
rotateRight branch@(Branch {lower = Tree lowerBranch, bound, higher}) =
  let ((maxBound, color), newLower) = popMax lowerBranch in
  branch {
      lower = newLower
    , bound = maxBound
    , higher = set (Below bound) color higher
  }

getWeight :: RangeMap a b -> Int
getWeight (Tree (Branch {weight})) = weight
getWeight _ = 0

balance :: Ord a => Branch a b -> Branch a b
balance branch@(Branch {lower, bound, higher})
  | lowerWeight - higherWeight > 1 = rotateRight branch
  | higherWeight - lowerWeight > 1 = rotateLeft branch
  | otherwise = branch
  where
    lowerWeight = getWeight lower
    higherWeight = getWeight higher