module Data.IntervalSet
    ( IntSet
    , null
    , empty
    , insert
    , delete
    , interval
    , toIntervals
    , nearestOutsider
    , Data.IntervalSet.lookup
    ) where

import Prelude hiding (null)
import qualified Data.IntMap.Strict as IntMap
         ;import Data.IntMap.Strict (IntMap)
import qualified Data.List as List
import Data.Ord


-- A set of integers.
newtype IntSet = IntSet (IntMap Interval)
 deriving Show

-- Note: the intervalMin is not stored here but is the lookup key in an IntMap.
data Interval = Interval
    { intervalMax  :: {-# UNPACK #-} !Int -- ^ Maximum value contained in this interval.
    , intervalNext :: {-# UNPACK #-} !Int -- ^ Minimum value in next interval if there is one.
    }
 deriving Show

null :: IntSet -> Bool
null (IntSet m) = IntMap.null m

empty :: IntSet
empty = IntSet IntMap.empty


insert :: Int -> IntSet -> IntSet
insert x (IntSet m) = IntSet $ case IntMap.lookupLE x m of
    Just (lb,Interval mx ub)
        | x <= mx   -> m
        | otherwise -> case ub == maxBound of

            True    | x == mx + 1 -> IntMap.insert lb (Interval x maxBound) m
                    | otherwise   -> IntMap.insert lb (Interval mx x)
                                     $ IntMap.insert x (Interval x maxBound) m

            False   | mx + 2 == ub -> let (Just v', m')
                                            = IntMap.updateLookupWithKey (\_ _ -> Nothing) ub m
                                      in IntMap.insert lb v' m'
                    | mx + 1 == x  -> IntMap.insert lb (Interval x ub) m
                    | otherwise -> IntMap.insert lb (Interval mx x)
                                   $ if ub == x + 1
                                        then let (Just v', m')
                                                   = IntMap.updateLookupWithKey
                                                        (\_ _ -> Nothing) ub m
                                             in IntMap.insert x v' m'
                                        else IntMap.insert x (Interval x ub) m

    Nothing     -> case IntMap.minViewWithKey m of

        Just ((ub,v),m')
            | x + 1 == ub -> IntMap.insert x v m'
            | otherwise   -> IntMap.insert x (Interval x ub) m

        Nothing           -> IntMap.singleton x (Interval x maxBound)

member :: Int -> IntSet -> Bool
member x (IntSet m) = case IntMap.lookupLE x m of
    Just (lb,Interval mx _) -> x <= mx
    Nothing                 -> False

nearestOutsider :: Int -> IntSet -> Maybe Int
nearestOutsider x (IntSet m)
    | List.null xs = Nothing -- There are no integers outside the set!
    | otherwise    = Just $ List.minimumBy (comparing (\y -> abs (x - y))) xs
 where
    xs = case IntMap.lookupLE x m of
            Nothing -> [x]
            Just (lb,Interval mx ub)
             -> if ub < maxBound
                    then case () of
                     () | x > mx        -> [x]
                        | minBound < lb -> [lb-1, mx+1, ub-1]
                        | otherwise     -> [mx+1, ub-1]
                    else case () of
                     () | x > mx                         -> [x]
                        | minBound < lb && mx < maxBound -> [lb-1, mx+1]
                        | minBound < lb                  -> [lb-1]
                        | mx < maxBound                  -> [mx+1]
                        | otherwise                      -> []

-- Note this could possibly benefit from a intervalPrev field.
delete :: Int -> IntSet -> IntSet
delete x (IntSet m) = IntSet $ case IntMap.lookupLE x m of
    Nothing -> m
    Just (lb,Interval mx nxt) -> case compare x mx of

        GT -> m

        EQ | lb < mx   -> IntMap.insert lb (Interval (mx - 1) nxt) m
           | otherwise -> case IntMap.lookupLE (x-1) m of -- no intervalPrev
                    Just (lb',Interval mx' _) -> IntMap.insert lb' (Interval mx' nxt)
                                                 $ IntMap.delete lb m
                    Nothing                   -> IntMap.delete lb m

        LT -> let m' = IntMap.insert (x+1) (Interval mx nxt) m
              in if lb < x
                then IntMap.insert lb (Interval (x - 1) (x+1)) m'
                else if x == minBound
                    then IntMap.delete minBound m'
                    else case IntMap.lookupLE (x-1) m' of -- no intervalPrev
                        Just (lb',Interval mx' _) -> IntMap.insert lb' (Interval mx' (x+1))
                                                     $ IntMap.delete lb m'
                        Nothing                   -> IntMap.delete lb m'

toIntervals :: IntSet -> [(Int,Int)]
toIntervals (IntSet m) = List.map (\(lb,(Interval mx _)) -> (lb,mx))
                                    $ IntMap.toList m

interval :: Int -> Int -> IntSet
interval lb mx
    | lb <= mx  = IntSet $ IntMap.singleton lb (Interval mx maxBound)
    | otherwise = IntSet IntMap.empty

lookup :: Int -> IntSet -> Maybe (Int,Int)
lookup k (IntSet m) = case IntMap.lookupLE k m of
    Nothing                 -> Nothing
    Just (lb,Interval mx _) -> Just (lb,mx)