{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Bio.PDB.Writer
  ( pdbToFile
  , pdbToText
  ) where


import           Bio.PDB.Type           (Atom (..), Chain, Model, PDB (..))
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Data.Text              (Text)
import qualified Data.Text              as T (cons, drop, init, intercalate,
                                              last, length, pack, replicate,
                                              singleton, splitAt, take)
import qualified Data.Text.IO           as TIO (writeFile)
import           Data.Vector            (Vector)
import qualified Data.Vector            as V (all, foldl', fromList, last,
                                              length, toList, zip)
import           Text.Printf            (printf)

-- | Writes 'PDB' to the given path.
--
pdbToFile :: MonadIO m => PDB -> FilePath -> m ()
pdbToFile :: PDB -> FilePath -> m ()
pdbToFile PDB
pdb = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (FilePath -> IO ()) -> FilePath -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath -> Text -> IO ()) -> Text -> FilePath -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip FilePath -> Text -> IO ()
TIO.writeFile (PDB -> Text
pdbToText PDB
pdb)

-- | Converts 'PDB' to 'Text'.
--
pdbToText :: PDB -> Text
pdbToText :: PDB -> Text
pdbToText PDB{Text
Map RemarkCode RemarkData
Map FieldType RemarkData
Vector Model
otherFields :: PDB -> Map FieldType RemarkData
remarks :: PDB -> Map RemarkCode RemarkData
models :: PDB -> Vector Model
title :: PDB -> Text
otherFields :: Map FieldType RemarkData
remarks :: Map RemarkCode RemarkData
models :: Vector Model
title :: Text
..} = (Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
newLine Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
toPDBLine Text
end)
                  (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Text -> [Text] -> Text
T.intercalate Text
newLine ([Text] -> Text)
-> (Vector (Model, Int) -> [Text]) -> Vector (Model, Int) -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RemarkData -> [Text]
forall a. Vector a -> [a]
V.toList (RemarkData -> [Text])
-> (Vector (Model, Int) -> RemarkData)
-> Vector (Model, Int)
-> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Model, Int) -> Text) -> Vector (Model, Int) -> RemarkData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> (Model, Int) -> Text
modelToText Bool
separateModels)
                  (Vector (Model, Int) -> Text) -> Vector (Model, Int) -> Text
forall a b. (a -> b) -> a -> b
$ Vector Model -> Vector Int -> Vector (Model, Int)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector Model
models (Vector Int -> Vector (Model, Int))
-> Vector Int -> Vector (Model, Int)
forall a b. (a -> b) -> a -> b
$ [Int] -> Vector Int
forall a. [a] -> Vector a
V.fromList [Int
1 ..]
  where
    separateModels :: Bool
separateModels = Vector Model -> Int
forall a. Vector a -> Int
V.length Vector Model
models Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1

    end :: Text
    end :: Text
end = Text
"END   "

type TerAtom = Atom

modelToText :: Bool -> (Model, Int) -> Text
modelToText :: Bool -> (Model, Int) -> Text
modelToText Bool
separateModels (Model
pdbModel, Int
modelInd) = Text
modelPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
atomsT Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
modelSuffix
  where
    atomsT :: Text
atomsT = Text -> [Text] -> Text
T.intercalate Text
newLine ([Text] -> Text) -> (Model -> [Text]) -> Model -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RemarkData -> [Text]
forall a. Vector a -> [a]
V.toList (RemarkData -> [Text]) -> (Model -> RemarkData) -> Model -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Either Atom Atom -> Text)
-> Vector (Either Atom Atom) -> RemarkData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Either Atom Atom -> Text
atomOrTer (Vector (Either Atom Atom) -> RemarkData)
-> (Model -> Vector (Either Atom Atom)) -> Model -> RemarkData
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Model -> Vector (Either Atom Atom)
withTers (Model -> Text) -> Model -> Text
forall a b. (a -> b) -> a -> b
$ Model
pdbModel

    modelPrefix :: Text
