<commit_before># -*- coding: utf-8 -*-
import os.path as op
import gzip
import string
import logging
import urllib.request
import re

from functools import wraps
from tempfile import NamedTemporaryFile
from collections import defaultdict, OrderedDict

import six
from fastcache import clru_cache

import numpy as np

import Bio
from Bio.PDB import PDBIO, Select, NeighborSearch
from Bio.PDB.PDBParser import PDBParser
from Bio.PDB.MMCIFParser import MMCIFParser
from Bio.PDB.Polypeptide import PPBuilder
from Bio.Alphabet import IUPAC
from Bio.Seq import Seq

from . import errors, helper

logger = logging.getLogger(__name__)

A_DICT = {
    'A': 'ALA', 'R': 'ARG', 'N': 'ASN', 'D': 'ASP', 'C': 'CYS', 'E': 'GLU',
    'Q': 'GLN', 'G': 'GLY', 'H': 'HIS', 'I': 'ILE', 'L': 'LEU', 'K': 'LYS',
    'M': 'MET', 'F': 'PHE', 'P': 'PRO', 'S': 'SER', 'T': 'THR', 'W': 'TRP',
    'Y': 'TYR', 'V': 'VAL', 'U': 'SEC', 'O': 'PYL',
    'B': 'ASX', 'Z': 'GLX', 'J': 'XLE', 'X': 'XAA', '*': 'TER'
}
AAA_DICT = dict([(value, key) for key, value in list(A_DICT.items())])
AAA_DICT['UNK'] = 'X'
AAA_DICT['MSE'] = 'M'
AAA_DICT['CSD'] = 'C'
# Phosphorylated residues
# AAA_DICT['SEP'] = 'S' # PHOSPHOSERINE
# AAA_DICT['TPO'] = 'T' # PHOSPHOTHREONINE
# AAA_DICT['SEP'] = 'Y' # O-PHOSPHOTYROSINE

# Methylated lysines
AAA_DICT['MLZ'] = 'K'
AAA_DICT['MLY'] = 'K'
AAA_DICT['M3L'] = 'K'

AMINO_ACIDS = list(AAA_DICT.keys())

METHYLATED_LYSINES = ['MLZ', 'MLY', 'M3L']
LYSINE_ATOMS = ['N', 'CA', 'CB', 'CG', 'CD', 'CE', 'NZ', 'C', 'O']


# %% Functions for downloading and parsing pdb files
class MMCIFParserMod(MMCIFParser):
    def __init__(self, temp_dir):
        self.temp_dir = temp_dir

    def get_structure(self, structure_id, gzip_fh):
        """
        Altered `get_structure` method which accepts gzip file handles as input.
        """
        with NamedTemporaryFile(mode='w', dir=self.temp_dir) as temp_fh:
            temp_fh.writelines(gzip_fh.readlines())
            temp_fh.flush()
            temp_fh.seek(0)
            return super(MMCIFParserMod, self).get_structure(structure_id, temp_fh.name)


def get_pdb_id(pdb_file):
    pdb_id = op.basename(pdb_file)
    for ext in ['\.gz$', '\.pdb$', '\.ent$', '\.cif$', '^pdb']:
        pdb_id = re.sub(ext, '', pdb_id)
    pdb_id = pdb_id.upper()
    return pdb_id.upper()


def get_pdb_file(pdb_id, pdb_database_dir, pdb_type='ent'):
    """Get PDB file from a local mirror of the PDB database.
    """
    if pdb_type == 'ent':
        # Original PDB structure.
        prefix = 'pdb'
        suffix = '.ent.gz'
        relative_pdb_file = (
            pdb_id[1:3].lower() + '/' + prefix + pdb_id.lower() + suffix
        )
    elif pdb_type == 'cif':
        # mmCIF pdb structure.
        prefix = ''
        suffix = '.cif.gz'
        relative_pdb_file = (
            '../mmCIF/' +
            pdb_id[1:3].lower() + '/' + prefix + pdb_id.lower() + suffix
        )
    elif pdb_type == 'pdb':
        # The first biological unit.
        prefix = ''
        suffix = '.pdb1.gz'
        relative_pdb_file = (
            '../../../biounit/coordinates/divided/' +
            pdb_id[1:3].lower() + '/' + prefix + pdb_id.lower() + suffix
        )
    elif pdb_type == 'raw':
        # Just a PDB file in some folder.
        relative_pdb_file = ''
    else:
        raise Exception

    pdb_file = op.join(pdb_database_dir, relative_pdb_file)
    return pdb_file


