module Bio.Utils.Overlap
    ( overlapFragment
    , overlapNucl
    , coverage
    ) where

import           Bio.Data.Bed
import           Conduit
import           Lens.Micro                ((^.))
import           Control.Monad
import qualified Data.ByteString.Char8       as B
import           Data.Function
import qualified Data.HashMap.Strict         as M
import qualified Data.IntervalMap.Strict     as IM
import           Data.List
import qualified Data.Vector.Unboxed         as V
import qualified Data.Vector.Unboxed.Mutable as VM

-- | convert lines of a BED file into a data structure - A hashmap of which the
-- | chromosomes, and values are interval maps.
toMap :: [(B.ByteString, (Int, Int))] -> M.HashMap B.ByteString (IM.IntervalMap Int Int)
toMap :: [(ByteString, (Int, Int))]
-> HashMap ByteString (IntervalMap Int Int)
toMap [(ByteString, (Int, Int))]
input = [(ByteString, IntervalMap Int Int)]
-> HashMap ByteString (IntervalMap Int Int)
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList([(ByteString, IntervalMap Int Int)]
 -> HashMap ByteString (IntervalMap Int Int))
-> ([((ByteString, (Int, Int)), Int)]
    -> [(ByteString, IntervalMap Int Int)])
-> [((ByteString, (Int, Int)), Int)]
-> HashMap ByteString (IntervalMap Int Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.([((ByteString, (Int, Int)), Int)]
 -> (ByteString, IntervalMap Int Int))
-> [[((ByteString, (Int, Int)), Int)]]
-> [(ByteString, IntervalMap Int Int)]
forall a b. (a -> b) -> [a] -> [b]
map [((ByteString, (Int, Int)), Int)]
-> (ByteString, IntervalMap Int Int)
forall e a v.
Ord e =>
[((a, (e, e)), v)] -> (a, IntervalMap (Interval e) v)
create([[((ByteString, (Int, Int)), Int)]]
 -> [(ByteString, IntervalMap Int Int)])
-> ([((ByteString, (Int, Int)), Int)]
    -> [[((ByteString, (Int, Int)), Int)]])
-> [((ByteString, (Int, Int)), Int)]
-> [(ByteString, IntervalMap Int Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(((ByteString, (Int, Int)), Int)
 -> ((ByteString, (Int, Int)), Int) -> Bool)
-> [((ByteString, (Int, Int)), Int)]
-> [[((ByteString, (Int, Int)), Int)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
(==) (ByteString -> ByteString -> Bool)
-> (((ByteString, (Int, Int)), Int) -> ByteString)
-> ((ByteString, (Int, Int)), Int)
-> ((ByteString, (Int, Int)), Int)
-> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` ((ByteString, (Int, Int)) -> ByteString
forall a b. (a, b) -> a
fst((ByteString, (Int, Int)) -> ByteString)
-> (((ByteString, (Int, Int)), Int) -> (ByteString, (Int, Int)))
-> ((ByteString, (Int, Int)), Int)
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
.((ByteString, (Int, Int)), Int) -> (ByteString, (Int, Int))
forall a b. (a, b) -> a
fst)) ([((ByteString, (Int, Int)), Int)]
 -> HashMap ByteString (IntervalMap Int Int))
-> [((ByteString, (Int, Int)), Int)]
-> HashMap ByteString (IntervalMap Int Int)
forall a b. (a -> b) -> a -> b
$ [(ByteString, (Int, Int))]
-> [Int] -> [((ByteString, (Int, Int)), Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(ByteString, (Int, Int))]
input [Int
0..]
    where
        f :: ((a, (a, a)), b) -> (Interval a, b)
f ((a
_, (a, a)
x), b
i) = ((a, a) -> Interval a
forall a. (a, a) -> Interval a
toInterval (a, a)
x, b
i)
        create :: [((a, (e, e)), v)] -> (a, IntervalMap (Interval e) v)
create [((a, (e, e)), v)]
xs = ((a, (e, e)) -> a
forall a b. (a, b) -> a
fst((a, (e, e)) -> a)
-> ([((a, (e, e)), v)] -> (a, (e, e))) -> [((a, (e, e)), v)] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.((a, (e, e)), v) -> (a, (e, e))
forall a b. (a, b) -> a
fst(((a, (e, e)), v) -> (a, (e, e)))
-> ([((a, (e, e)), v)] -> ((a, (e, e)), v))
-> [((a, (e, e)), v)]
-> (a, (e, e))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.[((a, (e, e)), v)] -> ((a, (e, e)), v)
forall a. [a] -> a
head ([((a, (e, e)), v)] -> a) -> [((a, (e, e)), v)] -> a
forall a b. (a -> b) -> a -> b
$ [((a, (e, e)), v)]
xs, [(Interval e, v)] -> IntervalMap (Interval e) v
forall k e v. Interval k e => [(k, v)] -> IntervalMap k v
IM.fromDistinctAscList([(Interval e, v)] -> IntervalMap (Interval e) v)
-> ([((a, (e, e)), v)] -> [(Interval e, v)])
-> [((a, (e, e)), v)]
-> IntervalMap (Interval e) v
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(((a, (e, e)), v) -> (Interval e, v))
-> [((a, (e, e)), v)] -> [(Interval e, v)]
forall a b. (a -> b) -> [a] -> [b]
map ((a, (e, e)), v) -> (Interval e, v)
forall a a b. ((a, (a, a)), b) -> (Interval a, b)
f ([((a, (e, e)), v)] -> IntervalMap (Interval e) v)
-> [((a, (e, e)), v)] -> IntervalMap (Interval e) v
forall a b. (a -> b) -> a -> b
$ [((a, (e, e)), v)]
xs)
{-# INLINE toMap #-}

coverage :: [BED]  -- ^ genomic locus in BED format
         -> ConduitT () BED IO ()  -- ^ reads in BED format
         -> IO (V.Vector Double, Int)
coverage :: [BED] -> ConduitT () BED IO () -> IO (Vector Double, Int)
coverage [BED]
bin ConduitT () BED IO ()
tags = (Vector Int -> (Vector Double, Int))
-> IO (Vector Int) -> IO (Vector Double, Int)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM Vector Int -> (Vector Double, Int)
forall b c.
(Unbox b, Unbox c, Fractional c, Integral b) =>
Vector b -> (Vector c, b)
getResult (IO (Vector Int) -> IO (Vector Double, Int))
-> IO (Vector Int) -> IO (Vector Double, Int)
forall a b. (a -> b) -> a -> b
$ ConduitT () Void IO (Vector Int) -> IO (Vector Int)
forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
runConduit (ConduitT () Void IO (Vector Int) -> IO (Vector Int))
-> ConduitT () Void IO (Vector Int) -> IO (Vector Int)
forall a b. (a -> b) -> a -> b
$ ConduitT () BED IO ()
tags ConduitT () BED IO ()
-> ConduitM BED Void IO (Vector Int)
-> ConduitT () Void IO (Vector Int)
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM BED Void IO (Vector Int)
forall o. ConduitT BED o IO (Vector Int)
sink
  where
    sink :: ConduitT BED o IO (Vector Int)
sink = do
        MVector (PrimState IO) Int
v <- IO (MVector (PrimState IO) Int)
-> ConduitT BED o IO (MVector (PrimState IO) Int)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (MVector (PrimState IO) Int)
 -> ConduitT BED o IO (MVector (PrimState IO) Int))
-> IO (MVector (PrimState IO) Int)
-> ConduitT BED o IO (MVector (PrimState IO) Int)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
0
        (BED -> IO ()) -> ConduitT BED o IO ()
forall (m :: * -> *) a o.
Monad m =>
(a -> m ()) -> ConduitT a o m ()
mapM_C ((BED -> IO ()) -> ConduitT BED o IO ())
-> (BED -> IO ()) -> ConduitT BED o IO ()
forall a b. (a -> b) -> a -> b
$ \BED
t -> do
                let set :: Maybe (IntervalMap Int Int)
set = ByteString
-> HashMap ByteString (IntervalMap Int Int)
-> Maybe (IntervalMap Int Int)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup (BED
tBED -> Getting ByteString BED ByteString -> ByteString
forall s a. s -> Getting a s a -> a
^.Getting ByteString BED ByteString
forall b. BEDLike b => Lens' b ByteString
chrom) HashMap ByteString (IntervalMap Int Int)
featMap
                    s :: Int
s = BED
tBED -> Getting Int BED Int -> Int
forall s a. s -> Getting a s a -> a
^.Getting Int BED Int
forall b. BEDLike b => Lens' b Int
chromStart
                    e :: Int
e = BED
tBED -> Getting Int BED Int -> Int
forall s a. s -> Getting a s a -> a
^.Getting Int BED Int
forall b. BEDLike b => Lens' b Int
chromEnd
                    b :: (Int, Int)
b = (Int
s, Int
e)
                    l :: Int
l = Int
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                    intervals :: [(Interval Int, Int)]
intervals = case Maybe (IntervalMap Int Int)
set of
                        Just IntervalMap Int Int
iMap -> IntervalMap Int Int -> [(Interval Int, Int)]
forall k v. IntervalMap k v -> [(k, v)]
IM.toList (IntervalMap Int Int -> [(Interval Int, Int)])
-> ((Int, Int) -> IntervalMap Int Int)
-> (Int, Int)
-> [(Interval Int, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntervalMap Int Int -> Interval Int -> IntervalMap Int Int
forall k e v.
Interval k e =>
IntervalMap k v -> k -> IntervalMap k v
IM.intersecting IntervalMap Int Int
iMap (Interval Int -> IntervalMap Int Int)
-> ((Int, Int) -> Interval Int)
-> (Int, Int)
-> IntervalMap Int Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Interval Int
forall a. (a, a) -> Interval a
toInterval ((Int, Int) -> [(Interval Int, Int)])
-> (Int, Int) -> [(Interval Int, Int)]
forall a b. (a -> b) -> a -> b
$ (Int, Int)
b
                        Maybe (IntervalMap Int Int)
_ -> []
                [(Interval Int, Int)] -> ((Interval Int, Int) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Interval Int, Int)]
intervals (\(Interval Int, Int)
interval -> do
                    let i :: Int
i = (Interval Int, Int) -> Int
forall a b. (a, b) -> b
snd (Interval Int, Int)
interval
                        nucl :: Int
nucl = (Int, Int) -> Interval Int -> Int
forall a. (Ord a, Num a) => (a, a) -> Interval a -> a
overlap (Int, Int)
b (Interval Int -> Int)
-> ((Interval Int, Int) -> Interval Int)
-> (Interval Int, Int)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Interval Int, Int) -> Interval Int
forall a b. (a, b) -> a
fst ((Interval Int, Int) -> Int) -> (Interval Int, Int) -> Int
forall a b. (a -> b) -> a -> b
$ (Interval Int, Int)
interval
                    MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector (PrimState IO) Int
v Int
i (Int -> IO ()) -> (Int -> Int) -> Int -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
nucl) (Int -> IO ()) -> IO Int -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector (PrimState IO) Int
v Int
i
                    )
                MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector (PrimState IO) Int
v Int
n (Int -> IO ()) -> (Int -> Int) -> Int -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
l) (Int -> IO ()) -> IO Int -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector (PrimState IO) Int
v Int
n
        IO (Vector Int) -> ConduitT BED o IO (Vector Int)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (Vector Int) -> ConduitT BED o IO (Vector Int))
-> IO (Vector Int) -> ConduitT BED o IO (Vector Int)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
V.freeze MVector (PrimState IO) Int
v
    getResult :: Vector b -> (Vector c, b)
getResult Vector b
v = ((b -> Int -> c) -> Vector b -> Vector Int -> Vector c
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith b -> Int -> c
forall a a a. (Fractional a, Integral a, Integral a) => a -> a -> a
normalize (Int -> Int -> Vector b -> Vector b
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
V.slice Int
0 Int
n Vector b
v) Vector Int
featWidth, Vector b
v Vector b -> Int -> b
forall a. Unbox a => Vector a -> Int -> a
V.! Int
n)
    featMap :: HashMap ByteString (IntervalMap Int Int)
featMap = [(ByteString, (Int, Int))]
-> HashMap ByteString (IntervalMap Int Int)
toMap([(ByteString, (Int, Int))]
 -> HashMap ByteString (IntervalMap Int Int))
-> ([BED] -> [(ByteString, (Int, Int))])
-> [BED]
-> HashMap ByteString (IntervalMap Int Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(BED -> (ByteString, (Int, Int)))
-> [BED] -> [(ByteString, (Int, Int))]
forall a b. (a -> b) -> [a] -> [b]
map (\BED
x -> (BED
xBED -> Getting ByteString BED ByteString -> ByteString
forall s a. s -> Getting a s a -> a
^.Getting ByteString BED ByteString
forall b. BEDLike b => Lens' b ByteString
chrom, (BED
xBED -> Getting Int BED Int -> Int
forall s a. s -> Getting a s a -> a
^.Getting Int BED Int
forall b. BEDLike b => Lens' b Int
chromStart, BED
xBED -> Getting Int BED Int -> Int
forall s a. s -> Getting a s a -> a
^.Getting Int BED Int
forall b. BEDLike b => Lens' b Int
chromEnd))) ([BED] -> HashMap ByteString (IntervalMap Int Int))
-> [BED] -> HashMap ByteString (IntervalMap Int Int)
forall a b. (a -> b) -> a -> b
$ [BED]
bin
    featWidth :: Vector Int
featWidth = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
V.fromList ([Int] -> Vector Int) -> [Int] -> Vector Int
forall a b. (a -> b) -> a -> b
$ (BED -> Int) -> [BED] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map BED -> Int
forall b. BEDLike b => b -> Int
size [BED]
bin
    n :: Int
n = [BED] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [BED]
bin
    overlap :: (a, a) -> Interval a -> a
overlap (a
l, a
u) (IM.ClosedInterval a
l' a
u')
        | a
l' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
l = if a
u' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
u then a
u'a -> a -> a
forall a. Num a => a -> a -> a
-a
l'a -> a -> a
forall a. Num a => a -> a -> a
+a
1 else a
ua -> a -> a
forall a. Num a => a -> a -> a
-a
l'a -> a -> a
forall a. Num a => a -> a -> a
+a
1
        | Bool
otherwise = if a
u' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
u then a
u'a -> a -> a
forall a. Num a => a -> a -> a
-a
la -> a -> a
forall a. Num a => a -> a -> a
+a
1 else a
ua -> a -> a
forall a. Num a => a -> a -> a
-a
la -> a -> a
forall a. Num a => a -> a -> a
+a
1
    overlap (a, a)
_ Interval a
_ = a
0
    normalize :: a -> a -> a
normalize a
a a
b = a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
b

overlapFragment, overlapNucl ::
                  [(Int, Int)] -- ^ Ascending order list
                -> [(Int, Int)] -- ^ tags in any order
                -> V.Vector Int
overlapFragment :: [(Int, Int)] -> [(Int, Int)] -> Vector Int
overlapFragment [(Int, Int)]
xs [(Int, Int)]
ts = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
V.create (Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate Int
n Int
0 ST s (MVector s Int)
-> (MVector s Int -> ST s (MVector s Int)) -> ST s (MVector s Int)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [(Int, Int)]
-> MVector (PrimState (ST s)) Int
-> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) (t :: * -> *) a.
(Foldable t, PrimMonad m, Unbox a, Num a) =>
t (Int, Int)
-> MVector (PrimState m) a -> m (MVector (PrimState m) a)
go [(Int, Int)]
ts)
    where
        n :: Int
n = [(Int, Int)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int)]
xs
        iMap :: IntervalMap Int Int
iMap = [(Interval Int, Int)] -> IntervalMap Int Int
forall k e v. (Interval k e, Eq k) => [(k, v)] -> IntervalMap k v
IM.fromAscList ([(Interval Int, Int)] -> IntervalMap Int Int)
-> [(Interval Int, Int)] -> IntervalMap Int Int
forall a b. (a -> b) -> a -> b
$ [Interval Int] -> [Int] -> [(Interval Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Int, Int) -> Interval Int) -> [(Int, Int)] -> [Interval Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Interval Int
forall a. (a, a) -> Interval a
toInterval [(Int, Int)]
xs) [Int
0..]
        go :: t (Int, Int)
-> MVector (PrimState m) a -> m (MVector (PrimState m) a)
go t (Int, Int)
ts' MVector (PrimState m) a
v = do
            t (Int, Int) -> ((Int, Int) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t (Int, Int)
ts' (\(Int, Int)
x -> do
                let indices :: [Int]
indices = IntervalMap Int Int -> [Int]
forall k v. IntervalMap k v -> [v]
IM.elems (IntervalMap Int Int -> [Int])
-> ((Int, Int) -> IntervalMap Int Int) -> (Int, Int) -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntervalMap Int Int -> Interval Int -> IntervalMap Int Int
forall k e v.
Interval k e =>
IntervalMap k v -> k -> IntervalMap k v
IM.intersecting IntervalMap Int Int
iMap (Interval Int -> IntervalMap Int Int)
-> ((Int, Int) -> Interval Int)
-> (Int, Int)
-> IntervalMap Int Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Interval Int
forall a. (a, a) -> Interval a
toInterval ((Int, Int) -> [Int]) -> (Int, Int) -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int, Int)
x
                [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int]
indices (\Int
i -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector (PrimState m) a
v Int
i (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a
forall a. Num a => a -> a -> a
+a
1) (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector (PrimState m) a
v Int
i)
                )
            MVector (PrimState m) a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a. Monad m => a -> m a
return MVector (PrimState m) a
v

overlapNucl :: [(Int, Int)] -> [(Int, Int)] -> Vector Int
overlapNucl [(Int, Int)]
xs [(Int, Int)]
ts = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
V.create (Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate Int
n Int
0 ST s (MVector s Int)
-> (MVector s Int -> ST s (MVector s Int)) -> ST s (MVector s Int)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [(Int, Int)]
-> MVector (PrimState (ST s)) Int
-> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) (t :: * -> *).
(Foldable t, PrimMonad m) =>
t (Int, Int)
-> MVector (PrimState m) Int -> m (MVector (PrimState m) Int)
go [(Int, Int)]
ts)
    where
        n :: Int
n = [(Int, Int)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int)]
xs
        iMap :: IntervalMap Int Int
iMap = [(Interval Int, Int)] -> IntervalMap Int Int
forall k e v. (Interval k e, Eq k) => [(k, v)] -> IntervalMap k v
IM.fromAscList ([(Interval Int, Int)] -> IntervalMap Int Int)
-> [(Interval Int, Int)] -> IntervalMap Int Int
forall a b. (a -> b) -> a -> b
$ [Interval Int] -> [Int] -> [(Interval Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Int, Int) -> Interval Int) -> [(Int, Int)] -> [Interval Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Interval Int
forall a. (a, a) -> Interval a
toInterval [(Int, Int)]
xs) [Int
0..]
        go :: t (Int, Int)
-> MVector (PrimState m) Int -> m (MVector (PrimState m) Int)
go t (Int, Int)
ts' MVector (PrimState m) Int
v = do
            t (Int, Int) -> ((Int, Int) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t (Int, Int)
ts' (\(Int, Int)
x -> do
                let intervals :: [(Interval Int, Int)]
intervals = IntervalMap Int Int -> [(Interval Int, Int)]
forall k v. IntervalMap k v -> [(k, v)]
IM.toList (IntervalMap Int Int -> [(Interval Int, Int)])
-> ((Int, Int) -> IntervalMap Int Int)
-> (Int, Int)
-> [(Interval Int, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntervalMap Int Int -> Interval Int -> IntervalMap Int Int
forall k e v.
Interval k e =>
IntervalMap k v -> k -> IntervalMap k v
IM.intersecting IntervalMap Int Int
iMap (Interval Int -> IntervalMap Int Int)
-> ((Int, Int) -> Interval Int)
-> (Int, Int)
-> IntervalMap Int Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Interval Int
forall a. (a, a) -> Interval a
toInterval ((Int, Int) -> [(Interval Int, Int)])
-> (Int, Int) -> [(Interval Int, Int)]
forall a b. (a -> b) -> a -> b
$ (Int, Int)
x
                [(Interval Int, Int)] -> ((Interval Int, Int) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Interval Int, Int)]
intervals (\(Interval Int, Int)
interval -> do
                    let i :: Int
i = (Interval Int, Int) -> Int
forall a b. (a, b) -> b
snd (Interval Int, Int)
interval
                        nucl :: Int
nucl = (Int, Int) -> Interval Int -> Int
forall a. (Ord a, Num a) => (a, a) -> Interval a -> a
overlap (Int, Int)
x (Interval Int -> Int)
-> ((Interval Int, Int) -> Interval Int)
-> (Interval Int, Int)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Interval Int, Int) -> Interval Int
forall a b. (a, b) -> a
fst ((Interval Int, Int) -> Int) -> (Interval Int, Int) -> Int
forall a b. (a -> b) -> a -> b
$ (Interval Int, Int)
interval
                    MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector (PrimState m) Int
v Int
i (Int -> m ()) -> (Int -> Int) -> Int -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
nucl) (Int -> m ()) -> m Int -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector (PrimState m) Int
v Int
i
                    )
                )
            MVector (PrimState m) Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a. Monad m => a -> m a
return MVector (PrimState m) Int
v
        overlap :: (a, a) -> Interval a -> a
overlap (a
l, a
u) (IM.ClosedInterval a
l' a
u')
            | a
l' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
l = if a
u' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
u then a
u'a -> a -> a
forall a. Num a => a -> a -> a
-a
l'a -> a -> a
forall a. Num a => a -> a -> a
+a
1 else a
ua -> a -> a
forall a. Num a => a -> a -> a
-a
l'a -> a -> a
forall a. Num a => a -> a -> a
+a
1
            | Bool
otherwise = if a
u' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
u then a
u'a -> a -> a
forall a. Num a => a -> a -> a
-a
la -> a -> a
forall a. Num a => a -> a -> a
+a
1 else a
ua -> a -> a
forall a. Num a => a -> a -> a
-a
la -> a -> a
forall a. Num a => a -> a -> a
+a
1
        overlap (a, a)
_ Interval a
_ = a
0

toInterval :: (a, a) -> IM.Interval a
toInterval :: (a, a) -> Interval a
toInterval (a
l, a
u) = a -> a -> Interval a
forall a. a -> a -> Interval a
IM.ClosedInterval a
l a
u
{-# INLINE toInterval #-}