{-# 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 :: ByteString -> m MMTF
decode = ByteString -> m MMTF
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 :: String -> m MMTF
fetch String
pdbid = do let url :: Request
url = String -> Request
forall a. IsString a => String -> a
fromString (String -> Request) -> String -> Request
forall a b. (a -> b) -> a -> b
$ String
"https://mmtf.rcsb.org/v1.0/full/" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
pdbid
                 Response ByteString
resp <- Request -> m (Response ByteString)
forall (m :: * -> *).
MonadIO m =>
Request -> m (Response ByteString)
httpLBS Request
url
                 ByteString -> m MMTF
forall (m :: * -> *). MonadFail m => ByteString -> m MMTF
decode (Response ByteString -> ByteString
forall a. Response a -> a
getResponseBody Response ByteString
resp)

instance StructureModels MMTF where
    -- TODO: add global bonds
    modelsOf :: MMTF -> Vector Model
modelsOf MMTF
m = [Model] -> Vector Model
forall a. [a] -> Vector a
l2v ((Vector Chain -> Vector (Bond GlobalID) -> Model)
-> Vector (Bond GlobalID) -> Vector Chain -> Model
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vector Chain -> Vector (Bond GlobalID) -> Model
Model Vector (Bond GlobalID)
forall a. Vector a
empty (Vector Chain -> Model)
-> ([Chain] -> Vector Chain) -> [Chain] -> Model
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Chain] -> Vector Chain
forall a. [a] -> Vector a
l2v ([Chain] -> Model) -> [[Chain]] -> [Model]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Text] -> [Vector Residue] -> [Chain])
-> [[Text]] -> [[Vector Residue]] -> [[Chain]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Text -> Vector Residue -> Chain)
-> [Text] -> [Vector Residue] -> [Chain]
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 = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> [Int32] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector Int32 -> [Int32]
forall a. Vector a -> [a]
toList (ModelData -> Vector Int32
chainsPerModel (MMTF -> ModelData
model MMTF
m))
        groupsCnts :: [Int]
groupsCnts = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> [Int32] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector Int32 -> [Int32]
forall a. Vector a -> [a]
toList (ChainData -> Vector Int32
groupsPerChain (MMTF -> ChainData
chain MMTF
m))
        groupsRaws :: [[(GroupType, SecondaryStructure, [Atom])]]
groupsRaws = ((Int, Int), [[(GroupType, SecondaryStructure, [Atom])]])
-> [[(GroupType, SecondaryStructure, [Atom])]]
forall a b. (a, b) -> b
snd (((Int, Int), [[(GroupType, SecondaryStructure, [Atom])]])
 -> [[(GroupType, SecondaryStructure, [Atom])]])
-> ((Int, Int), [[(GroupType, SecondaryStructure, [Atom])]])
-> [[(GroupType, SecondaryStructure, [Atom])]]
forall a b. (a -> b) -> a -> b
$ ((Int, Int)
 -> Int -> ((Int, Int), [(GroupType, SecondaryStructure, [Atom])]))
-> (Int, Int)
-> [Int]
-> ((Int, Int), [[(GroupType, SecondaryStructure, [Atom])]])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL (Int, Int)
-> Int -> ((Int, Int), [(GroupType, SecondaryStructure, [Atom])])
getGroups (Int
0, Int
0) [Int]
groupsCnts
        groups :: [[[(GroupType, SecondaryStructure, [Atom])]]]
groups     = [Int]
-> [[(GroupType, SecondaryStructure, [Atom])]]
-> [[[(GroupType, SecondaryStructure, [Atom])]]]
forall a. [Int] -> [a] -> [[a]]
cutter [Int]
chainsCnts [[(GroupType, SecondaryStructure, [Atom])]]
groupsRaws
        chainNames :: [[Text]]
chainNames = [Int] -> [Text] -> [[Text]]
forall a. [Int] -> [a] -> [[a]]
cutter [Int]
chainsCnts (Vector Text -> [Text]
forall a. Vector a -> [a]
toList (Vector Text -> [Text]) -> Vector Text -> [Text]
forall a b. (a -> b) -> a -> b
$ ChainData -> Vector Text
chainNameList (ChainData -> Vector Text) -> ChainData -> Vector Text
forall a b. (a -> b) -> a -> b
$ MMTF -> ChainData
chain MMTF
m)
        chainResis :: [[Vector Residue]]