modelPrefix | Bool
separateModels = Text -> Text
toPDBLine (Text
mdl Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText Int
4 Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Int -> Text
forall a. Show a => Int -> a -> Text
prependToS Int
4 Int
modelInd) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
newLine
                | Bool
otherwise      = Text
""

    modelSuffix :: Text
modelSuffix | Bool
separateModels = Text
newLine Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
toPDBLine Text
endmdl
                | Bool
otherwise      = Text
""

    mdl :: Text
    mdl :: Text
mdl = Text
"MODEL "

    endmdl :: Text
    endmdl :: Text
endmdl = Text
"ENDMDL "

    withTers :: Vector Chain -> Vector (Either Atom TerAtom)
    withTers :: Model -> Vector (Either Atom Atom)
withTers = (Int, Vector (Either Atom Atom)) -> Vector (Either Atom Atom)
forall a b. (a, b) -> b
snd ((Int, Vector (Either Atom Atom)) -> Vector (Either Atom Atom))
-> (Model -> (Int, Vector (Either Atom Atom)))
-> Model
-> Vector (Either Atom Atom)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Vector (Either Atom Atom))
 -> Chain -> (Int, Vector (Either Atom Atom)))
-> (Int, Vector (Either Atom Atom))
-> Model
-> (Int, Vector (Either Atom Atom))
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl' (Int, Vector (Either Atom Atom))
-> Chain -> (Int, Vector (Either Atom Atom))
addTer (Int
0, Vector (Either Atom Atom)
forall a. Monoid a => a
mempty)
      where
        addTer :: (Int, Vector (Either Atom TerAtom)) -> Chain -> (Int, Vector (Either Atom TerAtom))
        addTer :: (Int, Vector (Either Atom Atom))
-> Chain -> (Int, Vector (Either Atom Atom))
addTer (Int
add, Vector (Either Atom Atom)
res) Chain
cur = if (Atom -> Bool) -> Chain -> Bool
forall a. (a -> Bool) -> Vector a -> Bool
V.all (Text -> Bool
isHetatm (Text -> Bool) -> (Atom -> Text) -> Atom -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Atom -> Text
atomResName) Chain
cur then (Int
add, Vector (Either Atom Atom)
newRes) else (Int, Vector (Either Atom Atom))
withTer
          where
            ter :: Atom
ter    = Int -> Atom -> Atom
addSerial (Int
add Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Atom -> Atom) -> Atom -> Atom
forall a b. (a -> b) -> a -> b
$ Chain -> Atom
forall a. Vector a -> a
V.last Chain
cur
            newRes :: Vector (Either Atom Atom)
newRes = Vector (Either Atom Atom)
res Vector (Either Atom Atom)
-> Vector (Either Atom Atom) -> Vector (Either Atom Atom)
forall a. Semigroup a => a -> a -> a
<> (Atom -> Either Atom Atom) -> Chain -> Vector (Either Atom Atom)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Atom -> Either Atom Atom
forall a b. a -> Either a b
Left (Atom -> Either Atom Atom)
-> (Atom -> Atom) -> Atom -> Either Atom Atom
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Atom -> Atom
addSerial Int
add) Chain
cur

            withTer :: (Int, Vector (Either Atom Atom))
withTer = (Int
add Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Vector (Either Atom Atom)
newRes Vector (Either Atom Atom)
-> Vector (Either Atom Atom) -> Vector (Either Atom Atom)
forall a. Semigroup a => a -> a -> a
<> Either Atom Atom -> Vector (Either Atom Atom)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Atom -> Either Atom Atom
forall a b. b -> Either a b
Right Atom
ter))

            addSerial :: Int -> Atom -> Atom
            addSerial :: Int -> Atom -> Atom
