module Util.DynArray (DynArray,new,newMatrix,newLinearArray,read
                     ,write,writes
                     ,resizeTo,updateIO,update,bounds,index,indexBy
                     ,elements,showDynArray
                     ,deleteInLinearArray
                     ,deleteColumnsInMatrix,deleteRowsInMatrix
                     ,insertEmptyIntoLinearArrayBefore
                     ,insertEmptyRowIntoMatrixBefore
                     ,insertEmptyColumnIntoMatrixBefore)
    where

import Prelude hiding (init,read)
import Control.Applicative ((<$>))
import Control.Monad (forM_,when,foldM)
import Control.Exception (assert)
import Data.Array.IO (IOArray,readArray,writeArray,getBounds,getElems)
import qualified Data.Array.IO as ArrayIO
import Data.Ix (Ix,range,inRange)
import Data.IORef
import Data.List (sort)

data DynArray i e = DynArray { ref :: IORef (IOArray i e)
                             , grow :: (i,i) -> i -> (i,i)
                             , init :: e
                             }

defaultGrowingSize = 10

showDynArray :: (Ix i,Show e) => DynArray i e -> IO String
showDynArray dynArray = show <$> (readIORef (ref dynArray) >>= getElems)

new :: Ix i => (i,i) -> ((i,i) -> i -> (i,i)) -> e -> IO (DynArray i e)
new ix grow e = do
  array <- ArrayIO.newArray ix e
  ref <- newIORef array
  return $ DynArray ref grow e

newMatrix :: (Int,Int) -> e -> IO (DynArray (Int,Int) e)
newMatrix (initRows,initColumns) = 
    let grow (_,(r,c)) (i,j) = ( (0,0)
                               , ( max r $ i + defaultGrowingSize - 1
                                 , max c $ j + defaultGrowingSize - 1))
    in
      new ((0,0),(initRows-1,initColumns-1)) grow

newLinearArray :: Int -> e -> IO (DynArray Int e)
newLinearArray numElements =
    let grow _ i = (0,i + defaultGrowingSize - 1)
    in
      new (0,numElements-1) grow

read :: Ix i => i -> DynArray i e -> IO e
read i dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  if inRange bounds i then readArray array i
                      else return $ init dynArray

growArray :: Ix i => i -> DynArray i e -> IO ()
growArray i dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  array' <- ArrayIO.newArray (grow dynArray bounds i) $ init dynArray
  copyUnsafe array array'
  writeIORef (ref dynArray) array'

write :: Ix i => i -> e -> DynArray i e -> IO ()
write i e dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  if inRange bounds i 
   then writeArray array i e
   else growArray i dynArray >> write i e dynArray

writes :: (Ix i,Enum i) => i -> [e] -> DynArray i e -> IO ()
writes i es dynArray = 
    forM_ (zip [i..] es) $ \(pos,e) -> write pos e dynArray

copyUnsafe :: Ix i => IOArray i e -> IOArray i e -> IO ()
copyUnsafe from to = do
  r <- getBounds from
  forM_ (range r) $ \i -> readArray from i >>= writeArray to i

resizeTo :: Ix i => (i,i) -> DynArray i e -> IO ()
resizeTo range dynArray = 
    ArrayIO.newArray range (init dynArray) >>= writeIORef (ref dynArray)

updateIO :: Ix i => i -> (e -> IO e) -> DynArray i e -> IO ()
updateIO i f dynArray = 
    read i dynArray >>= f >>= \e -> write i e dynArray

update ::  Ix i => i -> (e -> e) -> DynArray i e -> IO ()
update i f = updateIO i (return . f)
  
bounds :: Ix i => DynArray i e -> IO (i,i)
bounds dynArray = readIORef (ref dynArray) >>= getBounds

indexBy :: (Enum i, Ix i) => (e -> Bool) -> DynArray i e -> IO (Maybe i)
indexBy equals dynArray = do
  array <- readIORef $ ref dynArray
  bounds <- getBounds array
  let indexRec i = 
          if inRange bounds i 
          then do value <- readArray array i
                  if equals value then return $ Just i
                                  else indexRec $ succ i
          else return Nothing
  indexRec $ fst bounds

index :: (Enum i, Ix i, Eq e) => e -> DynArray i e -> IO (Maybe i)
index e = indexBy ((==) e)

elements :: Ix i => DynArray i e -> IO [e]
elements dynArray = readIORef (ref dynArray) >>= getElems

insertEmptyIntoLinearArrayBefore :: Int -> Int -> DynArray Int e -> IO ()
insertEmptyIntoLinearArrayBefore i maxCopy dynArray = do
  maxBound <- snd <$> bounds dynArray
  when (maxCopy >= maxBound) $ growArray maxCopy dynArray
  array <- readIORef $ ref dynArray
  forM_ [maxCopy,maxCopy-1..i] $ \n -> 
      readArray array n >>= writeArray array (n+1)
  writeArray array i $ init dynArray

