{-# LANGUAGE CPP #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}
module Bio.MMTF
  ( module Bio.MMTF.Type
  , decode
  , fetch
  ) where

import           Bio.MMTF.Decode        (l2v)
import           Bio.MMTF.MessagePack   ()
import           Bio.MMTF.Type
import           Bio.Structure

import           Control.Monad.IO.Class (MonadIO)
import           Data.Bifunctor         (Bifunctor (..))
import           Data.ByteString.Lazy   (ByteString)
import           Data.Int               (Int32)
import           Data.List              (mapAccumL, zip4)
import           Data.MessagePack       (unpack)
import           Data.String            (IsString (..))
import           Data.Text              (Text)
import           Data.Vector            (Vector, empty, toList, (!))
import           Linear.V3              (V3 (..))
import           Network.HTTP.Simple    (getResponseBody, httpLBS)
#if !MIN_VERSION_base(4,13,0)
import           Control.Monad.Fail     (MonadFail (..))
import           Prelude                hiding (fail)
#endif

-- | Decodes a 'ByteString' to 'MMTF'
--
decode :: MonadFail m => ByteString -> m MMTF
decode :: forall (m :: * -> *). MonadFail m => ByteString -> m MMTF
decode = forall (m :: * -> *) a.
(Applicative m, Monad m, MonadFail m, MessagePack a) =>
ByteString -> m a
unpack

-- | Fetches MMTF structure from RSCB
fetch :: (MonadFail m, MonadIO m) => String -> m MMTF
fetch :: forall (m :: * -> *). (MonadFail m, MonadIO m) => String -> m MMTF
fetch String
pdbid = do let url :: Request
url = forall a. IsString a => String -> a
fromString forall a b. (a -> b) -> a -> b
$ String
"https://mmtf.rcsb.org/v1.0/full/" forall a. Semigroup a => a -> a -> a
<> String
pdbid
                 Response ByteString
resp <- forall (m :: * -> *).
MonadIO m =>
Request -> m (Response ByteString)
httpLBS Request
url
                 forall (m :: * -> *). MonadFail m => ByteString -> m MMTF
decode (forall a. Response a -> a
getResponseBody Response ByteString
resp)

instance StructureModels MMTF where
    -- TODO: add global bonds
    modelsOf :: MMTF -> Vector Model
modelsOf MMTF
m = forall a. [a] -> Vector a
l2v (forall a b c. (a -> b -> c) -> b -> a -> c
flip Vector Chain -> Vector (Bond GlobalID) -> Model
Model forall a. Vector a
empty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> Vector a
l2v forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Text -> Vector Residue -> Chain
Chain) [[Text]]
chainNames [[Vector Residue]]
chainResis)
      where
        chainsCnts :: [Int]
chainsCnts = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Vector a -> [a]
toList (ModelData -> Vector Int32
chainsPerModel (MMTF -> ModelData
model MMTF
m))
        groupsCnts :: [Int]
groupsCnts = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Vector a -> [a]
toList (ChainData -> Vector Int32
groupsPerChain (MMTF -> ChainData
chain MMTF
m))
        groupsRaws :: [[(GroupType, SecondaryStructure, [Atom])]]
groupsRaws = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (Int, Int)
-> Int -> ((Int, Int), [(GroupType, SecondaryStructure, [Atom])])
getGroups (Int
0, Int
0) [Int]
groupsCnts
        groups :: [[[(GroupType, SecondaryStructure, [Atom])]]]
groups     = forall a. [Int] -> [a] -> [[a]]
cutter [Int]
chainsCnts [[(GroupType, SecondaryStructure, [Atom])]]
groupsRaws
        chainNames :: [[Text]]
chainNames = forall a. [Int] -> [a] -> [[a]]
cutter [Int]
chainsCnts (forall a. Vector a -> [a]
toList forall a b. (a -> b) -> a -> b
$ ChainData -> Vector Text
chainNameList forall a b. (a -> b) -> a -> b
$ MMTF -> ChainData
chain MMTF
m)
        chainResis :: [[Vector Residue]]
chainResis = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. [a] -> Vector a
l2v forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (GroupType, SecondaryStructure, [Atom]) -> Residue
mkResidue)) [[[(GroupType, SecondaryStructure, [Atom])]]]
groups

        getGroups :: (Int, Int) -> Int -> ((Int, Int), [(GroupType, SecondaryStructure, [Atom])])
        getGroups :: (Int, Int)
-> Int -> ((Int, Int), [(GroupType, SecondaryStructure, [Atom])])
getGroups (Int
chOffset, Int
atOffset) Int
sz = let chEnd :: Int
chEnd        = Int
chOffset forall a. Num a => a -> a -> a
+ Int
sz
                                                gtl :: Vector Int32
gtl          = GroupData -> Vector Int32
groupTypeList (MMTF -> GroupData
group MMTF
m)
                                                gl :: Vector GroupType
gl           = GroupData -> Vector GroupType
groupList (MMTF -> GroupData
group MMTF
m)
                                                ssl :: Vector SecondaryStructure
ssl          = GroupData -> Vector SecondaryStructure
secStructList (MMTF -> GroupData
group MMTF
m)
                                                chr :: [Int]
chr          = [Int
chOffset .. Int
chEnd forall a. Num a => a -> a -> a
- Int
1]
                                                rgt :: [GroupType]
rgt          = (Vector GroupType
gl forall a. Vector a -> Int -> a
!) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Int32
gtl forall a. Vector a -> Int -> a
!) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
chr
                                                rss :: [SecondaryStructure]
