"""
Retrieve MedDRA:UMLS CUI's and COSTART:UMLS CUI's for each
adverse event term in the AERS database.
"""

import os
import csv
import sys
import urllib
import xml.dom.minidom
from xml.dom.minidom import Node

NCBO_ID_COSTART = 1341
NCBO_ID_MEDDRA = 1422

class Bean(object):
    """A Ontology Bean"""
    def __init__(self):
        super(Bean, self).__init__()
        self.id = ""
        self.label = ""
        self.umls_ids = []
    
    def __repr__(self):
        return """
Id: %s
Label: %s
UMLS IDs: %s
""" % (self.id, self.label, ", ".join(self.umls_ids))

class SearchBean(object):
    """A SearchBean"""
    def __init__(self):
        super(SearchBean, self).__init__()
        self.record_type = ""
        self.concept_id = ""
        self.name = ""
    
    def __repr__(self):
        return """
Name: %s
RecordType: %s
ConceptId: %s
""" % (self.name, self.record_type, self.concept_id)

class NCBORest(object):
    """
    A narrow class for accessing the NCBO Rest Services.
    """
    
    search_url = "http://rest.bioontology.org/bioportal/search/%s?ontologyids=%d&includeproperties=1&email=nick.tatonetti@stanford.edu"
    retrieve_url = "http://rest.bioontology.org/bioportal/virtual/ontology/%d/%s?email=nick.tatonetti@stanford.edu"
    
    @staticmethod
    def search(term, ontology_id):
        """ Search for $term in $ontology_id """
        xml_url = NCBORest.search_url % (urllib.quote(term), ontology_id)
        # print >> sys.stderr, xml_url
        raw_xml = urllib.urlopen(xml_url).read()
        doc = xml.dom.minidom.parseString(raw_xml)
        beans = []
        for xml_bean in doc.getElementsByTagName("searchBean"):
            
            bean = SearchBean()
            bean.record_type = xml_bean.getElementsByTagName('recordType')[0].childNodes[0].data
            bean.concept_id = xml_bean.getElementsByTagName('conceptId')[0].childNodes[0].data
            bean.name = xml_bean.getElementsByTagName('preferredName')[0].childNodes[0].data
            beans.append(bean)
        
        return beans
    
    @staticmethod
    def retrieve(concept_id, ontology_id):
        """ Retrieve the bean for $concept_id in $ontology_id """
        
        xml_url = NCBORest.retrieve_url % (ontology_id, concept_id)
        # print >> sys.stderr, xml_url
        raw_xml = urllib.urlopen(xml_url).read()
        doc = xml.dom.minidom.parseString(raw_xml)
        class_beans = doc.getElementsByTagName("classBean")
        if len(class_beans) == 0:
            # print >> sys.stderr, "Failed to find class beans at %s" % xml_url
            return None
        
        class_bean = class_beans[0]
        
        bean = Bean()
        bean.id = class_bean.getElementsByTagName('id')[0].childNodes[0].data
        bean.label = class_bean.getElementsByTagName('label')[0].childNodes[0].data
        
        relations = class_bean.getElementsByTagName('relations')[0]
        for entry in relations.getElementsByTagName('entry'):
            if entry.getElementsByTagName('string')[0].childNodes[0].data == u'UMLS_CUI':
                cui_list = entry.getElementsByTagName('list')[0]
                for cui_xml in cui_list.getElementsByTagName('string'):
                    umls_id = cui_xml.childNodes[0].data
                    bean.umls_ids.append( umls_id )
        
        return bean

if __name__ == '__main__':
    
    import MySQLdb
    # 
    db = MySQLdb.connect(host="localhost", port=3307, user="root", passwd="enter_your_password",db="effect_aers")
    c = db.cursor()
    query = """
    select distinct event_descrip_term
    from reactions;
    """
    c.execute(query)

    all_terms = set([row[0] for row in c.fetchall()])
    
    query = """
    select distinct event_descrip_term
    from aers2umls;
    """
    c.execute(query)
    
    mapped_terms = set([row[0] for row in c.fetchall()])
    
    terms = list(all_terms - mapped_terms)
    # terms = ['BLOOD GLUCOSE DECREASED'] #, 'BLOOD GLUCOSE INCREASED', 'HYPERGLYCAEMIA', 'HYPOGLYCAEMIA']
    
    query = """
    select distinct umls_id
    from aers2umls;
    """
    c.execute(query)
    
    umls_ids = [row[0] for row in c.fetchall()]
    
    completed_terms = []
    if os.path.exists('aers_medra_umls.txt'):
        completed_terms.extend([r.split('\t')[0] for r in open('aers_medra_umls.txt')])
    
    if os.path.exists('aers_costart_umls.txt'):
        completed_terms.extend([r.split('\t')[0] for r in open('aers_costart_umls.txt')])
    
    out_meddra = open('aers_meddra_umls.txt','a')
    out_costart = open('aers_costart_umls.txt','a')
    
    terms = list(set(terms) - set(completed_terms))
    
    total_terms = len(terms)
    for i,term in enumerate(terms):
        print >> sys.stderr, "%d of %d: %50s:\t" % (i, total_terms, term), ; sys.stderr.flush()
        mapping = {"MEDDRA_UMLS_IDS":set([])}
        
        found_map = None
        stage = 0
        
        for meddra_search_bean in NCBORest.search(term, NCBO_ID_MEDDRA):
            medra_bean = NCBORest.retrieve(meddra_search_bean.concept_id, NCBO_ID_MEDDRA)
            if stage == 0:
                print >> sys.stderr, ".",; sys.stderr.flush()
                stage = 1
            
            if not medra_bean is None:
                for medra_umls_id in medra_bean.umls_ids:
                    mapping['MEDDRA_UMLS_IDS'].add(medra_umls_id)
                    for costart_search_bean in NCBORest.search(medra_umls_id, NCBO_ID_COSTART):
                        costart_bean = NCBORest.retrieve(costart_search_bean.concept_id, NCBO_ID_COSTART)
                        if not costart_bean is None:
                            if stage == 1:
                                print >> sys.stderr, ".",; sys.stderr.flush()
                                stage = 2
                            
                            for costart_umls_id in costart_bean.umls_ids:
                                if costart_umls_id in umls_ids:
                                    print >> sys.stderr, "*"; sys.stderr.flush()
                                    found_map = costart_umls_id
                                    break
                        
                        if not found_map is None:
                            break
                    
                    if not found_map is None:
                        break
            
            if not found_map is None:
                break
                                    
        
        # for umls_id in mapping['COSTART_UMLS_IDS']:
        if not found_map is None:
            print >> out_costart, "%s\t%s" % (term, found_map)
        
        if found_map is None and len(mapping['MEDDRA_UMLS_IDS']) > 0:
            print >> sys.stderr, "-"; sys.stderr.flush()
            umls_id = mapping['MEDDRA_UMLS_IDS'].pop()
            print >> out_meddra, "%s\t%s" % (term, umls_id)
        
        out_meddra.flush()
        out_costart.flush()