Functions.hs

module Bio.PDB.Functions
  ( groupChainByResidue
  ) where

import qualified Bio.PDB.Type  as PDB (Atom (..))
import qualified Bio.Utils.Map as M ((!?!))

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M (fromList)
import           Data.List       (groupBy,
                                  sortOn)
import           Data.Vector     (Vector)
import qualified Data.Vector      as V (toList)
import           Data.Char       (toUpper)

groupChainByResidue :: Vector PDB.Atom -> [[PDB.Atom]]
groupChainByResidue = sortOn sortOnResidue . groupBy atomsFromSameResidue . V.toList
  where
    atomsFromSameResidue :: PDB.Atom -> PDB.Atom -> Bool
    atomsFromSameResidue atom1 atom2 = PDB.atomResSeq atom1 == PDB.atomResSeq atom2 && PDB.atomICode atom1 == PDB.atomICode atom2

    sortOnResidue :: [PDB.Atom] -> Int
    sortOnResidue [] = -1000000
    sortOnResidue (PDB.Atom{..} : _) = atomSerial * 100 + (insertionCodeSortingCorrections M.!?! toUpper atomICode)

    insertionCodeSortingCorrections :: Map Char Int
    insertionCodeSortingCorrections = M.fromList $ zip (' ':['A'..'Z']) [0..]