# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Downloads PDB Structure Information
# **[Work in progress]**
#
# This notebook downloads 3D-structure information for SARS-CoV-2 proteins
#
# Data sources: 
# [RCSB Protein Data Bank](https://www.rcsb.org), 
# [PDBe Protein Data Bank Europe](https://www.ebi.ac.uk/pdbe/)
#
# Author: <NAME> (<EMAIL>)

import os
import numpy as np
import pandas as pd
import requests
import json
import dateutil
from pathlib import Path
from py2neo import Graph
from rcsbsearch import TextQuery
from rcsbsearch import rcsb_attributes as attrs

pd.options.display.max_rows = None  # display all rows
pd.options.display.max_columns = None  # display all columsns

NEO4J_IMPORT = Path(os.getenv('NEO4J_IMPORT'))
print(NEO4J_IMPORT)

taxonomy_id = '2697049'

# ### Find PDB Structures containing SARS-CoV-2 proteins or nucleic acids

query = attrs.rcsb_entity_source_organism.taxonomy_lineage.id == taxonomy_id
entries = query('entry')
df = pd.DataFrame(entries, columns=['pdbId'])

df.head()

print("Number of structures:", df['pdbId'].shape[0])


# ### Get structure metadata

def get_pdb_entry_data(row):
    pdb_id = row['pdbId']
    
    response = json.loads(requests.get(f'http://data.rcsb.org/rest/v1/core/entry/{pdb_id}').text)
    
    data = np.empty(9, dtype=object)

    data[0] = pdb_id
    data[1] = response['struct']['title']
    data[2] = response['struct'].get('pdbx_descriptor', '')
    data[3] = response['rcsb_accession_info']['deposit_date']
    data[4] = response['rcsb_accession_info']['initial_release_date']
    data[5] = response['pdbx_vrpt_summary'].get('pdbresolution', '')
    data[6] = response['pdbx_vrpt_summary'].get('pdbr', '')
    data[7] = response['pdbx_vrpt_summary'].get('pdbrfree', '')      
    data[8] = response['exptl'][0]['method']
    
    return data


structures = df.apply(get_pdb_entry_data, axis=1, result_type='expand')

structures.columns = ['pdbId', 'title', 'description', 'depositDate', 'releaseDate', 'resolution', 'rFactor', 'rFree', 'method']

structures['depositDate'] = structures['depositDate'].apply(lambda d: dateutil.parser.parse(d[:10]))
structures['releaseDate'] = structures['releaseDate'].apply(lambda d: dateutil.parser.parse(d[:10]))

structures['pdbId'] = 'pdb:' + structures['pdbId']

structures.head()

structures.to_csv(NEO4J_IMPORT / "01f-PDBStructure.csv", index=False)

# ### Get PDB Chain - UniProt sequence mappings

sifts_url = 'http://ftp.ebi.ac.uk/pub/databases/msd/sifts/flatfiles/tsv/uniprot_segments_observed.tsv.gz'

segments = pd.read_csv(sifts_url, sep='\t', skiprows=1, dtype=str)
segments.head()

segments.rename(columns={'PDB': 'pdbId', 'CHAIN': 'chainId', 'SP_PRIMARY': 'accession'}, inplace=True)
segments.rename(columns={'RES_BEG': 'seqresStart', 'RES_END': 'seqresEnd'}, inplace=True)
segments.rename(columns={'PDB_BEG': 'pdbStart', 'PDB_END': 'pdbEnd'}, inplace=True)
segments.rename(columns={'SP_BEG': 'uniprotStart', 'SP_END': 'uniprotEnd'}, inplace=True)

segments.head()

segments['pdbId'] = segments['pdbId'].str.upper()

segments['pdbChainId'] = segments['pdbId'] + "." + segments['chainId']

coverage = segments.merge(df, on='pdbId')

# ### Sort segments by uniprot residue number

coverage['uniprotStart'] = coverage['uniprotStart'].astype(int)
coverage['uniprotEnd'] = coverage['uniprotEnd'].astype(int)
coverage['length'] = coverage['uniprotEnd'] - coverage['uniprotStart'] + 1
coverage.sort_values(by='uniprotStart', inplace=True)
coverage['uniprotStart'] = coverage['uniprotStart'].astype(str)
coverage['uniprotEnd'] = coverage['uniprotEnd'].astype(str)

# ### Assign CURIES

coverage['accession'] = 'uniprot:' + coverage['accession']
coverage['pdbId'] = 'pdb:' + coverage['pdbId']
coverage['pdbChainId'] = 'pdb:' + coverage['pdbChainId']

coverage.head()

# ### Group data by PDB chains

coverage = coverage.groupby(['pdbId','chainId','pdbChainId','accession']).agg(list).reset_index()

coverage.tail()

# ### Create semicolon separated string of residue numbers so they can be represented in a csv file

coverage['uniprotStart'] = coverage['uniprotStart'].apply(lambda x: ';'.join(x))
coverage['uniprotEnd'] = coverage['uniprotEnd'].apply(lambda x: ';'.join(x))
coverage['seqresStart'] = coverage['seqresStart'].apply(lambda x: ';'.join(x))
coverage['seqresEnd'] = coverage['seqresEnd'].apply(lambda x: ';'.join(x))
coverage['pdbStart'] = coverage['pdbStart'].apply(lambda x: ';'.join(x))
coverage['pdbEnd'] = coverage['pdbEnd'].apply(lambda x: ';'.join(x))

coverage['residues'] = coverage['length'].apply(lambda x: sum(x))

coverage.tail()

coverage.to_csv(NEO4J_IMPORT / "01f-PDBChain.csv", index=False)

StarcoderdataJupyterScriptsDedupFiltered