import pandas as pd
import numpy as np

d=pd.read_csv('histone_genes.csv')

hist_genes=(d[d['Biotype'].isin(['COD'])])[['Histone type','Histone variant','HGNC symbol','NCBI gene ID','Ensembl gene ID']]

from pybiomart import Dataset

dataset = Dataset(name='hsapiens_gene_ensembl',
                  host='http://www.ensembl.org')

genedata = dataset.query(attributes=['ensembl_gene_id','ensembl_transcript_id','ensembl_peptide_id','refseq_mrna','refseq_peptide','transcript_biotype'], only_unique=False)

histone_proteins=hist_genes.merge(genedata,left_on='Ensembl gene ID',right_on='Gene stable ID',suffixes=('','_y'),how='left').drop(columns='Gene stable ID')

import requests
seqs=[]
seqlen=[]
refs=[]
for index, row in histone_proteins.iterrows():
#     print(row['Protein stable ID'])
    seq=requests.get('http://rest.ensembl.org/sequence/id/%s?content-type=text/plain'%row['Protein stable ID']).content
    seqs.append(seq)
    seqlen.append(len(seq))
    refs.append(d.loc[d['NCBI gene ID']==row['NCBI gene ID']]['PMIDs'].values[0])


#     refs.append(hist_genes.loc(hist_genes['NCBI gene ID']==row['NCBI gene ID'])values[0])
histone_proteins['Protein sequence']=seqs
histone_proteins['Protein length']=seqlen
histone_proteins['References']=refs


#response = requests.get('http://rest.ensembl.org/sequence/id/ENST00000340857?content-type=text/plain')
#print (response.status_code)
#print (response.content)

# hp=(histone_proteins[histone_proteins['Transcript type'].isin(['protein_coding'])]).drop_duplicates(subset=['Transcript stable ID'])
hp=(histone_proteins[histone_proteins['Transcript type'].isin(['protein_coding',np.nan])]).sort_values('RefSeq peptide ID', ascending=True).drop_duplicates(subset=['Ensembl gene ID','Protein sequence']).sort_index()

hp.drop(columns=['Protein sequence','Transcript type','Protein stable ID']).to_csv('human_histone_proteins_autogenerated.csv',index=False)

!cp human_histone_proteins_autogenerated.csv docs/human_histones.csv

dataset.list_attributes().to_csv('atr.csv')