import os

import json

import socket

import struct

import argparse

from datetime import datetime

from biomedner_init import BioMedNER

from ops import filter_entities, pubtator2dict_list

def count_entities(data):

    num_entities = 0

    for d in data:

        if 'entities' not in d:

            continue

        for ent_type, entities in d['entities'].items():

            num_entities += len(entities)

    return num_entities

def biomedner_recognize(model, dict_path, base_name, args):

    input_mt_ner = os.path.join(args.biomedner_home, 'input',

                                f'{dict_path[2:]}.PubTator')

    output_mt_ner = os.path.join(args.biomedner_home, 'output',

                                f'{dict_path[2:]}.json')

    dict_list = pubtator2dict_list(input_mt_ner)

    res = model.recognize(

        input_dl=dict_list,

        base_name=base_name

    )

    if res is None:

        return None, 0

    num_filtered_species_per_doc = filter_entities(res)

    for n_f_spcs in num_filtered_species_per_doc:

        if n_f_spcs[1] > 0:

            print(datetime.now().strftime(args.time_format),

                  '[{}] Filtered {} species'

                  .format(base_name, n_f_spcs[1]))

    num_entities = count_entities(res)

    res[0]['num_entities'] = num_entities

    # Write output str to a .PubTator format file

    with open(output_mt_ner, 'w', encoding='utf-8') as f:

        json.dump(res[0], f)

def run_server(model, args):

    host = args.biomedner_host

    port = args.biomedner_port

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:

        s.bind((host, port))

        s.listen(1)

        while True:

            conn, addr = s.accept()

            dict_path = conn.recv(512).decode('utf-8')

            base_name = dict_path.split('.')[0]

            # hotfix

            base_name = base_name.replace("\x00A","")

            biomedner_recognize(model, dict_path, base_name, args)

            output_stream = struct.pack('>H', len(dict_path)) + dict_path.encode(

                'utf-8')

            conn.send(output_stream)

            conn.close()

if __name__ == "__main__":

    argparser = argparse.ArgumentParser()

    argparser.add_argument('--seed', type=int, help='random seed for initialization',

                            default=1)

    argparser.add_argument('--model_name_or_path', default='dmis-lab/bern2-ner')

    argparser.add_argument('--max_seq_length', type=int, help='The maximum total input sequence length after tokenization. Sequences longer than this will be truncated, sequences shorter will be padded.',

                            default=128)

    argparser.add_argument('--biomedner_home',

                           help='biomedical language model home')         

    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('--no_cuda', action="store_true", help="Avoid using CUDA when available")

    args = argparser.parse_args()

    mt_ner = BioMedNER(args)

    run_server(mt_ner, args)