chainResis = ([[(GroupType, SecondaryStructure, [Atom])]] -> [Vector Residue])
-> [[[(GroupType, SecondaryStructure, [Atom])]]]
-> [[Vector Residue]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([(GroupType, SecondaryStructure, [Atom])] -> Vector Residue)
-> [[(GroupType, SecondaryStructure, [Atom])]] -> [Vector Residue]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Residue] -> Vector Residue
forall a. [a] -> Vector a
l2v ([Residue] -> Vector Residue)
-> ([(GroupType, SecondaryStructure, [Atom])] -> [Residue])
-> [(GroupType, SecondaryStructure, [Atom])]
-> Vector Residue
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((GroupType, SecondaryStructure, [Atom]) -> Residue)
-> [(GroupType, SecondaryStructure, [Atom])] -> [Residue]
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 Int -> Int -> Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                                                rgt :: [GroupType]
rgt          = (Vector GroupType
gl Vector GroupType -> Int -> GroupType
forall a. Vector a -> Int -> a
!) (Int -> GroupType) -> (Int -> Int) -> Int -> GroupType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> (Int -> Int32) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Int32
gtl Vector Int32 -> Int -> Int32
forall a. Vector a -> Int -> a
!) (Int -> GroupType) -> [Int] -> [GroupType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
chr
                                                rss :: [SecondaryStructure]
rss          = (Vector SecondaryStructure
ssl Vector SecondaryStructure -> Int -> SecondaryStructure
forall a. Vector a -> Int -> a
!) (Int -> SecondaryStructure) -> [Int] -> [SecondaryStructure]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
chr
                                                (Int
atEnd, [[Atom]]
ats) = (Int -> GroupType -> (Int, [Atom]))
-> Int -> [GroupType] -> (Int, [[Atom]])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL Int -> GroupType -> (Int, [Atom])
getAtoms Int
atOffset [GroupType]
rgt
                                            in  ((Int
chEnd, Int
atEnd), [GroupType]
-> [SecondaryStructure]
-> [[Atom]]
-> [(GroupType, SecondaryStructure, [Atom])]
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  = (Int32 -> Int) -> [Int32] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int32] -> [Int]) -> (GroupType -> [Int32]) -> GroupType -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int32 -> [Int32]
forall a. Vector a -> [a]
toList (Vector Int32 -> [Int32])
-> (GroupType -> Vector Int32) -> GroupType -> [Int32]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupType -> Vector Int32
gtFormalChargeList (GroupType -> [Int]) -> GroupType -> [Int]
forall a b. (a -> b) -> a -> b
$ GroupType
gt
                                 nl :: [Text]
nl  = Vector Text -> [Text]
forall a. Vector a -> [a]
toList (Vector Text -> [Text])
-> (GroupType -> Vector Text) -> GroupType -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupType -> Vector Text
gtAtomNameList (GroupType -> [Text]) -> GroupType -> [Text]
forall a b. (a -> b) -> a -> b
$ GroupType
gt
                                 el :: [Text]
el  = Vector Text -> [Text]
forall a. Vector a -> [a]
toList (Vector Text -> [Text])
-> (GroupType -> Vector Text) -> GroupType -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GroupType -> Vector Text
gtElementList (GroupType -> [Text]) -> GroupType -> [Text]
forall a b. (a -> b) -> a -> b
$ GroupType
gt
                                 ics :: [Int]
ics = [Int
offset .. Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                                 end :: Int
end = Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
cl
                             in  (Int
end, (Int, Text, Text, Int) -> Atom
mkAtom ((Int, Text, Text, Int) -> Atom)
-> [(Int, Text, Text, Int)] -> [Atom]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> [Text] -> [Text] -> [Int] -> [(Int, Text, Text, Int)]
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
' ' ([Atom] -> Vector Atom
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 = (Int32 -> LocalID)
-> (Int32 -> LocalID) -> (Int32, Int32) -> (LocalID, LocalID)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Int -> LocalID
LocalID (Int -> LocalID) -> (Int32 -> Int) -> Int32 -> LocalID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Int -> LocalID
LocalID (Int -> LocalID) -> (Int32 -> Int) -> Int32 -> LocalID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral) ((Int32, Int32) -> (LocalID, LocalID))
-> [(Int32, Int32)] -> [(LocalID, LocalID)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector (Int32, Int32) -> [(Int32, Int32)]
forall a. Vector a -> [a]
toList Vector (Int32, Int32)
bal
                              boll :: [Int]
boll = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> [Int32] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector Int32 -> [Int32]
forall a. Vector a -> [a]
toList Vector Int32
bol
                              res :: [Bond LocalID]
res  = ((LocalID, LocalID) -> Int -> Bond LocalID)
-> [(LocalID, LocalID)] -> [Int] -> [Bond LocalID]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(LocalID
f, LocalID
t) Int
o -> LocalID -> LocalID -> Int -> Bond LocalID
forall m. m -> m -> Int -> Bond m
Bond LocalID
f LocalID
t Int
o) [(LocalID, LocalID)]
ball [Int]
boll
                          in  [Bond LocalID] -> Vector (Bond LocalID)
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)
                                          (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32 -> Int) -> Int32 -> Int
forall a b. (a -> b) -> a -> b
$ Vector Int32
i Vector Int32 -> Int -> Int32
forall a. Vector a -> Int -> a
! Int
idx)
                                          Text
n
                                          Text
e
                                          (Float -> Float -> Float -> V3 Float
forall a. a -> a -> a -> V3 a
V3 (Vector Float
x Vector Float -> Int -> Float
forall a. Vector a -> Int -> a
! Int
idx) (Vector Float
y Vector Float -> Int -> Float
forall a. Vector a -> Int -> a
! Int
idx) (Vector Float
z Vector Float -> Int -> Float
forall a. Vector a -> Int -> a
! Int
idx))
                                          Int
fc
                                          (Vector Float
b Vector Float -> Int -> Float
forall a. Vector a -> Int -> a
! Int
idx)
                                          (Vector Float
o Vector Float -> Int -> Float
forall a. Vector a -> Int -> a
! Int
idx)

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