# -*- coding: utf-8 -*-
from __future__ import unicode_literals, print_function
import os
import sys
import pandas as pd
import json
reload(sys);
sys.setdefaultencoding("utf8")
from bson.son import SON
from bson.codec_options import CodecOptions
from os import path
# nlputis codes.
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from protolib.python import document_pb2, rpc_pb2, edgRules_pb2
from utils.rpc.iterator import request_iter_docs, edg_request_iter_docs
from utils.rpc import grpcapi
from utils.helper import DocHelper
from utils.param_helper import ParamHelper
from utils.edg_relations import EdgArg, EdgRelation, EdgRelations
import glob
from collections import defaultdict
import re
from pymongo import MongoClient

'''
columns = ("doc_id", "sent_index", "token_index", "Token","base_np", "base_np_offset", \
           "full_np", "full_np_offset")
'''
class np_generator_for_abstrat:
    def __init__(self,pmid,edgServer,edgServerPort):
        self.pmid_list=pmid
        self.edgServer=edgServer
        self.edgServerPort=edgServerPort


    def get_offset(self,proto_obj, doc):
        if type(proto_obj) == document_pb2.Sentence.Constituent:
            token_start = doc.token[proto_obj.token_start]
            token_end = doc.token[proto_obj.token_end]
            char_start = token_start.char_start
            char_end = token_end.char_end
            return  str(char_start)+":"+str(char_end)
        else:
            return str(proto_obj.char_start)+":"+str(proto_obj.char_end)

    def generate_all_np(self,doc):
        np_list=[]
        relation_id = 0
        helper = DocHelper(doc)
        sentences = doc.sentence
        len_doc=len(sentences)

        for sent_index in range(len_doc):

            sentence = sentences[sent_index]
            constituents = sentence.constituent
            token_start=0
            for constituent in constituents:
                #print(constituent)
                if constituent.label=="S":
                    token_start=constituent.token_start
                    token_end=constituent.token_end
                    #print("Found ROOT!!")
                    break

            for token_index in range(token_start,token_end+1):

                np_cst_index = helper.getParentNPIndexFromLeafTokenIndex(sentence,token_index)
                base_np_cst_index = helper.getParentNPIndexFromLeafTokenIndex1(sentence,token_index)
                np_cst = sentence.constituent[np_cst_index]
                base_np_cst = sentence.constituent[base_np_cst_index]

                base_noun_phrase = re.sub("\n"," ",helper.text(base_np_cst))
                full_noun_phrase = re.sub("\n"," ",helper.text(np_cst))
                sentence_text = re.sub("\n"," ",helper.text(sentence))

                to_print = (doc.doc_id, str(sent_index), str(token_index),doc.token[token_index].word, base_noun_phrase, \
                            self.get_offset(base_np_cst, doc), full_noun_phrase, self.get_offset(np_cst, doc))
                #print ("\t".join(to_print))
                np_list.append(to_print)
        return np_list


    def generate_np_list(self):

        mongodb_host = os.environ.get("MONGODB_HOST","0.0.0.0") # change to biotm2.cis.udel.edu before dockerizing
        mongodb_port = os.environ.get("MONGODB_PORT","27017")

        db_text_from = os.environ.get("DBNAME_FROM_TEXT",'medline_current')
        fromCollectionText = os.environ.get("COLLECTION_FROM_TEXT",'text')
        # Database URI
        MONGODB_URI = 'mongodb://'+mongodb_host+':'+mongodb_port+'/'

        # Database object
        client = MongoClient(MONGODB_URI)
        opts = CodecOptions(document_class=SON)

        dbTextFrom = client[db_text_from]
        db = dbTextFrom[fromCollectionText].with_options(codec_options=opts)
        #####Iterate through all files in Input directory and create doc_list

        np_list_for_every_word=[]
        document_list = list()

        pmids_count = 0

        for pmid in self.pmid_list:
            pmid = pmid.strip()
            #print ("#"+pmid+"#")
            doc_id = pmid
            db_doc = db.find_one({'docId': pmid})
            if db_doc:
                pmids_count += 1
                doc_text = db_doc['text']
                raw_doc = document_pb2.Document()
                raw_doc.text = doc_text
                raw_doc.doc_id = doc_id
                document_list.append(raw_doc)
            else:
                raw_doc = document_pb2.Document()
                raw_doc.text = ''
                raw_doc.doc_id = doc_id
                document_list.append(raw_doc)
                #print('Empty pmid',pmid)
                #self.pmid_list.remove(pmid)


        # This is a simple function to make requests out of a list of documents. We
        # put 5 documents in each request.
        requests = request_iter_docs(document_list,
                                     request_size=5,
                                     request_type=rpc_pb2.Request.PARSE_BLLIP)

        # Given a request iterator, send requests in parallel and get responses.
        responses_queue = grpcapi.get_queue(server=self.edgServer,
                                            port=self.edgServerPort,
                                            request_thread_num=10,
                                            iterable_request=requests,
                                            edg_request_processor=True)
        count = 0
        for response in responses_queue:
            for doc in response.document:
                #print(doc)
                helper = DocHelper(doc)
                sentences = doc.sentence
                doc_id = doc.doc_id
                np_list_for_every_word.append(self.generate_all_np(doc))
                count += 1
        np_split_text_all_dic={}
        for doc_i in np_list_for_every_word:
            #for each doc

            np_split_text=[]
            previous_start=-1
            previous_end=-1
            for nw in doc_i:
                try:
                    pmid_text=document_list[self.pmid_list.index(nw[0])].text
                except:
                    print('ERROR!!!')
                    continue
                    #print('document list len:',len(document_list))
                    #print('nw:',nw)
                    #print('len of pmids:',len(self.pmid_list))
                    #print(self.pmid_list[:20])


                offset_np=nw[5]
                offset_np=offset_np.split(':')
                np_start=int(offset_np[0])
                np_end=int(offset_np[1])
                sent_index=int(nw[1])

                #if this word alreayd considered previously, then skip it
                if np_start==previous_start or np_start<previous_end:
                    continue
                #find the max np contain this word
                for wii in doc_i:

                    offset_np_candidate=wii[5]
                    offset_np_candidate=offset_np_candidate.split(':')
                    np_candidate_start=int(offset_np_candidate[0])
                    np_candidate_end=int(offset_np_candidate[1])
                    candidate_sent_index=int(wii[1])

                    if candidate_sent_index>sent_index:
                        break
                    if np_candidate_start!=np_start:
                        continue
                    else:
                        if np_candidate_end>np_end:
                            np_end=np_candidate_end

                previous_start=np_start
                previous_end=np_end
                np_split_text.append((nw[0],str(sent_index),pmid_text[np_start:np_end+1],str(np_start)+':'+str(np_end)))
            if len(np_split_text)>0:
                np_split_text_all_dic[str(np_split_text[0][0])]=np_split_text
        return np_split_text_all_dic

