################################################################################
# Classes to use the UMLS::Similarity web interface, these were adapted from   #
# the perl module at https://metacpan.org/pod/distribution/UMLS-Similarity/    #
# utils/query-umls-similarity-webinterface.pl                                  #
################################################################################
from __future__ import print_function
import urllib
from bs4 import BeautifulSoup
import re
import time
import sys


class UmlsSimilarityWeb(object):
    '''
    Get similarity and relatedness scores between terms or concepts CUIs from 
    the remote UMLS similarity webserver at:
    http://atlas.ahc.umn.edu/cgi-bin/umls_similarity.cgi
    server=The server to query (as far as I know there is only one
    sleep=the sleep time after the query (the default is the same time as in the
          perl script (2 seconds)
    '''

    ### CLASS wide variables, not static but it would be silly to edit them ###
    # Allowed permutations for the relatedness SAB:REL
    allowed_rel_sabrel={'SNOMEDCT': ['CUI/PAR/CHD/RB/RN','CUI'],
                        'MSH':['CUI/PAR/CHD/RB/RN','CUI'],
                        'UMLS_ALL': ['CUI/PAR/CHD/RB/RN','CUI']}

    # Allowed permutations for the similarity SAB:REl
    allowed_sim_sabrel={'SNOMEDCT': ['PAR/CHD','RB/RN'],
                        'MSH': ['PAR/CHD','RB/RN'],
                        'FMA': ['PAR/CHD','RB/RN'],
                        'OMIM': ['PAR/CHD','RB/RN']}

    # Allowed algorithm abbreviations for the similarity measure
    allowed_sim_method={'all': 'All Similarity Methods',
                        'cdist': 'Conceptual Distance',
                        'jcn': 'Jiang & Conrath',
                        'lch': 'Leacock & Chodorow',
                        'lin': 'Lin',
                        'nam': 'Nguyen & Al-Mubaid',
                        'path': 'Path Length',
                        'random': 'Random Measure',
                        'wup': 'Wu & Palmer',
                        'res': 'Resnik'}

    # Allowed algorithm abbreviations for the relatedness measure
    allowed_rel_method={'all': 'All Relatedness Methods',
                        'lesk': 'Lesk',
                        'vector': 'Vector Measure'}

    # The defaults for the similarity and relatedness get methods
    sim_defaults={'sab':'MSH','rel':allowed_sim_sabrel['MSH'][0],'similarity':'path'}
    rel_defaults={'sab':'UMLS_ALL','rel':allowed_rel_sabrel['MSH'][0],'relatedness':'vector'}

    # The keys in the results
    results_keys=['MEASURE','TERM1','CUI1','TERM2','CUI2','METHOD','SCORE']

    def __init__(self,
                 server='http://atlas.ahc.umn.edu',
                 cgi='cgi-bin/umls_similarity.cgi',
                 sleep=2,
                 verbose=False):
        self.server=server
        self.cgi=cgi
        self.sleep=sleep
        self.verbose=verbose
        self.results_re=re.compile(r'The\s*(?P<MEASURE>similarity|relatedness)\s*of\s*(?P<TERM1>.+)\s*\((?P<CUI1>C\d{7})\s*\)\s*and\s*(?P<TERM2>.+)\s*\((?P<CUI2>C\d{7})\s*\)\s*using.+\((?P<METHOD>\w+)\)\s+is\s+(?P<SCORE>-?\d+(?:\.\d+)?)\.',re.IGNORECASE)

    def version(self):
        '''
        Return the UMLS,SIMLS::Similarity and UMLS::Interface verion
        from the server
        '''
        
        # Build the URL to check the version number
        data = {'version':'yes'}
        url_value = urllib.urlencode(data)
        full_url = self.server+'/'+self.cgi+'?'+url_value

        # Open the URL and get the data, I relly should have some
        # error handling in here
        data = urllib.urlopen(full_url)
        soup = BeautifulSoup(data, 'html.parser')

        # Some regular expressions for the version numbers
        interface_re=re.compile(r'UMLS::Interface version (?P<VERSION>\d+(?:\.\d+))?')
        similarity_re=re.compile(r'UMLS::Similarity version (?P<VERSION>\d+(?:\.\d+))?')
        umls_re=re.compile(r'UMLS version (?P<VERSION>.*)')

        # Lables to associate with the regular expressions
        labels=['umls_interface','umls_similarity','umls']
        
        versions={} # Will hold the verion information

        # Loop through the regular expressions and labels
        for c,r in enumerate([interface_re,similarity_re,umls_re]):
            # The default version label
            versions[labels[c]]='UNKNOWN'

            # Now go through the p tags
            for i in soup.find_all('p'):
                try:
                    # If we find the match then move to the next regexp
                    found=r.match(i.string)
                    versions[labels[c]]=found.group('VERSION')
                    break
                except TypeError:
                    pass
                except AttributeError:
                    pass
        
        return versions

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_similarity(self,
                       term1,
                       term2,
                       sab=sim_defaults['sab'],
                       rel=sim_defaults['rel'],
                       method=sim_defaults['similarity']):
        '''
        Return the similarity measure between 2 terms
        sab=The sources used to determine the similarity 
        rel=The relationship types to determine the similarity
        measure=The method or algorithm to determine similarity
        '''
        self._check_allowed(self.allowed_sim_sabrel,self.allowed_sim_method,sab,rel,method)

        # Build the URL to determine the similarity
        data = {'word1':term1,'word2':term2,'sab':sab,'rel':rel,
                'similarity':method,'button':'Compute Similarity',
                'sabdef':self.rel_defaults['sab'],'reldef':self.rel_defaults['rel'],
                'relatedness':self.rel_defaults['relatedness']}
        html=self._do_query(data)
        return self._extract_results(html)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_relatedness(self,
                        term1,
                        term2,
                        sab=rel_defaults['sab'],
                        rel=rel_defaults['rel'],
                        method=rel_defaults['relatedness']):
        '''
        Return the relatedness measure between 2 terms
        sab=The sources used to determine the similarity 
        rel=The relationship types to determine the similarity
        measure=The method or algorithm to determine similarity
        '''
        self._check_allowed(self.allowed_rel_sabrel,self.allowed_rel_method,sab,rel,method)

        # Build the URL to determine the similarity
        data = {'word1':term1,'word2':term2,'sab':self.sim_defaults['sab'],
                'rel':self.sim_defaults['rel'],'similarity':self.sim_defaults['similarity'],
                'button':'Compute Relatedness','sabdef':sab,'reldef':rel,
                'relatedness':method}
        html=self._do_query(data)
        return self._extract_results(html)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _extract_results(self,html):
        '''
        Extract the results data from the html output from UMLS similarity
        '''
        soup=BeautifulSoup(html,'html.parser')
        all_results=[]

        #It looks like all the results are in p elements with the class results
        for i in soup.find_all('p',"results"):
            result={}
            matches=self.results_re.match(i.get_text())
            
            # Loop through all the results fields and record the data for each one
            # It is an error if we can't find one
            for r in self.results_keys:
                try:
                    result[r]=matches.group(r).strip()
                except AttributeError:
                    print(i)
                    print(i.get_text())
                    print("[error] unable to extract results for %s maybe the structure of the web page has changed?" % r,file=sys.stderr)
                    raise
            
            # Make the score a float
            result['SCORE']=float(result['SCORE'])

            # Store all the results
            all_results.append(result)

            # Have a quick sleep before retuning results. This just makes sure that
            # the server does not get hammered
            time.sleep(self.sleep)
        return all_results

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _do_query(self,data):
        '''
        when provided with the components of the url as a dict this builds the 
        url and executes the query before returning the html
        '''
        # Encode the data and build the URL
        url_value=urllib.urlencode(data)
        full_url=self.server+'/'+self.cgi+'?'+url_value

        if self.verbose:
            print("Executing: %s" % full_url)

        # Open the URL and get the data, I relly should have some
        # error handling in here
        data = urllib.urlopen(full_url)
        return data

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _check_allowed(self,allowed_sabrel,allowed_method,sab,rel,method):
        '''
        Determines if the passed arguments are allowed and raises errors if
        they are not
        '''
        try:
            # Check that the supplied sab and rel are allowed
            checked=allowed_sabrel[sab.upper()].index(rel)
        except KeyError:
            print('[error] sab %s is not allowed expecting either of these %s'
                  % (sab," | ".join(allowed_sabrel.keys())),file=sys.stderr)
            raise
        except IndexError:
            print('[error] rel %s is not allowed expecting either of these %s'
                  % (rel," | ".join(allowed_sabrel[sab.upper()])),file=sys.stderr)
            raise

        # Now check if the measure is allowed
        try:
            checked=allowed_method[method.lower()]
        except KeyError:
            print('[error] method %s is not allowed expecting either of these %s'
                  % (method," | ".join(allowed_method.keys())),file=sys.stderr)
            raise
            

            
