{-# LANGUAGE TupleSections #-}
module Bio.Chain
    ( ChainLike (..)
    , Chain
    , chain, fromList
    , (!), (//)
    ) where

import           Control.Lens
import           Data.Array      (Array, Ix, array, listArray, (!), (//))
import qualified Data.Array      as A (assocs, bounds)
import           Data.Array.Base (unsafeAt)
import qualified Data.Vector     as V

type Chain i a = Array i a

-- | Construct new chain from list
--
chain :: Ix i => (i, i) -> [(i, a)] -> Chain i a
chain :: forall i a. Ix i => (i, i) -> [(i, a)] -> Chain i a
chain = forall i a. Ix i => (i, i) -> [(i, a)] -> Chain i a
array

-- | Construct new int-labeled chain from list
--
fromList :: [a] -> Chain Int a
fromList :: forall a. [a] -> Chain Int a
fromList [a]
lst = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
lst forall a. Num a => a -> a -> a
- Int
1) [a]
lst

-- | Chain-like sequence, by default it is an array or a list
--
class (Ixed m, Enum (Index m)) => ChainLike m where
    bounds       :: m -> (Index m, Index m)
    assocs       :: m -> [(Index m, IxValue m)]
    modify       :: Index m -> (IxValue m -> IxValue m) -> m -> m
    modifyBefore :: Index m -> (IxValue m -> IxValue m) -> m -> m
    modifyAfter  :: Index m -> (IxValue m -> IxValue m) -> m -> m

    unsafeRead   :: m -> Index m -> IxValue m
    unsafeRead m
ch Index m
i = m
ch forall s a. HasCallStack => s -> Getting (Endo a) s a -> a
^?! forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix Index m
i

instance ChainLike [a] where
    bounds :: [a] -> (Index [a], Index [a])
bounds = (Int
0,) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Enum a => a -> a
pred forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length

    assocs :: [a] -> [(Index [a], IxValue [a])]
assocs  = forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..]

    modify :: Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
modify       Index [a]
_ IxValue [a] -> IxValue [a]
_ []     = []
    modify       Index [a]
0 IxValue [a] -> IxValue [a]
f (a
x:[a]
xs) = IxValue [a] -> IxValue [a]
f a
xforall a. a -> [a] -> [a]
:[a]
xs
    modify       Index [a]
i IxValue [a] -> IxValue [a]
f (a
x:[a]
xs) = a
xforall a. a -> [a] -> [a]
:forall m.
ChainLike m =>
Index m -> (IxValue m -> IxValue m) -> m -> m
modify (Index [a]
i forall a. Num a => a -> a -> a
- Int
1) IxValue [a] -> IxValue [a]
f [a]
xs

    modifyBefore :: Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
modifyBefore Index [a]
i IxValue [a] -> IxValue [a]
f [a]
lst = (IxValue [a] -> IxValue [a]
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Int -> [a] -> [a]
take Index [a]
i [a]
lst) forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
drop Index [a]
i [a]
lst
    modifyAfter :: Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
modifyAfter  Index [a]
i IxValue [a] -> IxValue [a]
f [a]
lst = forall a. Int -> [a] -> [a]
take (Index [a]
i forall a. Num a => a -> a -> a
+ Int
1) [a]
lst forall a. [a] -> [a] -> [a]
++ (IxValue [a] -> IxValue [a]
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Int -> [a] -> [a]
drop (Index [a]
i forall a. Num a => a -> a -> a
+ Int
1) [a]
lst)

    unsafeRead :: [a] -> Index [a] -> IxValue [a]
unsafeRead = forall a. [a] -> Int -> a
(!!)

instance (Ix i, Enum i) => ChainLike (Array i a) where
    bounds :: Array i a -> (Index (Array i a), Index (Array i a))
bounds = forall i e. Array i e -> (i, i)
A.bounds

    assocs :: Array i a -> [(Index (Array i a), IxValue (Array i a))]
assocs = forall i e. Ix i => Array i e -> [(i, e)]
A.assocs

    modify :: Index (Array i a)
-> (IxValue (Array i a) -> IxValue (Array i a))
-> Array i a
-> Array i a
modify       Index (Array i a)
i IxValue (Array i a) -> IxValue (Array i a)
f Array i a
ar = Array i a
ar forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(Index (Array i a)
i, IxValue (Array i a) -> IxValue (Array i a)
f (Array i a
ar forall i e. Ix i => Array i e -> i -> e
! Index (Array i a)
i))]

    modifyBefore :: Index (Array i a)
-> (IxValue (Array i a) -> IxValue (Array i a))
-> Array i a
-> Array i a
modifyBefore Index (Array i a)
i IxValue (Array i a) -> IxValue (Array i a)
f Array i a
ar = let (Index (Array i a)
mi, Index (Array i a)
_) = forall m. ChainLike m => m -> (Index m, Index m)
bounds Array i a
ar
                          in Array i a
ar forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
j, IxValue (Array i a) -> IxValue (Array i a)
f (Array i a
ar forall i e. Ix i => Array i e -> i -> e
! i
j)) | i
j <- [Index (Array i a)
mi .. forall a. Enum a => a -> a
pred Index (Array i a)
i]]
    modifyAfter :: Index (Array i a)
