<commit_before>#!/usr/bin/python
# encoding: utf-8
"""
ligand.py
Basic parsing for the PDB ligand .cif and XML formats.

Created by <NAME> 2015
"""

import os
from io import BytesIO
import re
import copy
from PIL import Image

from klab.bio.rcsb import retrieve_ligand_cif, retrieve_pdb_ligand_info, retrieve_ligand_diagram
from klab.bio.basics import three_letter_ion_codes
from klab.fs.fsio import read_file, write_file
from klab import colortext



class Ligand(object):

    '''A class used to store ligand information from the RCSB. The most useful way to use this class is to call
       Ligand.retrieve_data_from_rcsb to lookup the RCSB for the ligand data.'''

    def __init__(self, ligand_code):
        self.PDBCode = ligand_code
        self.LigandCode = None
        self.Formula = None
        self.MolecularWeight = None
        self.LigandType = None
        self.Diagram = None
        self.SimilarCompoundsDiagram = None
        self.InChI = None
        self.InChIKey = None
        self.pdb_id = None
        self.has_many_atoms = None # The PDB file for 1BIZ lists CAC as having one atom. This property is intended to simply record when a ligand file has many atoms so we treat it as a ligand rather than a charged atom.

        if ligand_code == 'UNL':
            # Unknown ligand
            self.has_many_atoms = True
        elif ligand_code == 'UNX':
            # Unknown atom or ion
            self.has_many_atoms = False

         # These fields are not filled in automatically
        self.Solubility = None
        self.CellPermeability = None
        self.AssaysToDetermineConcentrationInCells = None
        self.ProductionInCells = None
        self.ProductionInCellsNotes = None

        # These fields are not guaranteed to be filled in
        self.Diagram = None
        self.SimilarCompoundsDiagram = None

        self.descriptors = []
        self.identifiers = []
        self.synonyms = []


    def __repr__(self):
        s =  'Ligand      : {0}, {1}, {2}\n'.format(self.PDBCode, self.LigandCode, self.Formula)
        if self.synonyms:
            s += '              AKA {0}\n'.format(', '.join(self.synonyms))
        s += 'Type        : {0}\n'.format(self.LigandType)
        s += 'Weight      : {0} g/mol\n'.format(self.MolecularWeight)
        s += 'InChI       : {0} ({1})\n'.format(self.InChI, self.InChIKey)
        if self.descriptors:
            s += 'Descriptors :\n'
            for d in self.descriptors:
                s += '              {0}    ({1}, {2} {3})\n'.format(d['Descriptor'], d['DescriptorType'], d['Program'], d['Version'])
        if self.identifiers:
            s += 'Identifiers :\n'
            for i in self.identifiers:
                s += '              {0}    ({1}, {2} {3})\n'.format(i['Identifier'], i['IDType'], i['Program'], i['Version'])
        if self.Diagram:
            file = BytesIO(self.Diagram)
            img = Image.open(file)
            w, h = img.size
            s += 'Diagram     : {0}x{1}'.format(w,h)
        if self.pdb_id:
            s += '\nAss. PDB    : {0}'.format(self.pdb_id)
        return s


    @classmethod
    def retrieve_data_from_rcsb(cls, ligand_code, pdb_id = None, silent = True, cached_dir = None):
        '''Retrieve a file from the RCSB.'''
        if not silent:
            colortext.printf("Retrieving data from RCSB")
        if cached_dir:
            assert(os.path.exists(cached_dir))

        ligand_info_path, ligand_info, pdb_ligand_info, pdb_ligand_info_path = None, None, None, None
        if cached_dir:
            ligand_info_path = os.path.join(cached_dir, '{0}.cif'.format(ligand_code))
            if os.path.exists(ligand_info_path):
                ligand_info = read_file(ligand_info_path)
        if not ligand_info:
            ligand_info = retrieve_ligand_cif(ligand_code)
            if cached_dir:
                write_file(ligand_info_path, ligand_info)

        # Parse .cif
        l = cls(ligand_code)
        l.parse_cif(ligand_info)
        l.pdb_id = pdb_id or l.pdb_id

        # Parse PDB XML
        pdb_ligand_info_path = None
        if cached_dir:
            pdb_ligand_info_path = os.path.join(cached_dir, '{0}.pdb.ligandinfo'.format(l.pdb_id.lower()))
        if len(l.pdb_id) == 4:
            ### TODO: This (somewhat nasty) hack was made by Shane, so should be fixed by Shane, despite git blame blaming Kyle
            if cached_dir:
                if os.path.exists(pdb_ligand_info_path):
                    pdb_ligand_info = read_file(pdb_ligand_info_path)
                else:
                    pdb_ligand_info = retrieve_pdb_ligand_info(l.pdb_id)
                    write_file(pdb_ligand_info_path, pdb_ligand_info)
            else:
                pdb_ligand_info = retrieve_pdb_ligand_info(l.pdb_id)
        if pdb_ligand_info:
            l.parse_pdb_ligand_info(pdb_ligand_info)

        # Retrive the diagram image
        l.get_diagram()

        return l


    def parse_cif(self, cif):
        '''See http://www.iucr.org/__data/iucr/cif/standard/cifstd4.html, Acta Cryst. (1991). A47, 655-685.'''
        found_cif_header = False
        found_cif_descriptors = False
        found_cif_identifiers = False

        blocks = []
        blocklines = []<commit_msg>Fixing a bug where the rain got in. (Hopefully, I need to test it now)<commit_after>#!/usr/bin/python