addSerial Int
i at :: Atom
at@Atom{Char
Float
Int
Text
atomCharge :: Atom -> Text
atomElement :: Atom -> Text
atomTempFactor :: Atom -> Float
atomOccupancy :: Atom -> Float
atomZ :: Atom -> Float
atomY :: Atom -> Float
atomX :: Atom -> Float
atomICode :: Atom -> Char
atomResSeq :: Atom -> Int
atomChainID :: Atom -> Char
atomAltLoc :: Atom -> Char
atomName :: Atom -> Text
atomSerial :: Atom -> Int
atomCharge :: Text
atomElement :: Text
atomTempFactor :: Float
atomOccupancy :: Float
atomZ :: Float
atomY :: Float
atomX :: Float
atomICode :: Char
atomResSeq :: Int
atomChainID :: Char
atomResName :: Text
atomAltLoc :: Char
atomName :: Text
atomSerial :: Int
atomResName :: Atom -> Text
..} = Atom
at { atomSerial :: Int
atomSerial = Int
atomSerial Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i }

    atomOrTer :: Either Atom TerAtom -> Text
    atomOrTer :: Either Atom Atom -> Text
atomOrTer = (Atom -> Text) -> (Atom -> Text) -> Either Atom Atom -> Text
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Atom -> Text
atomToText Atom -> Text
terAtomToText

