import requests

import logging

logger = logging.getLogger(__name__)

from nltk.stem import WordNetLemmatizer

lemmatizer = WordNetLemmatizer()

def solr_escape(query: str) -> str:

    # These special Solr characters need to be escaped. We deal with some of them

    # "+ - && || ! ( ) { } [ ] ^ " ~ * ? : \ /"

    char_translation_table = str.maketrans({

        '[': '',

        ']': '',

        '{': '',

        '}': '',

        '^': '',

        '~': '',

        '*': '',

        '\\': '',

        '/': '',

        '"': '',

        '!': '\!',

        '?': '\?',

        ':': '\:',

    })

    return query.translate(char_translation_table)

def get_solr_response_generic(solr_formatted_query: str, solr_core_name: str):

    solr_response_json = requests.get(

        url ='http://tangra.cs.yale.edu:8983/solr/{}/select?'.format(solr_core_name),

        params={

            'indent': 'on',

            'q': solr_formatted_query,

            'rows': '100',

            'wt': 'json',

            'fl': 'id, score'

        }

    ).json()

    solr_response = solr_response_json['response']

    num_rows = solr_response['numFound']

    rows = solr_response['docs']

    solr_matched_ids = []

    id_to_score_map = {}

    for row in rows:

        item_id = int(row['id'])

        item_solr_score = float(row['score'])

        solr_matched_ids.append(item_id)

        id_to_score_map[item_id] = item_solr_score

    return solr_matched_ids, id_to_score_map

def get_solr_response_mimic(raw_query):

    escaped_query = solr_escape(raw_query)

    query_words = ['"' + lemmatizer.lemmatize(word) + '"' for word in escaped_query.split()]

    if len(query_words) > 0:

            query = ' AND '.join(query_words)

            solr_formatted_query = 'abbreviations:({}) OR abbreviations_sent_id:({}) OR doctext_not_stored:({})'.format(

                query, query, query)

    else:

        solr_formatted_query = '*:*'

    return get_solr_response_generic(solr_formatted_query, 'ehr_abbs_mimic')

def get_solr_response_umn_wrap(raw_query):

    escaped_query = solr_escape(raw_query)

    query_words = ['"' + lemmatizer.lemmatize(word) + '"' for word in escaped_query.split()]

    if len(query_words) > 0:

            query = ' AND '.join(query_words)

            solr_formatted_query = 'short_form:({})'.format(

                query)

    else:

        solr_formatted_query = '*:*'

    return get_solr_response_umn(solr_formatted_query, 'ehr_abbsense_umn')

def get_solr_response_umn(solr_formatted_query: str, solr_core_name: str):

    solr_response_json = requests.get(

        url ='http://tangra.cs.yale.edu:8983/solr/{}/select?'.format(solr_core_name),

        params={

            'indent': 'on',

            'q': solr_formatted_query,

            'rows': '100',

            'wt': 'json',

            'fl': 'id, long_form, score'

        }

    ).json()

    solr_response = solr_response_json['response']

    num_rows = solr_response['numFound']

    rows = solr_response['docs']

    solr_matched_longforms = []

    long_form_to_score_map = {}

    for row in rows:

        item = row['long_form']

        item_solr_score = float(row['score'])

        solr_matched_longforms.append(item)

        long_form_to_score_map[item] = item_solr_score

    return solr_matched_longforms, long_form_to_score_map