# encoding: utf-8
"""
ligand.py
Basic parsing for the PDB ligand .cif and XML formats.

Created by <NAME> 2015
"""

import os
from io import BytesIO
import re
import copy
from PIL import Image

from klab.bio.rcsb import retrieve_ligand_cif, retrieve_pdb_ligand_info, retrieve_ligand_diagram
from klab.bio.basics import three_letter_ion_codes
from klab.fs.fsio import read_file, write_file
from klab import colortext



class Ligand(object):

    '''A class used to store ligand information from the RCSB. The most useful way to use this class is to call
       Ligand.retrieve_data_from_rcsb to lookup the RCSB for the ligand data.'''

    def __init__(self, ligand_code):
        self.PDBCode = ligand_code
        self.LigandCode = None
        self.Formula = None
        self.MolecularWeight = None
        self.LigandType = None
        self.Diagram = None
        self.SimilarCompoundsDiagram = None
        self.InChI = None
        self.InChIKey = None
        self.pdb_id = None
        self.has_many_atoms = None # The PDB file for 1BIZ lists CAC as having one atom. This property is intended to simply record when a ligand file has many atoms so we treat it as a ligand rather than a charged atom.

        if ligand_code == 'UNL':
            # Unknown ligand
            self.has_many_atoms = True
        elif ligand_code == 'UNX':
            # Unknown atom or ion
            self.has_many_atoms = False

         # These fields are not filled in automatically
        self.Solubility = None
        self.CellPermeability = None
        self.AssaysToDetermineConcentrationInCells = None
        self.ProductionInCells = None
        self.ProductionInCellsNotes = None

        # These fields are not guaranteed to be filled in
        self.Diagram = None
        self.SimilarCompoundsDiagram = None

        self.descriptors = []
        self.identifiers = []
        self.synonyms = []


    def __repr__(self):
        s =  'Ligand      : {0}, {1}, {2}\n'.format(self.PDBCode, self.LigandCode, self.Formula)
        if self.synonyms:
            s += '              AKA {0}\n'.format(', '.join(self.synonyms))
        s += 'Type        : {0}\n'.format(self.LigandType)
        s += 'Weight      : {0} g/mol\n'.format(self.MolecularWeight)
        s += 'InChI       : {0} ({1})\n'.format(self.InChI, self.InChIKey)
        if self.descriptors:
            s += 'Descriptors :\n'
            for d in self.descriptors:
                s += '              {0}    ({1}, {2} {3})\n'.format(d['Descriptor'], d['DescriptorType'], d['Program'], d['Version'])
        if self.identifiers:
            s += 'Identifiers :\n'
            for i in self.identifiers:
                s += '              {0}    ({1}, {2} {3})\n'.format(i['Identifier'], i['IDType'], i['Program'], i['Version'])
        if self.Diagram:
            file = BytesIO(self.Diagram)
            img = Image.open(file)
            w, h = img.size
            s += 'Diagram     : {0}x{1}'.format(w,h)
        if self.pdb_id:
            s += '\nAss. PDB    : {0}'.format(self.pdb_id)
        return s


    @classmethod
    def retrieve_data_from_rcsb(cls, ligand_code, pdb_id = None, silent = True, cached_dir = None):
        '''Retrieve a file from the RCSB.'''
        if not silent:
            colortext.printf("Retrieving data from RCSB")
        if cached_dir:
            assert(os.path.exists(cached_dir))

        ligand_info_path, ligand_info, pdb_ligand_info, pdb_ligand_info_path = None, None, None, None
        if cached_dir:
            ligand_info_path = os.path.join(cached_dir, '{0}.cif'.format(ligand_code))
            if os.path.exists(ligand_info_path):
                ligand_info = read_file(ligand_info_path)
        if not ligand_info:
            ligand_info = retrieve_ligand_cif(ligand_code)
            if cached_dir:
                write_file(ligand_info_path, ligand_info)

        # Parse .cif
        l = cls(ligand_code)
        l.parse_cif(ligand_info)
        l.pdb_id = pdb_id or l.pdb_id
        has_pdb_id = l.pdb_id and (len(l.pdb_id) == 4) and (l.pdb_id != '?')  # the last case is unnecessary and will be short-cut but I included it to show possible values

        # Parse PDB XML
        if has_pdb_id:
            if cached_dir:
                pdb_ligand_info_path = os.path.join(cached_dir, '{0}.pdb.ligandinfo'.format(l.pdb_id.lower()))
                if os.path.exists(pdb_ligand_info_path):
                    pdb_ligand_info = read_file(pdb_ligand_info_path)
                else:
                    pdb_ligand_info = retrieve_pdb_ligand_info(l.pdb_id)
                    write_file(pdb_ligand_info_path, pdb_ligand_info)
            else:
                pdb_ligand_info = retrieve_pdb_ligand_info(l.pdb_id)
        if pdb_ligand_info:
            l.parse_pdb_ligand_info(pdb_ligand_info)

        # Retrive the diagram image
        l.get_diagram()

        return l


    def parse_cif(self, cif):
        '''See http://www.iucr.org/__data/iucr/cif/standard/cifstd4.html, Acta Cryst. (1991). A47, 655-685.'''
        found_cif_header = False
        found_cif_descriptors = False
        found_cif_identifiers = False

        blocks = []
        blocklines = []
{'fraction_non_alphanumeric': 0.09970674486803519, 'fraction_numerical': 0.007757071232617539, 'mean_word_length': 2.1139911634756996, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '763171', 'n_tokens_mistral': 3562, 'n_tokens_neox': 3377, 'n_words': 835}