import time
import torch
import string
import pdb
import argparse

from transformers import BertTokenizer, BertForMaskedLM
import BatchInference as bd
import batched_main_NER as ner
import aggregate_server_json as aggr
import json


DEFAULT_TOP_K = 20
SPECIFIC_TAG=":__entity__"
DEFAULT_MODEL_PATH="ajitrajasekharan/biomedical"
DEFAULT_RESULTS="results.txt"


def perform_inference(text,bio_model,ner_bio,aggr_ner):
    print("Getting predictions from BIO model...")
    bio_descs = bio_model.get_descriptors(text,None)
    print("Computing BIO results...")
    bio_ner = ner_bio.tag_sentence_service(text,bio_descs)
    obj = json.loads(bio_ner)
    combined_arr = [obj,obj]
    aggregate_results = aggr_ner.fetch_all(text,combined_arr)
    return aggregate_results
  
   
def process_input(results):
    try:
        input_file = results.input
        output_file = results.output
        print("Initializing BIO module...")
        bio_model = bd.BatchInference("bio/desc_a100_config.json",'ajitrajasekharan/biomedical',False,False,DEFAULT_TOP_K,True,True,       "bio/","bio/a100_labels.txt",False)
        ner_bio = ner.UnsupNER("bio/ner_a100_config.json")
        print("Initializing Aggregation module...")
        aggr_ner = aggr.AggregateNER("./ensemble_config.json")
        wfp = open(output_file,"w")
        with open(input_file) as fp:
            for line in fp:
                text_input = line.strip().split()
                print(text_input)
                text_input = [t + ":__entity__" for t in text_input]  
                text_input = ' '.join(text_input)
                start = time.time()
                results = perform_inference(text_input,bio_model,ner_bio,aggr_ner)
                print(f"prediction took {time.time() - start:.2f}s")
                pdb.set_trace()
                wfp.write(json.dumps(results))
                wfp.write("\n\n")
        wfp.close()
    except Exception as e:
        print("Some error occurred in batch processing") 
	
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Batch handling of NER ',formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-model', action="store", dest="model", default=DEFAULT_MODEL_PATH,help='BERT pretrained models, or custom model path')
    parser.add_argument('-input', action="store", dest="input", required=True,help='Input file with sentences')
    parser.add_argument('-output', action="store", dest="output", default=DEFAULT_RESULTS,help='Output file with sentences')
    results = parser.parse_args()
    process_input(results)