# TNRS

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 TNRS server

# Protocol params

url_base = 'http://tnrs.iplantcollaborative.org/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+

'||org.iplantc.tnrs.demo.client.SearchService|doSearch|\

java.lang.String/2004016611|{"sources":"gcc,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":"false", "source_sorting":"false", "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|')

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":"Best", "type":"Detailed", "encoding":"utf8", "dirty":"false", \

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

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

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

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

class InvalidResponse(Exception): pass

def gwt_encode(str_):

    return strings.esc_quotes(strings.json_encode(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

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

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

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

]

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()

def tnrs_request(names, debug=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.

    '''

    name_ct = len(names)

    assert name_ct <= max_names

    # Logging

    def debug_log(label, str_=''):

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

    ## HTTP

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

    def do_request(request):

        debug_log('request', str(request))

        response = urllib2.urlopen(urllib2.Request(url, request, headers))

        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 has

        # different lengths depending on the taxonomic ranks present in the

        # provided taxon name. The extra download step is therefore necessary.

        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')

        response = urllib2.urlopen(urllib2.Request(csv_url))

        debug_log('response info', str(response.info()))

        return TnrsOutputStream(response)

    finally:

        profiler.stop(name_ct)

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

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

    for try_num in xrange(2):

        try: return 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