{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}
-- |
-- A Word64RangeMap is actually a sorted array of arrays
-- but lookup is done via binary search testing the range.
-- One way of implementing a sparse array i suppose.
module Data.Word64RangeMap where

import Data.Word
import Data.Array.Unboxed
import qualified Data.Array.MArray as MA
import Data.Array.MArray (MArray(..))
import Data.Reference
import Debug.Trace
import Control.Concurrent.STM
import Data.Array.IO
import Data.IORef

type OuterIndex = Int
type Index = Word64
type InnerArray b = UArray Index b

-- | Although this type includes a parameter for index, the code assumes bounds start at 0
-- and the index has 'Integral', and 'Num' instances.
newtype RefArray r ma i e = RefArray (r (ma i e))

-- convenient contraint kind
type RangeArray rangeArray m b ref = (IArray UArray b, MArray rangeArray (InnerArray b) m, Reference ref m)
-- The RangeMap type, to be used with the above constraint
type RangeMap rangeArray b ref = RefArray ref rangeArray OuterIndex (InnerArray b)

insertArrayAt :: (MA.MArray ma e m, Reference r m, Ix i, Num i, Integral i) =>  RefArray r ma i e -> i -> e -> m ()
insertArrayAt (RefArray ref) pos e = do
    ma <- readRef ref
    (0,n) <- MA.getBounds ma -- irrefutable pattern, assume array starts at 0
    (esBefore,esAt) <- splitAt (fromIntegral pos) <$> (getElems ma)
    newMA <- MA.newListArray (0,n+1) (esBefore ++ [e] ++ esAt)
    writeRef ref newMA

deleteArrayAt :: (MA.MArray ma e m, Reference r m, Ix i, Num i, Integral i) =>  RefArray r ma i e -> i -> m ()
deleteArrayAt (RefArray ref) pos = do
    ma <- readRef ref
    (0,n) <- MA.getBounds ma -- irrefutable pattern, assume array starts at 0
    (esBefore,esAt) <- splitAt (fromIntegral pos) <$> getElems ma
    newMA <- MA.newListArray (0,n-1) (esBefore ++ drop 1 esAt)
    writeRef ref newMA


-- forall (ma :: * -> * -> *) i e (r :: * -> *) (m :: * -> *).
readRefArray :: Reference r m => RefArray r ma i e -> m (ma i e)
readRefArray (RefArray x) = readRef x

-- forall (ma :: * -> * -> *) i e (r :: * -> *) (m :: * -> *).
writeRefArray :: Reference r m => RefArray r ma i e -> ma i e -> m ()
writeRefArray (RefArray x) y = writeRef x y

{-
instance (Reference r m, MArray ma e m) => MArray (RefArray r ma) e m where
    getBounds (RefArray r) = readRef r >>= Base.getBounds
    getNumElements (RefArray r) = readRef r >>= Base.getNumElements
    unsafeRead (RefArray r) i = readRef r >>= \x -> Base.unsafeRead x i
    unsafeWrite (RefArray r) i e = modifyRef r (\x -> Base.unsafeWrite x i e >> return (x,()))
-}

emptySTMRangeMap :: STM (RangeMap TArray Word8 TVar)
emptySTMRangeMap = RefArray <$>
                     (newTVar =<<
                        newListArray (0,-1) [])

-- | a sample RangeMap for easier debugging
getX :: IO (RangeMap IOArray Word8 IORef)
getX = RefArray <$>
        (newIORef =<<
             newListArray (0,2)
                          [ listArray (15,20) [15..20]
                          , listArray (100,105) [100..105]
                          , listArray (106,107) [106..107]
                          ])
-- | a sample RangeMap for easier debugging
getY :: IO (RangeMap IOArray Word8 IORef)
getY = RefArray <$>
        (newIORef =<<
             newListArray (0,3)
                          [ listArray (0,5) [0..5]
                          , listArray (15,20) [15..20]
                          , listArray (100,105) [100..105]
                          , listArray (106,107) [106..107]
                          ])

lookupArray :: RangeArray ra m b ref => Index -> RangeMap ra b ref -> m (Either OuterIndex (OuterIndex,InnerArray b))
lookupArray i r = do
    ra <- readRefArray r
    (zr,nr) <- getBounds ra -- bounds of mutable range array
    let (dr,mr) = (nr+1 - zr) `divMod` 2
    lookupArrayHelper i r (zr+dr+mr-1) zr nr

lookupArrayHelper :: RangeArray ra m b ref => Index -> RangeMap ra b ref
                                       -> OuterIndex{- current position  -}
                                       -> OuterIndex{- smallest possible -}
                                       -> OuterIndex{- largest possible  -}
                                       -> m (Either OuterIndex (OuterIndex,InnerArray b))