rss          = (Vector SecondaryStructure
ssl forall a. Vector a -> Int -> a
!) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
chr
                                                (Int
atEnd, [[Atom]]
ats) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL Int -> GroupType -> (Int, [Atom])
getAtoms Int
atOffset [GroupType]
rgt
                                            in  ((Int
chEnd, Int
atEnd), forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [GroupType]
rgt [SecondaryStructure]
rss [[Atom]]
ats)

        getAtoms :: Int -> GroupType -> (Int, [Atom])
        getAtoms :: Int -> GroupType -> (Int, [Atom])
getAtoms Int
offset GroupType
gt = let cl :: [Int]
cl  = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Vector a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupType -> Vector Int32
gtFormalChargeList forall a b. (a -> b) -> a -> b
$ GroupType
gt
                                 nl :: [Text]
nl  = forall a. Vector a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupType -> Vector Text
gtAtomNameList forall a b. (a -> b) -> a -> b
$ GroupType
gt
                                 el :: [Text]
el  = forall a. Vector a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupType -> Vector Text
gtElementList forall a b. (a -> b) -> a -> b
$ GroupType
gt
                                 ics :: [Int]
ics = [Int
offset .. Int
end forall a. Num a => a -> a -> a
- Int
1]
                                 end :: Int
end = Int
offset forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
cl
                             in  (Int
end, (Int, Text, Text, Int) -> Atom
mkAtom forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 [Int]
cl [Text]
nl [Text]
el [Int]
ics)

        mkResidue :: (GroupType, SecondaryStructure, [Atom]) -> Residue
        -- TODO: support residue number here
        mkResidue :: (GroupType, SecondaryStructure, [Atom]) -> Residue
mkResidue (GroupType
gt, SecondaryStructure
ss, [Atom]
atoms') = Text
-> Int
-> Char
-> Vector Atom
-> Vector (Bond LocalID)
-> SecondaryStructure
-> Text
-> Residue
Residue (GroupType -> Text
gtGroupName GroupType
gt) (-Int
1) Char
' ' (forall a. [a] -> Vector a
l2v [Atom]
atoms')
                                             (Vector (Int32, Int32) -> Vector Int32 -> Vector (Bond LocalID)
mkBonds (GroupType -> Vector (Int32, Int32)
gtBondAtomList GroupType
gt) (GroupType -> Vector Int32
gtBondOrderList GroupType
gt))
                                              SecondaryStructure
ss (GroupType -> Text
gtChemCompType GroupType
gt)

        mkBonds :: Vector (Int32, Int32) -> Vector Int32 -> Vector (Bond LocalID)
        mkBonds :: Vector (Int32, Int32) -> Vector Int32 -> Vector (Bond LocalID)
mkBonds Vector (Int32, Int32)
bal Vector Int32
bol = let ball :: [(LocalID, LocalID)]
ball = forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Int -> LocalID
LocalID forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Int -> LocalID
LocalID forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Vector a -> [a]
toList Vector (Int32, Int32)
bal
                              boll :: [Int]
boll = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Vector a -> [a]
toList Vector Int32
bol
                              res :: [Bond LocalID]
res  = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(LocalID
f, LocalID
t) Int
o -> forall m. m -> m -> Int -> Bond m
Bond LocalID
f LocalID
t Int
o) [(LocalID, LocalID)]
ball [Int]
boll
                          in  forall a. [a] -> Vector a
l2v [Bond LocalID]
res

        mkAtom :: (Int, Text, Text, Int) -> Atom
        mkAtom :: (Int, Text, Text, Int) -> Atom
mkAtom (Int
fc, Text
n, Text
e, Int
idx) = let i :: Vector Int32
i = AtomData -> Vector Int32
atomIdList (MMTF -> AtomData
atom MMTF
m)
                                     x :: Vector Float
x = AtomData -> Vector Float
xCoordList (MMTF -> AtomData
atom MMTF
m)
                                     y :: Vector Float
y = AtomData -> Vector Float
yCoordList (MMTF -> AtomData
atom MMTF
m)
                                     z :: Vector Float
z = AtomData -> Vector Float
zCoordList (MMTF -> AtomData
atom MMTF
m)
                                     o :: Vector Float
o = AtomData -> Vector Float
occupancyList (MMTF -> AtomData
atom MMTF
m)
                                     b :: Vector Float
b = AtomData -> Vector Float
bFactorList (MMTF -> AtomData
atom MMTF
m)
                                 in  GlobalID
-> Int -> Text -> Text -> V3 Float -> Int -> Float -> Float -> Atom
Atom (Int -> GlobalID
GlobalID Int
idx)
                                          (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Vector Int32
i forall a. Vector a -> Int -> a
! Int
idx)
                                          Text
n
                                          Text
e
                                          (forall a. a -> a -> a -> V3 a
V3 (Vector Float
x forall a. Vector a -> Int -> a
! Int
idx) (Vector Float
y forall a. Vector a -> Int -> a
! Int
idx) (Vector Float
z forall a. Vector a -> Int -> a
! Int
idx))
                                          Int
fc
                                          (Vector Float
b forall a. Vector a -> Int -> a
! Int
idx)
                                          (Vector Float
o forall a. Vector a -> Int -> a
! Int
idx)

        cutter :: [Int] -> [a] -> [[a]]
        cutter :: forall a. [Int] -> [a] -> [[a]]
cutter []     []    = []
        cutter (Int
x:[Int]
xs) [a]
ys    = forall a. Int -> [a] -> [a]
take Int
x [a]
ys forall a. a -> [a] -> [a]
: forall a. [Int] -> [a] -> [[a]]
cutter [Int]
xs (forall a. Int -> [a] -> [a]
drop Int
x [a]
ys)
        cutter []     (a
_:[a]
_) = forall a. HasCallStack => String -> a
error String
"Cutter: you cannot be here"