insertEmptyRowIntoMatrixBefore :: Int -> Int -> DynArray (Int,Int) e -> IO ()
insertEmptyRowIntoMatrixBefore i maxCopyRow dynArray = do
  (maxRow,_) <- snd <$> bounds dynArray
  when (maxCopyRow >= maxRow) $ growArray (maxCopyRow,0) dynArray

  forM_ [maxCopyRow,maxCopyRow-1..i] $ \n -> 
      copyRowUnsafe n (n+1) dynArray
  writeEmptyRow i dynArray

insertEmptyColumnIntoMatrixBefore :: Int -> Int -> DynArray (Int,Int) e -> IO ()
insertEmptyColumnIntoMatrixBefore i maxCopyCol dynArray = do
  (_,maxCol) <- snd <$> bounds dynArray
  when (maxCopyCol >= maxCol) $ growArray (0,maxCopyCol) dynArray

  forM_ [maxCopyCol,maxCopyCol-1..i] $ \n -> 
      copyColumnUnsafe n (n+1) dynArray
  writeEmptyColumn i dynArray

deleteInLinearArray :: [Int] -> DynArray Int e -> IO ()
deleteInLinearArray is dynArray = do
  assert (is == sort is) $ return ()
  array <- readIORef $ ref dynArray
  (minB,maxB) <- getBounds array

  let deleteElement (x:xs,deleted) index | x == index = 
          return (xs,deleted+1)

      deleteElement (xs,deleted) index = do
        when (deleted > 0) $ 
          readArray array index >>= writeArray array (index - deleted)
        return (xs,deleted)

  (_,deleted) <- foldM deleteElement (is,0) [minB..maxB] 
  forM_ [maxB-deleted+1..maxB] $ \x -> writeArray array x $ init dynArray

deleteColumnsInMatrix :: [Int] -> DynArray (Int,Int) e -> IO ()
deleteColumnsInMatrix columns dynArray = do
  assert (columns == sort columns) $ return ()
  array <- readIORef $ ref dynArray
  ((_,minCol),(_,maxCol)) <- getBounds array
                                       
  let deleteColumn (col:cols,deleted) index | col == index = 
          return (cols,deleted+1)

      deleteColumn (cols,deleted) index = do
        when (deleted > 0) $ 
             copyColumnUnsafe index (index - deleted) dynArray
        return (cols,deleted)

  (_,deleted) <- foldM deleteColumn (columns,0) [minCol..maxCol] 
  forM_ [maxCol-deleted+1..maxCol] $ writeEmptyColumn `flip` dynArray

deleteRowsInMatrix :: [Int] -> DynArray (Int,Int) e -> IO ()
deleteRowsInMatrix rowsToDelete dynArray = do
  assert (rowsToDelete == sort rowsToDelete) $ return ()
  array <- readIORef $ ref dynArray
  ((minRow,_),(maxRow,_)) <- getBounds array
                                       
  let deleteRow (row:rows,deleted) index | row == index = 
          return (rows,deleted+1)

      deleteRow (rows,deleted) index = do
        when (deleted > 0) $ 
             copyRowUnsafe index (index - deleted) dynArray
        return (rows,deleted)

  (_,deleted) <- foldM deleteRow (rowsToDelete,0) [minRow..maxRow] 
  forM_ [maxRow-deleted+1..maxRow] $ writeEmptyRow `flip` dynArray

copyColumnUnsafe :: Int -> Int -> DynArray (Int,Int) e -> IO ()
copyColumnUnsafe from to dynArray = do
  array <- readIORef $ ref dynArray
  ((minRow,_),(maxRow,_)) <- getBounds array
  forM_ [minRow..maxRow] $ \r -> 
      readArray array (r,from) >>= writeArray array (r,to)

copyRowUnsafe :: Int -> Int -> DynArray (Int,Int) e -> IO ()
copyRowUnsafe from to dynArray = do
  array <- readIORef $ ref dynArray
  ((_,minCol),(_,maxCol)) <- getBounds array
  forM_ [minCol..maxCol] $ \c -> 
      readArray array (from,c) >>= writeArray array (to,c)

writeEmptyColumn :: Int -> DynArray (Int,Int) e -> IO ()
writeEmptyColumn i dynArray = do
  array <- readIORef $ ref dynArray
  ((minRow,_),(maxRow,_)) <- getBounds array
  forM_ [minRow..maxRow] $ \r -> 
      writeArray array (r,i) $ init dynArray

writeEmptyRow :: Int -> DynArray (Int,Int) e -> IO ()
writeEmptyRow i dynArray = do
  array <- readIORef $ ref dynArray
  ((_,minCol),(_,maxCol)) <- getBounds array
  forM_ [minCol..maxCol] $ \c -> 
      writeArray array (i,c) $ init dynArray