def download_pdb_file(pdb_id, output_dir):
    """Move PDB structure to the local working directory.
    """
    PDB_URL = 'http://www.rcsb.org/pdb/files/{}.pdb'

    output_pdb_filename = op.join(output_dir, pdb_id + '.pdb')

    # If the PDB already exists, do nothing...
    if op.isfile(output_pdb_filename):
        logger.debug('PDB file {} already exists...'.format(output_pdb_filename))
        return output_pdb_filename

    # Download the PDB file from the internet...
    logger.info('Downloading PDB {}...'.format(pdb_id + '.pdb'))
    response = urllib.request.urlopen(PDB_URL.format(pdb_id))
    with open(output_pdb_filename, 'wb') as ofh:
        ofh.write(response.read())

    return output_pdb_filename


def get_pdb_parser(pdb_type, temp_dir='/tmp'):
    """Get PDB parser that can work with structures of the specified type.
    """
    if pdb_type in {'ent', 'pdb', 'raw'}:
        parser = PDBParser(QUIET=True)
    elif pdb_type in {'cif'}:
        parser = MMCIFParserMod(temp_dir=temp_dir)
    else:
        raise Exception('Unsupported ``pdb_type`` {}'.format(pdb_type))
    return parser


@clru_cache(maxsize=128, typed=False)
def get_pdb(pdb_id, pdb_path, temp_dir='/tmp', pdb_type='ent', use_external=True):
    """
    Parse a pdb file with biopythons PDBParser() and return the structure.<commit_msg>Add PDB download fallback in case RCSB is offline<commit_after># -*- coding: utf-8 -*-
import os.path as op
import gzip
import string
import logging
import urllib.request
import re

from functools import wraps
from tempfile import NamedTemporaryFile
from collections import defaultdict, OrderedDict

import six
from fastcache import clru_cache

import numpy as np

import Bio
from Bio.PDB import PDBIO, Select, NeighborSearch
from Bio.PDB.PDBParser import PDBParser
from Bio.PDB.MMCIFParser import MMCIFParser
from Bio.PDB.Polypeptide import PPBuilder
from Bio.Alphabet import IUPAC
from Bio.Seq import Seq

from . import errors, helper

logger = logging.getLogger(__name__)

A_DICT = {
    'A': 'ALA', 'R': 'ARG', 'N': 'ASN', 'D': 'ASP', 'C': 'CYS', 'E': 'GLU',
    'Q': 'GLN', 'G': 'GLY', 'H': 'HIS', 'I': 'ILE', 'L': 'LEU', 'K': 'LYS',
    'M': 'MET', 'F': 'PHE', 'P': 'PRO', 'S': 'SER', 'T': 'THR', 'W': 'TRP',
    'Y': 'TYR', 'V': 'VAL', 'U': 'SEC', 'O': 'PYL',
    'B': 'ASX', 'Z': 'GLX', 'J': 'XLE', 'X': 'XAA', '*': 'TER'
}
AAA_DICT = dict([(value, key) for key, value in list(A_DICT.items())])
AAA_DICT['UNK'] = 'X'
AAA_DICT['MSE'] = 'M'
AAA_DICT['CSD'] = 'C'
# Phosphorylated residues
# AAA_DICT['SEP'] = 'S' # PHOSPHOSERINE
# AAA_DICT['TPO'] = 'T' # PHOSPHOTHREONINE
# AAA_DICT['SEP'] = 'Y' # O-PHOSPHOTYROSINE

# Methylated lysines
AAA_DICT['MLZ'] = 'K'
AAA_DICT['MLY'] = 'K'
AAA_DICT['M3L'] = 'K'

AMINO_ACIDS = list(AAA_DICT.keys())

METHYLATED_LYSINES = ['MLZ', 'MLY', 'M3L']
LYSINE_ATOMS = ['N', 'CA', 'CB', 'CG', 'CD', 'CE', 'NZ', 'C', 'O']


# %% Functions for downloading and parsing pdb files
class MMCIFParserMod(MMCIFParser):
    def __init__(self, temp_dir):
        self.temp_dir = temp_dir

    def get_structure(self, structure_id, gzip_fh):
        """
        Altered `get_structure` method which accepts gzip file handles as input.
        """
        with NamedTemporaryFile(mode='w', dir=self.temp_dir) as temp_fh:
            temp_fh.writelines(gzip_fh.readlines())
            temp_fh.flush()
            temp_fh.seek(0)
            return super(MMCIFParserMod, self).get_structure(structure_id, temp_fh.name)


def get_pdb_id(pdb_file):
    pdb_id = op.basename(pdb_file)
    for ext in ['\.gz$', '\.pdb$', '\.ent$', '\.cif$', '^pdb']:
        pdb_id = re.sub(ext, '', pdb_id)
    pdb_id = pdb_id.upper()
    return pdb_id.upper()


def get_pdb_file(pdb_id, pdb_database_dir, pdb_type='ent'):
    """Get PDB file from a local mirror of the PDB database.
    """
    if pdb_type == 'ent':
        # Original PDB structure.
        prefix = 'pdb'
        suffix = '.ent.gz'
        relative_pdb_file = (
            pdb_id[1:3].lower() + '/' + prefix + pdb_id.lower() + suffix
        )
    elif pdb_type == 'cif':
        # mmCIF pdb structure.
        prefix = ''
        suffix = '.cif.gz'
        relative_pdb_file = (
            '../mmCIF/' +
            pdb_id[1:3].lower() + '/' + prefix + pdb_id.lower() + suffix
        )
    elif pdb_type == 'pdb':
        # The first biological unit.
        prefix = ''
        suffix = '.pdb1.gz'
        relative_pdb_file = (
            '../../../biounit/coordinates/divided/' +
            pdb_id[1:3].lower() + '/' + prefix + pdb_id.lower() + suffix
        )
    elif pdb_type == 'raw':
        # Just a PDB file in some folder.
        relative_pdb_file = ''
    else:
        raise Exception

    pdb_file = op.join(pdb_database_dir, relative_pdb_file)
    return pdb_file


def download_pdb_file(pdb_id, output_dir):
    """Move PDB structure to the local working directory.
    """
    PDB_URL = 'http://www.rcsb.org/pdb/files/{}.pdb'
    PDB_EURO_URL = 'http://www.ebi.ac.uk/pdbe/entry-files/download/pdb{}.ent'

    output_pdb_filename = op.join(output_dir, pdb_id + '.pdb')

    # If the PDB already exists, do nothing...
    if op.isfile(output_pdb_filename):
        logger.debug('PDB file {} already exists...'.format(output_pdb_filename))
        return output_pdb_filename

    # Download the PDB file from the internet...
    logger.info('Downloading PDB {}...'.format(pdb_id + '.pdb'))
    try:
        pdb_url = PDB_URL.format(pdb_id)
        logger.debug(pdb_url)
        response = urllib.request.urlopen(pdb_url)
    except urllib.error.URLError as e:
        logger.debug(e)
        pdb_url = PDB_EURO_URL.format(pdb_id.lower())
        logger.debug(pdb_url)
        response = urllib.request.urlopen(pdb_url)

    with open(output_pdb_filename, 'wb') as ofh:
        ofh.write(response.read())

    return output_pdb_filename


def get_pdb_parser(pdb_type, temp_dir='/tmp'):
    """Get PDB parser that can work with structures of the specified type.
    """
    if pdb_type in {'ent', 'pdb', 'raw'}:
        parser = PDBParser(QUIET=True)
    elif pdb_type in {'cif'}:
        parser = MMCIFParserMod(temp_dir=temp_dir)
    else:
        raise Exception('Unsupported ``pdb_type`` {}'.format(pdb_type))
    return parser


@clru_cache(maxsize=128, typed=False)
def get_pdb(pdb_id, pdb_path, temp_dir='/tmp', pdb_type='ent', use_external=True):
    """
    Parse a pdb file with biopythons PDBParser() and return the structure.
{'fraction_non_alphanumeric': 0.1598501720996153, 'fraction_numerical': 0.0028345818991698723, 'mean_word_length': 3.0256723716381417, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2102515', 'n_tokens_mistral': 3825, 'n_tokens_neox': 3715, 'n_words': 871}