# TNRS

import json

import os.path

import re

import sys

import time

import urllib2

import csvs

import exc

import profiling

import streams

import strings

# Config

initial_pause = 0.35 # sec

pause_growth_factor = 1.3

max_pause = 30*60 # sec; = 30 min; TNRS sometimes freezes for ~10 min

assert initial_pause <= max_pause

#max_names = 5000 #according to http://tnrs.iplantcollaborative.org/TNRSapp.html

max_names = 500 # the maximum above crashes the live and dev TNRS servers

# Protocol params

server = streams.file_get_contents(os.path.dirname(__file__)+"/tnrs.url")

#server = 'tnrs.iplantcollaborative.org' # live server

url_base = 'http://'+server+'/tnrsdemo/'

url = url_base+'search'

initial_headers = {

    'Content-Type': 'text/x-gwt-rpc; charset=utf-8',

    'X-GWT-Module-Base': url_base,

    'X-GWT-Permutation': '574AA16D15D917C7704646FD92AFF6B3',

}

submission_request_template = ('7|0|7|'+url_base+

'|1E87C78041CEFBF0992F46BDF84D7D60|org.iplantc.tnrs.demo.client.SearchService|\

doSearch|java.lang.String/2004016611|{"sources":"gcc,tpl,tropicos,usda", \

"names":"[names]", "type":"matching", "taxonomic":"true", \

"classification":"tropicos", "match_to_rank":"true"}|0.05|1|2|3|4|2|5|5|6|7|')

submission_response_pattern = r'^//OK\[1,\["(\w+)"\],0,7\]$'

retrieval_request_template = ('7|0|15|'+url_base+

'|1E87C78041CEFBF0992F46BDF84D7D60|org.iplantc.tnrs.demo.client.SearchService\

|getRemoteData|com.extjs.gxt.ui.client.data.PagingLoadConfig|\

java.lang.String/2004016611|com.extjs.gxt.ui.client.data.BasePagingLoadConfig/\

2011366567|com.extjs.gxt.ui.client.data.RpcMap/3441186752|sortField|sortDir|\

com.extjs.gxt.ui.client.Style$SortDir/640452531|offset|java.lang.Integer/\

3438268394|limit|{"email":"tnrs@lka5jjs.orv", "key":"[key]", \

"taxonomic_constraint":"true", "source_sorting":"true", "first":"false"}\

|1|2|3|4|2|5|6|7|0|1|8|4|9|0|10|11|0|12|13|0|14|13|100|15|')

	# taxonomic_constraint (Constrain by Higher Taxonomy): selects lower ranks

	#  only from within the matched higher ranks. must be turned on, to ensure

	#  that higher ranks are always authoritative.

	# source_sorting (Constrain by Source): always puts matches in the order of

	#  the sources, regardless of match score. should never be turned on,

	#  because it selects worse matches instead of better ones. *however*, since

	#  this is currently broken and always forced on, we turn it on so that the

	#  download settings reflect what TNRS actually used.

retrieval_response_pattern = '^//OK\[.*?\["com.extjs.gxt.ui.client.data.\

BasePagingLoadResult/496878394","java.util.ArrayList/4159755760","org.iplantc.\

tnrs.demo.shared.BeanTNRSEntry/1039545748",".*"\],0,7\]$'

retrieval_response_info_pattern = r'(?ms).*^Set-Cookie: JSESSIONID=(\w+);'

download_request_template = ('7|0|6|'+url_base+

'|1E87C78041CEFBF0992F46BDF84D7D60|org.iplantc.tnrs.demo.client.SearchService|\

downloadRemoteResults|java.lang.String/2004016611|{"name":"tnrs_results.txt", \

"mode":"All", "type":"Detailed", "encoding":"utf8", "dirty":"false", \

"sources":"false", "taxonomic":"true", "email":"tnrs@lka5jjs.orv", \

"key":"[key]"}|1|2|3|4|1|5|6|')

	# dirty:

	#  when on, sometimes marks multiple names as Selected. must be turned *off*

	#  so that only one name is marked as Selected. note that this actually used

	#  to be on in the web app (see r9910, 2013-6-18), but does not appear to be

	#  needed (the source_sorting bug alluded to in r9910 is not fixed by

	#  enabling the dirty setting).

download_response_pattern = '^//OK\[1,\["(.*)"\],0,7\]$'

download_url_suffix = '&name=tnrs_results.txt&encoding=utf8'

tsv_to_json_col_map = {

	'Name_matched_author': 'Canonical_author',

	'Accepted_name_species': 'Accepted_species',

	'Accepted_name_family': 'Accepted_family',

	'Selected': 'selected',

}

class InvalidResponse(Exception): pass

def gwt_encode(str_):

    # embedded strs within GWT sections additionally have " JSON-escaped

    return strings.esc_quotes(strings.json_encode(str_), '|', quote_esc='\!')

def gwt_decode(str_):

    # embedded strs within GWT sections additionally have " JSON-escaped

    return strings.json_decode(strings.unesc_quotes(str_, '|', quote_esc='\!'))

def make_spliced_decode_map(decode_map):

    return [(r'(?: |(?<=\t)|^)'+re.escape(from_.strip())+r'(?: |(?=\t)|$)',

        strings.regexp_repl_esc(to)) for from_, to in decode_map]

padding = ' !pad ' # prepend to empty and whitespace-only strings

encode_map = [

    ('!', ' !exc '), # our escape char

    ('\t', ' !tab '), # TNRS replaces with " "

    ('\n', ' !nl '), # used to separate multiple names

    ('\r', ' !cr '), # used to separate multiple names

    ('"', ' !quo '), # TNRS removes it when at the beginning or end

    ('%', ' !pct '), # TNRS URL-decodes it in matched fields

    ("'", ' !apo '), # TNRS removes it when at the beginning or end

    ('+', ' !plus '), # TNRS removes it from the morphospecies

    (';', ' !sem '), # changes TNRS response format

    ('\\', ' !bsl '), # TNRS removes it

    ('_', ' !und '), # TNRS replaces with " "

    ('­', ' !sub '), # TNRS removes it

    ('×', ' !mul '), # TNRS replaces with "x"

]

decode_map = strings.flip_map(encode_map)

decode_map.append((padding, ''))

spliced_decode_map = make_spliced_decode_map(decode_map)

def encode(str_):

    str_ = strings.replace_all(encode_map, str_)

    # Empty and whitespace-only strings are ignored by TNRS (no response row)

    if str_.strip() == '': str_ = padding+str_

    return str_

def decode(str_): return strings.replace_all_re(spliced_decode_map, str_)

decode_for_tsv_map = make_spliced_decode_map([(from_, strings.replace_all(

    csvs.tsv_encode_map, to)) for from_, to in decode_map])

def decode_for_tsv(str_):

    return strings.replace_all_re(decode_for_tsv_map, str_)

class TnrsOutputStream(streams.FilterStream):

    '''Decodes a TNRS response whose names were encoded with encode()'''

    def __init__(self, stream):

        streams.FilterStream.__init__(self, decode_for_tsv, stream)

def parse_response(name, pattern, str_, response, response_info):

    match = re.match(pattern, str_)

    if not match:

        raise InvalidResponse('Invalid '+name+' response:\n'+response_info+'\n'

            +response)

    return match.groups()

export_cols = None

def single_tnrs_request(names, debug=False, cumulative_profiler=None,

    use_tnrs_export=False):

    '''

    Note that names containing only whitespace characters (after gwt_encode())

    are ignored by TNRS and do not receive a response row. Thus, you should

    always match up the Name_submitted returned by TNRS with the actual

    submitted name to determine the corresponding TNRS response row.

    @param use_tnrs_export whether to use TNRS's TSV export feature, which

        currently returns incorrect selected matches (vegpath.org/issues/943)

    '''

    global export_cols

    name_ct = len(names)

    assert name_ct <= max_names

    assert name_ct >= 1 # with no names, TNRS hangs indefinitely

    # Logging

    def debug_log(label, str_=''):

        if debug: sys.stderr.write('\n'+label+':\n'+str_+'\n')

    if not use_tnrs_export and export_cols == None: # need export_cols

        debug_log('fetching export columns')

        export_cols = csvs.header(single_tnrs_request([''], debug,

            use_tnrs_export=True)) # need at least 1 name

        debug_log('TSV export columns', strings.urepr(export_cols))

        json_cols = [tsv_to_json_col_map.get(c, c) for c in export_cols]

        debug_log('JSON columns', strings.urepr(json_cols))

    ## HTTP

    headers = initial_headers.copy() # don't modify global constant!

    def do_request(request):

        request_obj = urllib2.Request(url, request, headers)

        debug_log('request URL', str(url))

        debug_log('request info', str(request_obj.header_items()))

        debug_log('request str', str(request_obj.get_data()))

        response = urllib2.urlopen(request_obj)

        response_str = streams.read_all(response)

        response_info = str(response.info())

        debug_log('response info', response_info)

        debug_log('response str', response_str)

        return response_str, response_info

    def do_repeated_request(request):

        pause = initial_pause

        total_pause = 0

        while True:

            total_pause += pause

            if total_pause > max_pause: raise # error is not temporary

            debug_log('total_pause', str(total_pause)+'s')

            time.sleep(pause) # wait for job to complete

            try: return do_request(request)

            except urllib2.HTTPError: pass # try again

            pause *= pause_growth_factor

    profiler = profiling.ItersProfiler(start_now=True, iter_text='name')

    try:

        debug_log('Submit')

        request = submission_request_template.replace('[names]',

            gwt_encode('\n'.join(map(encode, names))))

        response, response_info = do_request(request)

        key, = parse_response('submission', submission_response_pattern,

            response, response, response_info)

        debug_log('key', key)

        key_enc = gwt_encode(key)

        debug_log('Retrieve')

        request = retrieval_request_template.replace('[key]', key_enc)

        response, response_info = do_repeated_request(request)

        parse_response('retrieval', retrieval_response_pattern, response,

            response, response_info)

        session_id, = parse_response('retrieval info',

            retrieval_response_info_pattern, response_info, response,

            response_info)

        debug_log('session_id', session_id)

        headers['Cookie'] = 'JSESSIONID='+session_id

        # the output of the retrieve step is unusable because the array does not

        # contain all the columns, contains no column names, and has different

        # lengths depending on the taxonomic ranks present in the provided taxon

        # name. the extra download step is therefore necessary.

        if use_tnrs_export:

            debug_log('Prepare download')

            request = download_request_template.replace('[key]', key_enc)

            response, response_info = do_request(request)

            csv_url, = parse_response('download', download_response_pattern,

                response, response, response_info)

            csv_url += download_url_suffix

            debug_log('csv_url', csv_url)

            debug_log('Download')

            request_obj = urllib2.Request(csv_url)

            debug_log('request URL', str(csv_url))

            debug_log('request info', str(request_obj.header_items()))

            debug_log('request str', str(request_obj.get_data()))

            response = urllib2.urlopen(request_obj)

            response_info = str(response.info())

            debug_log('response info', response_info)

            output_stream = response

        else:

            debug_log('Expand')

            request = ('7|0|6|http://tnrs.iplantcollaborative.org/tnrsdemo/|\

1E87C78041CEFBF0992F46BDF84D7D60|org.iplantc.tnrs.demo.client.SearchService|\

requestGroupMembers|java.lang.String/2004016611|{"group":"0", \

"email":"tnrs@lka5jjs.orv", "key":"[key]", "source_sorting":"false", \

"taxonomic_constraint":"false"}|1|2|3|4|1|5|6|'.replace('[key]', key_enc))

            response, response_info = do_request(request)

            data_gwt, = parse_response('expand', r'^//OK\[1,\["(.*)"\],0,7\]$',

                response, response, response_info)

            data_json = json.loads(gwt_decode(data_gwt))

            debug_log('JSON data', strings.urepr(data_json))

            output_stream = csvs.InputRewriter(

                csvs.RowNumFilter(

                    csvs.JsonReader(data_json['items'], json_cols[1:], ';'),

                    json_cols[0]

                )

            )

        return TnrsOutputStream(output_stream)

    finally:

        profiler.stop(name_ct)

        sys.stderr.write(profiler.msg()+'\n')

        if cumulative_profiler != None:

            cumulative_profiler.add_subprofiler(profiler)

            sys.stderr.write('Cumulatively: '+cumulative_profiler.msg()+'\n')

def tnrs_request(names, debug=False, **kw_args):

    for try_num in xrange(2):

        try: return single_tnrs_request(names, debug, **kw_args)

        except (urllib2.HTTPError, InvalidResponse), e:

            exc.print_ex(e, detail=False)

            debug = True

            # try again with debug turned on

    raise # error is not temporary