'''
Copyleft Jul 05, 2016 Arya Iranmehr, PhD Student, Bafna Lab, UC San Diego,  Email: airanmehr@gmail.com
'''
import numpy as np;

np.set_printoptions(linewidth=200, precision=5, suppress=True)
import pandas as pd;

pd.options.display.max_rows = 20;
pd.options.display.expand_frame_repr = False
import os;
import multiprocessing

home = os.path.expanduser('~') + '/'
import Utils.Util as utl
from Bio import Entrez

# Entrez.email = "Your.Name.Here@example.org"
# handle = Entrez.einfo() # or esearch, efetch, ...
# record = Entrez.read(handle)
# handle.close()
# print pd.Series(list(record['DbList'])).sort_values()
a = pd.read_csv(utl.outpath + 'real/SNP.csv', sep='\t').set_index(['CHROM', 'POS'])


def saveBatch(idx):
    try:
        print idx[-1]
        path = utl.outpath + 'dbsnp'
        Entrez.email = "Your.Name.Here@example.org"
        if not os.path.exists(path): os.mkdir(path)
        handle = Entrez.efetch("snp", id=','.join(a.iloc[idx].ID.apply(lambda x: x.replace('rs', '')).values),
                               retmode='xml')
        with open(path + '/batch_{:.0f}.xml'.format(int(idx[-1])), 'w') as f:
            print >> f, ''.join(handle.readlines())
    except:
        print 'error:', idx[-1]


def savedbSNP():
    batchSize = 10000
    print 'Num Batches=', a.shape[0] / batchSize + 1
    multiprocessing.Pool(6).map(saveBatch, np.array_split(range(a.shape[0]), a.shape[0] / batchSize + 1))


savedbSNP()


# field='ancestralAllele'
#
# e = xml.etree.ElementTree.parse('/home/arya/a').getroot()
#
# for x in e._children:
#     try:
#         print x.attrib['rsId'],pd.Series(x.find('{http://www.ncbi.nlm.nih.gov/SNP/docsum}Sequence').attrib[field].split(',')).unique()[0]
#     except:
#         pass
#
#
#
#
# from Bio.Entrez.Parser import expat
# records=Entrez.parse(handle,validate=False)
# for record in records:
#     print record
#     break
# c.tag.name
# len(records)
#
#
# def f(x):
#     print x
#     try:
#         records = Entrez.efetch("snp",  id=x.replace('rs',''),retmode='txt').readlines();
#         for record in records:
#             if  field in record: print record
#         return [pd.Series(record.split(field)[1].replace('"','').strip().split(',')).unique()[0]  for record in records if  field in record][0]
#     except:
#         print 'NA'
#         return None
# a.ID.iloc[:10].apply(lambda x: f(x))
#
