{-# LANGUAGE OverloadedStrings #-}
module Bio.GO
    ( GO(..)
    , GOId
    , GOMap
    , getGOLevel
    ) where

import qualified Data.HashMap.Strict as M
import Data.Maybe
import qualified Data.Text           as T

data GO = GO
    { GO -> GOId
_oboId        :: !GOId
    , GO -> Text
_label        :: !T.Text
    , GO -> [GOId]
_subProcessOf :: ![GOId]
    , GO -> Text
_oboNS        :: !T.Text
    } deriving (GOId -> GO -> ShowS
[GO] -> ShowS
GO -> String
(GOId -> GO -> ShowS)
-> (GO -> String) -> ([GO] -> ShowS) -> Show GO
forall a.
(GOId -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GO] -> ShowS
$cshowList :: [GO] -> ShowS
show :: GO -> String
$cshow :: GO -> String
showsPrec :: GOId -> GO -> ShowS
$cshowsPrec :: GOId -> GO -> ShowS
Show, ReadPrec [GO]
ReadPrec GO
GOId -> ReadS GO
ReadS [GO]
(GOId -> ReadS GO)
-> ReadS [GO] -> ReadPrec GO -> ReadPrec [GO] -> Read GO
forall a.
(GOId -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [GO]
$creadListPrec :: ReadPrec [GO]
readPrec :: ReadPrec GO
$creadPrec :: ReadPrec GO
readList :: ReadS [GO]
$creadList :: ReadS [GO]
readsPrec :: GOId -> ReadS GO
$creadsPrec :: GOId -> ReadS GO
Read)

type GOId = Int

type GOMap = M.HashMap GOId GO

type TermCount = M.HashMap GOId Int

-- | The top level is 0.
getGOLevel :: GOId -> GOMap -> Int
getGOLevel :: GOId -> GOMap -> GOId
getGOLevel GOId
gid GOMap
gm = GOId -> [GOId] -> GOId
forall t. Num t => t -> [GOId] -> t
loop GOId
0 [GOId
gid]
  where
    loop :: t -> [GOId] -> t
loop t
l [GOId]
ids | [GOId] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [GOId]
parents = t
l
               | Bool
otherwise = t -> [GOId] -> t
loop (t
lt -> t -> t
forall a. Num a => a -> a -> a
+t
1) [GOId]
parents
      where
        parents :: [GOId]
parents = (GO -> [GOId]) -> [GO] -> [GOId]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap GO -> [GOId]
_subProcessOf ([GO] -> [GOId]) -> [GO] -> [GOId]
forall a b. (a -> b) -> a -> b
$ ((GOId -> Maybe GO) -> [GOId] -> [GO])
-> [GOId] -> (GOId -> Maybe GO) -> [GO]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (GOId -> Maybe GO) -> [GOId] -> [GO]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [GOId]
ids ((GOId -> Maybe GO) -> [GO]) -> (GOId -> Maybe GO) -> [GO]
forall a b. (a -> b) -> a -> b
$ \GOId
i -> GOId -> GOMap -> Maybe GO
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup GOId
i GOMap
gm 

{-
getParentById :: GOId -> GOMap -> Maybe GO
getParentById gid goMap = M.lookup gid goMap >>= _subProcessOf
                                             >>= (`M.lookup` goMap)
{-# INLINE getParentById #-}

-- | Add a GO term to the count table. Term counts will propogate from child to
-- its parents. This function works for cyclical graph as well.
addTerm :: GO -> GOMap -> TermCount -> TermCount
addTerm g m t = loop S.empty g t
  where
    loop visited go table
        | _oboId go `S.member` visited = table
        | otherwise = case _subProcessOf go of
            Nothing -> table'
            Just gid -> loop (S.insert (_oboId go) visited)
                (M.lookupDefault undefined gid m) table'
      where
        table' = M.insertWith (+) (_oboId go) 1 table

enrichment :: (TermCount, Int)  -- ^ Background frequency and the total number
           -> (TermCount, Int)  -- ^ Foreground
           -> [(GOId, Double, Double)]
enrichment (bg, bg_total) (fg, fg_total) =
    flip map (M.toList fg) $ \(gid, fg_count) ->
        let enrich = fromIntegral (fg_count * bg_total) /
                     fromIntegral (fg_total * bg_count)
            bg_count = M.lookupDefault undefined gid bg
            p = 1 - hyperquick fg_count bg_count fg_total bg_total
        in (gid, enrich, p)
-}