{-# LANGUAGE BangPatterns #-}
module Bio.GO.GREAT
    ( AssocRule(..)
    , getRegulatoryDomains
    , enrichedTerms
    ) where

import           Conduit
import           Control.Monad.Primitive
import qualified Data.ByteString.Char8        as B
import           Data.Default.Class
import           Data.Function                (on)
import qualified Data.HashMap.Strict          as M
import qualified Data.IntervalMap             as IM
import           Data.List                    (foldl', group, sort, sortBy)
import           Data.Maybe
import           Data.Ord                     (comparing)
import qualified Data.Vector                  as V
import qualified Data.Vector.Algorithms.Intro as I

import           Bio.Data.Bed
import           Bio.GO
import           Bio.Utils.Functions

-- | how to associate genomic regions with genes
data AssocRule = BasalPlusExtension Int Int Int
               | TwoNearest
               | OneNearest

instance Default AssocRule where
    def = BasalPlusExtension 5000 1000 1000000

type Gene a = ( ( B.ByteString  -- ^ chromosome
              , Int           -- ^ tss
              , Bool          -- ^ is forward stranded
              ), a)

-- | given a gene list and the rule, compute the rulatory domain for each gene
getRegulatoryDomains :: AssocRule -> [Gene a] -> [(BED3, a)]
getRegulatoryDomains (BasalPlusExtension up dw ext) genes = (extendTail r ext, a) : rs
  where
    (rs, Just (r,a)) = foldl' f ([], Nothing) $ sortBy (compareBed `on` fst) basal
    f (acc, Nothing) (b,x) = (acc, Just (extendHead b ext, x))
    f (acc, Just (b', x')) (b,x)
        | chrom b' /= chrom b = ( (extendTail b' ext, x') : acc
                                , Just (extendHead b ext, x) )
        | chromEnd b' >= chromStart b = ((b',x') : acc, Just (b,x))
        | otherwise = let ext' = min ext $ (chromStart b - chromEnd b') `div` 2
                      in ((extendTail b' ext', x') : acc, Just (extendHead b ext', x))
    extendHead (BED3 chr s e) l | s - l >= 0 = BED3 chr (s-l) e
                                | otherwise = BED3 chr 0 e
    extendTail (BED3 chr s e) l = BED3 chr s (e+l)
    basal = flip map genes $ \((chr, tss, str), x) ->
        if str then (BED3 chr (tss - up) (tss + dw), x)
               else (BED3 chr (tss - dw) (tss + up), x)
getRegulatoryDomains _ _ = undefined
{-# INLINE getRegulatoryDomains #-}

-- | how many times a particular GO term is hit by given regions
countTerms :: (BEDLike b, Monad m)
           => BEDTree [GOId]
           -> Sink b m (Int, M.HashMap GOId Int)
countTerms tree = go 0 M.empty
  where
    go !n !termCount = do
        x <- await
        case x of
            Just bed ->do
                let chr = chrom bed
                    s = chromStart bed
                    e = chromEnd bed
                    terms = nub' . concat . IM.elems . IM.intersecting
                        (M.lookupDefault IM.empty chr tree) $ IM.IntervalCO s e
                go (n+1) $ foldl (\acc t -> M.insertWith (+) t 1 acc) termCount terms
            _ -> return (n, termCount)
{-
    expand = nub' . foldl f []
      where
        f acc i = traceBack i $ i : acc
        traceBack i acc = case getParentById i goMap of
            Just g -> traceBack (_oboId g) $ (_oboId g) : acc
            _ -> acc
-}
{-# INLINE countTerms #-}

nub' :: [B.ByteString] -> [B.ByteString]
nub' = map head . group . sort
{-# INLINE nub' #-}

{-
-- | since GO terms are organized as a tree, a node is considered as being hit
-- if any of its children is hit. During this step, we traverse GO tree to include
-- parents if the regions that hit parents are different from those hiting their
-- children
expandTermCounts :: M.HashMap GOId Int -> GOMap -> M.HashMap GOId Int
expandTermCounts counts goMap = M.fromList . prune . expand $ counts
  where
    prune m = loop goTree
      where
        loop (Node p children : xs) =
            let children' = filter (\(Node (_,c) _) -> c /= 0) children
                l = length children'
            in case () of
               _ | l == 0 -> p : loop xs
                 | snd p == (snd . rootLabel . head) children' -> loop $ children' ++ xs
                 | otherwise -> p : (loop $ children' ++ xs)
        loop _ = []

        goTree :: [Tree (GOId, Int)]
        goTree = map (fmap (\x -> (_oboId x, M.lookupDefault 0 (_oboId x) m))) . buildGOTree $ goMap
    expand m = M.foldlWithKey' f M.empty m
      where
        f acc i c = traceBack c i $ M.insertWith (+) i c acc
    traceBack c i m = case getParentById i goMap of
        Just g -> traceBack c (_oboId g) $ M.insertWith (+) (_oboId g) c m
        _ -> m
{-# INLINE expandTermCounts #-}
-}

enrichedTerms :: (BEDLike b1, BEDLike b2, PrimMonad m)
              => Source m b1
              -> Source m b2
              -> BEDTree [GOId]
              -> m (V.Vector (GOId, (Double, Double)))
enrichedTerms fg bg tree = do
    (n1, table1) <- fg $$ countTerms tree
    (n2, table2) <- bg $$ countTerms tree
    v <- V.unsafeThaw $ V.fromList $ M.toList $ flip M.mapWithKey table1 $ \t c ->
        let k = fromMaybe (error "x") $ M.lookup t table2
        in (1 - hyperquick c k n1 n2, fromIntegral (c*n2) / fromIntegral (n1*k))
    I.sortBy (comparing snd) v
    V.unsafeFreeze v