{-# 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.Unboxed where

import Data.Word
import Data.Array.Unboxed
import qualified Data.Array.MArray as MA
import Data.Array.MArray (MArray(..))
import qualified Data.Array.Base as Base
import Data.Reference
import Debug.Trace
import Control.Concurrent.STM
import Control.Concurrent.STM.TArray
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 a 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