lookupArrayHelper i r pos min max = do
    {- trace ("lookupArrayHelper " ++ show i ++ " "
                                ++ "ra "
                                ++ show pos ++ " "
                                ++ show min ++ " "
                                ++ show max) (return ()) -}
    ra <- readRefArray r
    ca <- MA.readArray ra pos -- current array
    let beforeOrAfter i (x,y) | i < x = LT
        beforeOrAfter i (x,y) | i > y = GT
        beforeOrAfter _ _ = EQ
        avg a b = (a + b) `divMod` 2
        isUnequalMod2 a b = if a `mod` 2 /= b `mod` 2 then 1 else 0
        f a b recurse finish
                = let (pos',r) = avg a b
                      u = isUnequalMod2 r (max - min) {- trace ("u = " ++ show t
                                                                       ++ " min=" ++ show min
                                                                       ++ " max=" ++ show max
                                                                       ++ " pos'=" ++ show pos'
                                                                       ++ " r=" ++ show r
                                                                       ) t -}
                      in if max - min + u > 1 && pos' /= pos then recurse pos' r else finish
        s = "beforeOrAfter " ++ show i ++ " " ++ show (bounds ca) ++ " --> "
        checkAndReturn pos = do
            boundsRA <- getBounds ra
            case beforeOrAfter pos boundsRA of
                EQ -> do
                    a <- MA.readArray ra pos
                    {- trace ("checkAndReturn " ++ show pos ++ " bounds a= " ++ show (bounds a)) (return ()) -}
                    case beforeOrAfter i (bounds a) of
                        EQ -> {- trace "EQ" $ -} return (Right (pos,a))
                        LT -> {- trace "LT" $ -} return (Left pos)
                        GT -> {- trace "GT" $ -} return (Left (pos+1))
                LT -> {- trace "LT" $ -} return (Left pos)
                GT -> {- trace "GT" $ -} return (Left (pos+1))
    case beforeOrAfter i (bounds ca) of
        GT -> {- trace (s ++ "GT") $ -} f pos max (\pos' r' -> lookupArrayHelper i r (pos'+r') pos max) (checkAndReturn (pos + 1))
        LT -> {- trace (s ++ "LT") $ -} f min pos (\pos' r' -> lookupArrayHelper i r pos'      min pos) (checkAndReturn pos)
        EQ -> {- trace (s ++ "EQ") $ -} return (Right (pos,ca))

insertWhereItGoes :: RangeArray ra m e ref => InnerArray e -> RangeMap ra e ref -> m (Either (OuterIndex, InnerArray e) ())
insertWhereItGoes e r = do
    let (i,j) = bounds e
    outerArray <- readRefArray r
    (outerZ,outerN) <- MA.getBounds outerArray
    let inside x (low,hi) = x >= low && x <= hi -- flip inRange
    lr <- lookupArray i r
    case lr of
        Left ii ->
            trace "(insertWhereItGoes) Left ii case"
                  $if ii <= outerN -- is there a next array?
                    then do
                        y <- MA.readArray outerArray (ii)
                        let boundsy = show (bounds y)
                            insideStr = show j ++ " inside " ++ boundsy
                            outsideStr = show j ++ " outside " ++ boundsy
                        if j `inside` bounds y -- does the final bound of inserted array overlap next array?
                            then trace insideStr  $ return (Left (ii,y)) -- TODO, actually could shrink y or e
                            else trace outsideStr $ Right <$> insertArrayAt r ii e
                    else Right <$> insertArrayAt r ii e
        Right x ->
            trace "(insertWhereItGoes) Right x case"
                  $return (Left x)

lookupInRangeMap :: RangeArray ra m e ref => Index -> RangeMap ra e ref -> m (Maybe e)
lookupInRangeMap k r = do
    b <- lookupArray k r
    case b of
        Right (_,ar) -> return (Just (ar ! k))
        Left _ -> return Nothing


deleteRange :: RangeArray ra m e ref => (Index,Index)
                                     -> RangeMap ra e ref
                                     -> m (Either OuterIndex (OuterIndex,InnerArray e))

deleteRange (z,n) r = do
    result <- lookupArray z r
    case result of
        Left _ -> return result
        Right (outerIndex,a) | bounds a /= (z,n) -> return (Left outerIndex)
        Right (outerIndex,a) | bounds a == (z,n) -> do
           outerArray <- readRefArray r
           (zz,nn) <- getBounds outerArray
           as <-   map (\(i,e) -> (if i > outerIndex then i - 1 else i, e))
               .   filter ((/=outerIndex) . fst)
               <$> getAssocs outerArray
           mutAr <- MA.newListArray (zz,nn-1) (map snd as)
           writeRefArray r mutAr
           return result