-> (IxValue (Array i a) -> IxValue (Array i a))
-> Array i a
-> Array i a
modifyAfter  Index (Array i a)
i IxValue (Array i a) -> IxValue (Array i a)
f Array i a
ar = let (Index (Array i a)
_, Index (Array i a)
ma) = forall m. ChainLike m => m -> (Index m, Index m)
bounds Array i a
ar
                          in Array i a
ar forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
j, IxValue (Array i a) -> IxValue (Array i a)
f (Array i a
ar forall i e. Ix i => Array i e -> i -> e
! i
j)) | i
j <- [forall a. Enum a => a -> a
succ Index (Array i a)
i .. Index (Array i a)
ma]]

    {-# INLINE unsafeRead #-}
    unsafeRead :: Array i a -> Index (Array i a) -> IxValue (Array i a)
unsafeRead = forall m. UnsafeReadArray m => m -> Index m -> IxValue m
unsafeReadArray

instance ChainLike (V.Vector a) where
    bounds :: Vector a -> (Index (Vector a), Index (Vector a))
bounds Vector a
v = (Int
0, forall a. Vector a -> Int
V.length Vector a
v forall a. Num a => a -> a -> a
- Int
1)

    assocs :: Vector a -> [(Index (Vector a), IxValue (Vector a))]
assocs = forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Vector a -> [a]
V.toList

    modify :: Index (Vector a)
-> (IxValue (Vector a) -> IxValue (Vector a))
-> Vector a
-> Vector a
modify Index (Vector a)
i IxValue (Vector a) -> IxValue (Vector a)
f Vector a
ar = Vector a
ar forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Index (Vector a)
i, IxValue (Vector a) -> IxValue (Vector a)
f (Vector a
ar forall a. Vector a -> Int -> a
V.! Index (Vector a)
i))]

    modifyBefore :: Index (Vector a)
-> (IxValue (Vector a) -> IxValue (Vector a))
-> Vector a
-> Vector a
modifyBefore Index (Vector a)
i IxValue (Vector a) -> IxValue (Vector a)
f Vector a
ar = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap IxValue (Vector a) -> IxValue (Vector a)
f Vector a
before forall a. Semigroup a => a -> a -> a
<> Vector a
after
      where
        (Vector a
before, Vector a
after) = forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt Index (Vector a)
i Vector a
ar

    modifyAfter :: Index (Vector a)
-> (IxValue (Vector a) -> IxValue (Vector a))
-> Vector a
-> Vector a
modifyAfter Index (Vector a)
i IxValue (Vector a) -> IxValue (Vector a)
f Vector a
ar = Vector a
before forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap IxValue (Vector a) -> IxValue (Vector a)
f Vector a
after
      where
        (Vector a
before, Vector a
after) = forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt (Index (Vector a)
i forall a. Num a => a -> a -> a
+ Int
1) Vector a
ar

    {-# INLINE unsafeRead #-}
    unsafeRead :: Vector a -> Index (Vector a) -> IxValue (Vector a)
unsafeRead = forall a. Vector a -> Int -> a
V.unsafeIndex 


class (Ixed m) => UnsafeReadArray m where
    unsafeReadArray :: m -> Index m -> IxValue m

instance (Ix i, Enum i) => UnsafeReadArray (Array i a) where
    {-# INLINE unsafeReadArray #-}
    unsafeReadArray :: Array i a -> Index (Array i a) -> IxValue (Array i a)
unsafeReadArray = forall i e. Ix i => Array i e -> i -> e
(!)

instance {-# OVERLAPPING #-} UnsafeReadArray (Array Int a) where
    {-# INLINE unsafeReadArray #-}
    unsafeReadArray :: Array Int a -> Index (Array Int a) -> IxValue (Array Int a)
unsafeReadArray = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
unsafeAt