import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
file_path = "../input/corona-virus-proteins-from-uniprot-database/corona.csv"
df = pd.read_csv(file_path)

df.head(5)

df.shape

(21876, 8)

df_organism = pd.DataFrame(df.groupby("Organism").count()['Entry'])
df_organism = df_organism.sort_values(by = "Entry", ascending = False)
df_organism[0:20].plot.barh(figsize = [15,10], fontsize =20)
plt.gca().invert_yaxis()

df_organism[0:20]

df['Virus hosts'] = df['Virus hosts'].apply(lambda x: str(x)[0:50] )
df_host = pd.DataFrame(df.groupby("Virus hosts").count()['Entry'])
df_host = df_host.sort_values(by = "Entry", ascending = False)
df_host[1:20].plot.barh(figsize = [15,10], fontsize =20)
plt.gca().invert_yaxis()

df_host[1:20]

def filter(line):
    proteins = set()
    line = str(line)
    line = line.lower()

    '''for lines without () or [] terms'''
    if "(" not in line or "[" not in line:
        proteins.add(line.strip().replace(' ', '_'))


    '''for line including () terms'''    
    if '(' in line:
        start = 0
        open_in = line.find('(')
        tmp = line[start:open_in].strip().replace(' ', '_')
        proteins.add(tmp)
        while open_in >=0:
            start = open_in+1
            end = line.find(')', start)
            proteins.add(line[start:end].strip().replace(' ', '_'))
            open_in = line.find('(', end)

    '''for lines including [] trems'''
    if '[' in line:
        raw = line[line.find('['):line.find(']')]
        #print("THIS IS RAW:", raw[15:-1])
        raw = raw[15:-1]
        lraw = raw.split("; ")
        for item in lraw:
            #print(item)
            if '(' in item:
                start = 0
                open_in = item.find('(')
                tmp = item[start:open_in].strip().replace(' ', '_')
                proteins.add(tmp)
            else:
                proteins.add(item.strip().replace(' ', '_'))
    return proteins

allProteins = []
i = 0
for u,p in zip(df['Entry'],df['Protein names']):
    print(u,"|",p)
    print("------------")
    print(u,"|",filter(p))
    print("===================================================")
    i += 1
    if i>4:
        break

A0A3R5SMJ6 | Genome polyprotein
------------
A0A3R5SMJ6 | {'genome_polyprotein'}
===================================================
M1UFP6 | Genome polyprotein
------------
M1UFP6 | {'genome_polyprotein'}
===================================================
P11223 | Spike glycoprotein (S glycoprotein) (E2) 
(Peplomer protein) [Cleaved into: Spike protein S1; 
Spike protein S2; Spike protein S2']
------------
P11223 | {'peplomer_protein', 'spike_protein_s1',
'e2', 'spike_protein_s2', 'spike_glycoprotein',
's_glycoprotein'}
===================================================
P11224 | Spike glycoprotein (S glycoprotein) (E2) 
(Peplomer protein) [Cleaved into: Spike protein S1;
Spike protein S2; Spike protein S2']
------------
P11224 | {'peplomer_protein', 'spike_protein_s1',
'e2', 'spike_protein_s2', 'spike_glycoprotein',
's_glycoprotein'}
===================================================

allProteins = []
for u,p in zip(df['Entry'],df['Protein names']):
    allProteins.append({"id":u, "names":list(filter(p))})

allProteins[0:5]

[{'id': 'A0A3R5SMJ6', 'names': ['genome_polyprotein']},
 {'id': 'M1UFP6', 'names': ['genome_polyprotein']},
 {'id': 'P11223',
  'names': ['peplomer_protein',
   'spike_protein_s1',
   'e2',
   'spike_protein_s2',
   'spike_glycoprotein',
   's_glycoprotein']},
 {'id': 'P11224',
  'names': ['peplomer_protein',
   'spike_protein_s1',
   'e2',
   'spike_protein_s2',
   'spike_glycoprotein',
   's_glycoprotein']},
 {'id': 'P0C6X9',
  'names': ['m-pro',
   'nsp16',
   'nsp7',
   'exon',
   'nsp14',
   'nsp10',
   'guanine-n7_methyltransferase',
   'ec_3.4.22.-',
   'nendou',
   'non-structural_protein_3',
   'non-structural_protein_2',
   'ec_2.1.1.-',
   "2'-o-methyltransferase",
   'hel',
   'pol',
   'ec_3.6.4.13',
   'growth_factor-like_peptide',
   'non-structural_protein_7',
   'gfl',
   'ec_3.4.22.69',
   'p22',
   'p27',
   'non-structural_protein_9',
   'orf1ab_polyprotein',
   'nsp4',
   'uridylate-specific_endoribonuclease',
   '3cl-pro',
   'pp1ab',
   'p65',
   'host_translation_inhibitor_nsp1',
   'p15',
   'nsp1',
   'ec_3.6.4.12',
   'nsp2',
   'ec_2.7.7.48',
   'p67',
   'nsp12',
   'peptide_hd2',
   'nsp5',
   'p210',
   'rdrp',
   'nsp9',
   'p100',
   'nsp3',
   'nsp8',
   'non-structural_protein_6',
   'rna-directed_rna_polymerase',
   'p35',
   '3c-like_proteinase',
   'papain-like_proteinase',
   'ec_3.1.13.-',
   'non-structural_protein_4',
   'pl-pro',
   '3clp',
   'p12',
   'non-structural_protein_10',
   'p28',
   'non-structural_protein_8',
   'helicase',
   'nsp6',
   'nsp13',
   'nsp15',
   'p44',
   'ec_3.1.-.-',
   'replicase_polyprotein_1ab',
   'ec_3.4.19.12',
   'p10']}]

import json
with open("virus-proteins.json", 'w') as fn:
    json.dump(allProteins,fn)