def run_np_generate_file(pmidFile,edgServer,edgServerPort):

    found_new_pmid=False
    pmidList = pd.read_csv(pmidFile,header=None).iloc[:,0].tolist()
    pmidList=list(set(pmidList))
    #for np dic
    if path.exists('glygen_set_abstract_np.json'):
        abstract_np_file='glygen_set_abstract_np.json'
        with open(abstract_np_file) as jfile:
            abstract_np_dic=json.load(jfile)
    else:
        abstract_np_file='glygen_set_abstract_np.json'
        with open(abstract_np_file,'w') as jfile:
            json.dump({},jfile)
        abstract_np_dic={}
    pmidSet_need_process=set()
    for pi in pmidList:
        if str(pi) in abstract_np_dic:
            continue
        found_new_pmid=True
        pmidSet_need_process.add(str(pi))

    ng=np_generator_for_abstrat(list(pmidSet_need_process),edgServer,edgServerPort)

    nps_dic=ng.generate_np_list()

    for pi in list(pmidSet_need_process):
        if str(pi) in nps_dic:
            nps=nps_dic[str(pi)]
            #add the generated NP to the cached file
            abstract_np_dic[str(pi)]=nps
    if found_new_pmid:
        print("Update NP file!")
        with open(abstract_np_file,'w') as jfile:
            json.dump(abstract_np_dic,jfile)

if __name__ == '__main__':

    pmidFile1='glygen_large.txt'
    pmidList1 = pd.read_csv(pmidFile1,header=None).iloc[:,0].tolist()
    pmidList1=list(set(pmidList1))
    pmidList1=[str(pi) for pi in pmidList1]
    ng=np_generator_for_abstrat(pmidList1)
    nps=ng.generate_np_list()
    print("pmid num:",len(pmidList1))
    print("nps dic len:",len(nps))
    #print(nps[:5])
    with open('glygen_large_abstract_np.json','w') as jfile:
        json.dump(nps,jfile)
    '''
    pmidList1=[12421832]
    pmidList1=[str(pi) for pi in pmidList1]
    ng=np_generator_for_abstrat(pmidList1)
    nps=ng.generate_np_list()['12421832']
    print("nps:",nps)
    '''