import random
import requests
import os
import string
import numpy as np
import hashlib
import time
import shutil
import asyncio
import socket
import struct
import json
import sys
from datetime import datetime
from collections import OrderedDict
import traceback
import bioregistry

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

class RunBioMedNER():
    def __init__(self, 
        biomedner_home,
        biomedner_port,
        biomedner_host='localhost',
        time_format='[%d/%b/%Y %H:%M:%S.%f]',
        max_word_len=50, 
        seed=2019,
        keep_files=False,
        no_cuda=False):

        self.time_format = time_format

        print(datetime.now().strftime(self.time_format), 'BioMedNER LOADING..')
        random.seed(seed)
        np.random.seed(seed)

        if not os.path.exists('./output'):
            os.mkdir('output')

        # delete prev. version outputs
        if not keep_files:
            delete_files('./output')
            delete_files(os.path.join('./multi_ner', 'input'))
            delete_files(os.path.join('./multi_ner', 'tmp'))
            delete_files(os.path.join('./multi_ner', 'output'))

        self.biomedner_home = biomedner_home
        self.biomedner_host = biomedner_host
        self.biomedner_port = biomedner_port

        self.max_word_len = max_word_len

        print(datetime.now().strftime(self.time_format), 'BioMedNER LOADED..')
    
    def annotate_text(self, text, pmid=None):
        try:
            text = text.strip()
            base_name = self.generate_base_name(text) # for the name of temporary files
            text = self.preprocess_input(text, base_name)
            output = self.tag_entities(text, base_name)
            output['error_code'], output['error_message'] = 0, ""
        except Exception as e:
            errStr = traceback.format_exc()
            print(errStr)

            output = {"error_code": 1, "error_message": "Something went wrong. Try again."}

        return output


    def preprocess_input(self, text, base_name):
        if '\r\n' in text:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a CRLF -> replace it w/ a space')
            text = text.replace('\r\n', ' ')

        if '\n' in text:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a line break -> replace it w/ a space')
            text = text.replace('\n', ' ')

        if '\t' in text:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a tab -> replace w/ a space')
            text = text.replace('\t', ' ')

        if '\xa0' in text:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a \\xa0 -> replace w/ a space')
            text = text.replace('\xa0', ' ')

        if '\x0b' in text:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a \\x0b -> replace w/ a space')
            text = text.replace('\x0b', ' ')
            
        if '\x0c' in text:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a \\x0c -> replace w/ a space')
            text = text.replace('\x0c', ' ')
        
        # remove non-ascii
        text = text.encode("ascii", "ignore").decode()

        found_too_long_words = 0
        tokens = text.split(' ')
        for idx, tk in enumerate(tokens):
            if len(tk) > self.max_word_len:
                tokens[idx] = tk[:self.max_word_len]
                found_too_long_words += 1
        if found_too_long_words > 0:
            print(datetime.now().strftime(self.time_format),
                  f'[{base_name}] Found a too long word -> cut the suffix of the word')
            text = ' '.join(tokens)

        return text

    def tag_entities(self, text, base_name):
        n_ascii_letters = 0
        for l in text:
            if l not in string.ascii_letters:
                continue
            n_ascii_letters += 1

        if n_ascii_letters == 0:
            text = 'No ascii letters. Please enter your text in English.'

        base_name = self.generate_base_name(text)
        print(datetime.now().strftime(self.time_format),
              f'id: {base_name}')

        pubtator_file = f'{base_name}.PubTator'

        input_biomedner = os.path.join(self.biomedner_home, 'input',
                                     f'{pubtator_file}.PubTator')
        output_biomedner = os.path.join(self.biomedner_home, 'output',
                                     f'{pubtator_file}.json')

        if not os.path.exists(self.biomedner_home + '/input'):
            os.mkdir(self.biomedner_home + '/input')
        if not os.path.exists(self.biomedner_home + '/output'):
            os.mkdir(self.biomedner_home + '/output')

        # Write input str to a .PubTator format file
        with open(input_biomedner, 'w', encoding='utf-8') as f:
            # only abstract
            f.write(f'{base_name}|t|\n')
            f.write(f'{base_name}|a|{text}\n\n')

        ner_start_time = time.time()
        
        arguments_for_coroutines = []
        loop = asyncio.new_event_loop()
        for ner_type in ['biomedner']:
            arguments_for_coroutines.append([ner_type, pubtator_file, output_biomedner, base_name, loop])
        async_result = loop.run_until_complete(self.async_ner(arguments_for_coroutines))
        loop.close()
        biomedner_elapse_time = async_result['biomedner_elapse_time']

        # get output result to merge
        tagged_docs = async_result['tagged_docs']
        num_entities = async_result['num_entities']
        
        ner_elapse_time = time.time() - ner_start_time
        print(datetime.now().strftime(self.time_format),
              f'[{base_name}] ALL NER {ner_elapse_time} sec')

        # time record
        tagged_docs[0]['elapse_time'] = {
            'biomedner_elapse_time':biomedner_elapse_time,
            'ner_elapse_time': ner_elapse_time,
        } 

        # Delete temp files
        os.remove(input_biomedner)
        os.remove(output_biomedner)
        
        return tagged_docs[0]

    # generate id for temporary files
    def generate_base_name(self, text):
        # add time.time() to avoid collision
        base_name = hashlib.sha224((text+str(time.time())).encode('utf-8')).hexdigest()
        return base_name

    async def async_ner(self, arguments):
        coroutines = [self._ner_wrap(*arg) for arg in arguments]
        result = await asyncio.gather(*coroutines)
        result = {k:v for e in result for k,v in e.items()} # merge
        return result

    async def _ner_wrap(self, ner_type, pubtator_file, output_biomedner, base_name, loop):
        if ner_type == 'biomedner':            
            # Run neural model
            start_time = time.time()
            biomedner_resp = await async_tell_inputfile(self.biomedner_host,
                                         self.biomedner_port,
                                         pubtator_file,
                                         loop)
            
            with open(output_biomedner, 'r', encoding='utf-8') as f:
                tagged_docs = [json.load(f)]

            num_entities = tagged_docs[0]['num_entities']
            if tagged_docs is None:
                return None

            assert len(tagged_docs) == 1
            biomedner_elapse_time = time.time() - start_time
            print(datetime.now().strftime(self.time_format),
                f'[{base_name}] Multi-task NER {biomedner_elapse_time} sec, #entities: {num_entities}')

            return {"biomedner_elapse_time": biomedner_elapse_time,
                    "tagged_docs": tagged_docs,
                    "num_entities": num_entities}