terAtomToText :: Atom -> Text
terAtomToText :: Atom -> Text
terAtomToText Atom
at = Text -> Text
toPDBLine (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Text
pref Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText Int
6 Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
suf
  where
    t :: Text
t           = (Text
ter Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Text -> Text
T.take Int
21 (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Text -> Text
T.drop Int
6 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Atom -> Text
atomToText Atom
at
    (Text
pref, Text
suf) = Int -> Text -> Text
T.drop Int
6 (Text -> Text) -> (Text, Text) -> (Text, Text)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Text -> (Text, Text)
T.splitAt Int
11 Text
t

    ter :: Text
    ter :: Text
ter = Text
"TER   "

atomToText :: Atom -> Text
atomToText :: Atom -> Text
atomToText Atom{Char
Float
Int
Text
atomCharge :: Text
atomElement :: Text
atomTempFactor :: Float
atomOccupancy :: Float
atomZ :: Float
atomY :: Float
atomX :: Float
atomICode :: Char
atomResSeq :: Int
atomChainID :: Char
atomResName :: Text
atomAltLoc :: Char
atomName :: Text
atomSerial :: Int
atomCharge :: Atom -> Text
atomElement :: Atom -> Text
atomTempFactor :: Atom -> Float
atomOccupancy :: Atom -> Float
atomZ :: Atom -> Float
atomY :: Atom -> Float
atomX :: Atom -> Float
atomICode :: Atom -> Char
atomResSeq :: Atom -> Int
atomChainID :: Atom -> Char
atomAltLoc :: Atom -> Char
atomName :: Atom -> Text
atomSerial :: Atom -> Int
atomResName :: Atom -> Text
..} = Text
res
  where
    recordName :: Text
recordName | Text -> Bool
isHetatm Text
atomResName = Text
hetatm
               | Bool
otherwise            = Text
atm

    serial :: Text
serial     = Int -> Int -> Text
forall a. Show a => Int -> a -> Text
prependToS Int
5 Int
atomSerial
    name :: Text
name       = (\Text
t -> if Text -> Char
T.last Text
t Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
space then Char -> Text -> Text
T.cons Char
space (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Text -> Text
T.init Text
t else Text
t) (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Text -> Text
appendTo Int
4 Text
atomName
    altLoc :: Text
altLoc     = Char -> Text
T.singleton Char
atomAltLoc
    resName :: Text
resName    = Int -> Text -> Text
prependTo Int
3 Text
atomResName
    chainID :: Text
chainID    = Char -> Text
T.singleton Char
atomChainID
    resSeq :: Text
resSeq     = Int -> Int -> Text
forall a. Show a => Int -> a -> Text
prependToS Int
4 Int
atomResSeq
    iCode :: Text
iCode      = Char -> Text
T.singleton Char
atomICode
    x :: Text
x          = Int -> Text -> Text
prependTo Int
8 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Float -> Text
printFloat Int
3 Float
atomX
    y :: Text
y          = Int -> Text -> Text
prependTo Int
8 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Float -> Text
printFloat Int
3 Float
atomY
    z :: Text
z          = Int -> Text -> Text
prependTo Int
8 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Float -> Text
printFloat Int
3 Float
atomZ
    occupancy :: Text
occupancy  = Int -> Text -> Text
prependTo Int
6 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Float -> Text
printFloat Int
2 Float
atomOccupancy
    tempFactor :: Text
tempFactor = Int -> Text -> Text
prependTo Int
6 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Int -> Float -> Text
printFloat Int
2 Float
atomTempFactor
    element :: Text
element    = Int -> Text -> Text
prependTo Int
2 Text
atomElement

    charge :: Text
charge | Text
atomCharge Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
zeroCharge = Int -> Text -> Text
prependTo Int
2 Text
atomCharge
           | Bool
otherwise                = Int -> Text
spaceText Int
2

    res :: Text
res = Text
recordName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
serial Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText Int
1 Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
altLoc
                     Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
resName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText Int
1 Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
chainID Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
resSeq Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
iCode Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText Int
3
                     Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
y Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
z Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
occupancy Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
tempFactor Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText Int
10
                     Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
element Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
charge

    atm :: Text
    atm :: Text
atm = Text
"ATOM  "

    hetatm :: Text
    hetatm :: Text
hetatm = Text
"HETATM"

    zeroCharge :: Text
    zeroCharge :: Text
zeroCharge = Text
"0"

    printFloat :: Int -> Float -> Text
    printFloat :: Int -> Float -> Text
printFloat Int
after Float
f = FilePath -> Text
T.pack (FilePath -> Text) -> FilePath -> Text
forall a b. (a -> b) -> a -> b
$ FilePath -> Int -> Float -> FilePath
forall r. PrintfType r => FilePath -> r
printf FilePath
"%.*f" Int
after Float
f

--------------------------------------------------------------------------------
-- Utility functions.
--------------------------------------------------------------------------------

isHetatm :: Text -> Bool
isHetatm :: Text -> Bool
isHetatm = (Text -> [Text] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Text]
canonicalAminoAcids)
  where
    canonicalAminoAcids :: [Text]
canonicalAminoAcids = [ Text
"ACE", Text
"ALA", Text
"ARG", Text
"ASN", Text
"ASP", Text
"CYS", Text
"GLU", Text
"GLN"
                          , Text
"GLY", Text
"HIS", Text
"HID", Text
"HIE", Text
"HIP", Text
"ILE", Text
"LEU", Text
"LYS", Text
"LYN"
                          , Text
"MET", Text
"NMA", Text
"PHE", Text
"PRO", Text
"SER", Text
"THR", Text
"TRP", Text
"TYR", Text
"VAL"
                          ]

toPDBLine :: Text -> Text
toPDBLine :: Text -> Text
toPDBLine = Int -> Text -> Text
appendTo Int
80

prependToS :: Show a => Int -> a -> Text
prependToS :: Int -> a -> Text
prependToS Int
i = Int -> Text -> Text
prependTo Int
i (Text -> Text) -> (a -> Text) -> a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FilePath -> Text
T.pack (FilePath -> Text) -> (a -> FilePath) -> a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> FilePath
forall a. Show a => a -> FilePath
show

prependTo :: Int -> Text -> Text
prependTo :: Int -> Text -> Text
prependTo Int
i Text
t = Int -> Text
spaceText (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
T.length Text
t) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
t

appendTo :: Int -> Text -> Text
appendTo :: Int -> Text -> Text
appendTo Int
i Text
t = Text
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
spaceText (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
T.length Text
t)

newLine :: Text
newLine :: Text
newLine = Text
"\n"

spaceText :: Int -> Text
spaceText :: Int -> Text
spaceText = (Int -> Text -> Text) -> Text -> Int -> Text
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Text -> Text
T.replicate Text
" "

space :: Char
space :: Char
space = Char
' '