async def async_tell_inputfile(host, port, inputfile, loop):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    try:
        sock.connect((host, port))
        input_str = inputfile
        input_stream = struct.pack('>H', len(input_str)) + input_str.encode(
            'utf-8')
        sock.send(input_stream)
        # output_stream = sock.recv(512)
        output_stream = await loop.run_in_executor(None, sock.recv, 512) # for async
        resp = output_stream.decode('utf-8')[2:]

        sock.close()
        return resp
    except ConnectionRefusedError as e:
        print(e)
        return None
    except TimeoutError as e:
        print(e)
        return None
    except ConnectionResetError as e:
        print(e)
        return None

def sync_tell_inputfile(host, port, inputfile):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    try:
        sock.connect((host, port))
        input_str = inputfile
        input_stream = struct.pack('>H', len(input_str)) + input_str.encode(
            'utf-8')
        sock.send(input_stream)
        output_stream = sock.recv(512) # for sync
        # output_stream = await loop.run_in_executor(None, sock.recv, 512)
        resp = output_stream.decode('utf-8')[2:]

        sock.close()
        return resp
    except ConnectionRefusedError as e:
        print(e)
        return None
    except TimeoutError as e:
        print(e)
        return None
    except ConnectionResetError as e:
        print(e)
        return None

def delete_files(dirname):
    if not os.path.exists(dirname):
        return

    n_deleted = 0
    for f in os.listdir(dirname):
        f_path = os.path.join(dirname, f)
        if not os.path.isfile(f_path):
            continue
        # print('Delete', f_path)
        os.remove(f_path)
        n_deleted += 1
    print(dirname, n_deleted)

if __name__ == '__main__':
    import argparse

    argparser = argparse.ArgumentParser()
    argparser.add_argument('--max_word_len', type=int, help='word max chars',
                           default=50)
    argparser.add_argument('--seed', type=int, help='seed value', default=2019)

    argparser.add_argument('--biomedner_home',
                           help='biomedical language model home',
                           default=os.path.join(os.path.expanduser('~'),
                                                'biomedner', 'biomednerHome'))
    argparser.add_argument('--biomedner_host',
                           help='biomedical language model host', default='localhost')
    argparser.add_argument('--biomedner_port', type=int, 
                           help='biomedical language model port', default=18894)
    argparser.add_argument('--time_format',
                           help='time format', default='[%d/%b/%Y %H:%M:%S.%f]')
    argparser.add_argument("--keep_files", action="store_true")
    argparser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
    args = argparser.parse_args()
    biomedner = RunBioMedNER(
        max_word_len=args.max_word_len,
        seed=args.seed,
        biomedner_home=args.biomedner_home,
        biomedner_host=args.biomedner_host,
        biomedner_port=args.biomedner_port,
        time_format=args.time_format,
        keep_files=args.keep_files,
        no_cuda=args.no_cuda,
    )
    text = "KRAS is a proto-oncogene involved in various cancers."
    result = biomedner.annotate_text(text.lower())
    print(result)