#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division,print_function
import time,re,sys,os,ast
import pandas as pd
import pymongo
from pymongo import MongoClient
import random,operator
from collections import OrderedDict
import shortuuid
from bson.son import SON
from bson.codec_options import CodecOptions

def get_duid(entityDuidList):
    while True:
        duid = "PU:" + shortuuid.ShortUUID().random(length=4)

        if duid not in entityDuidList:
            break

    # duid = "PU:" + shortuuid.ShortUUID().random(length=4)
    return duid

amino_acid_shrot=['Ala', 'Arg', 'Asn', 'Asp', 'Cys', 'Glu', 'Gln', 'Gly', 'His', 'Hyp', 'Ile', 'Leu', 'Lys', 'Met', 'Phe', 'Pro', 'Glp', 'Ser', 'Thr', 'Trp', 'Tyr', 'Val',]
amino_acid_shrot_lower=[i.lower() for i in amino_acid_shrot]
amino_acid_shrot_lower_con='|'.join(amino_acid_shrot_lower)

def remove_singleLetterSingleOrDoubleDigitSites_thatAreProteins(pmid,entityList,DBCollection,TextDBCollection):
    site_list=[]
    for duid,entity in entityList.items():

        if entity["entityType"] == "SpecificSite" or entity["entityType"] == "SiteSequence" or entity["entityType"] == "Site" or entity["entityType"] == "AminoAcid":
            # Throw out all single digit single letter,
            # Plus double digits for P,H,S,C
            site_list.append((entity["entityText"],entity["duid"],entity["charStart"],entity["charEnd"]))

            if re.match(r"(P|H|S|C)(-)?\d{1,2}$",entity["entityText"],re.I) or \
                    re.match(r"[A-MO-Z]( |-)?\d{1,3}$",entity["entityText"].upper(),re.I) or \
                    re.match(r"(N)(-)?\d{1}$",entity["entityText"],re.I):
                entityKey = "entity." + duid
                DBCollection.update_one({"docId":pmid},{ "$unset": { entityKey: 1}})
        elif entity["entityType"] == "Gene" or entity["entityType"] == "Protein":
            if re.match(r"("+amino_acid_shrot_lower_con+r")( |-)?\d{1,3}$",entity["entityText"].lower(),re.I):
                entityKey = "entity." + duid
                DBCollection.update_one({"docId":pmid},{ "$unset": { entityKey: 1}})
        elif entity["entityType"] == "Chemical" and re.match(r".*(\d+)$",entity["entityText"],re.I):
            site_list.append((entity["entityText"],entity["duid"],entity["charStart"],entity["charEnd"]))

    #check site sequence and delete them
    duid_set=set()
    abstract_raw_text = TextDBCollection.find_one({"docId":pmid})
    if abstract_raw_text:
        abstract_text= abstract_raw_text["text"]
    else:
        abstract_text=None

    if abstract_text and site_list:

        for si in range(len(site_list)):
            for sii in range(si+1,len(site_list)):
                if site_list[si][2]>site_list[sii][3]:
                    text_between=abstract_text[site_list[sii][3]+1:site_list[si][2]]
                else:
                    text_between=abstract_text[site_list[si][3]+1:site_list[sii][2]]

                if re.match(r"^(\s?-\s?X?\(?\s?\d*\s?\)?\s?-?\s?)$",text_between,re.I):
                    #print('pmid: ',pmid)
                    #print(abstract_text[min(site_list[si][2],site_list[sii][2]):max(site_list[si][3],site_list[sii][3])+1])
                    duid_set.add(site_list[si][1])
                    duid_set.add(site_list[sii][1])

    abstract_raw_doc = DBCollection.find_one({"docId":pmid}) # toDBCollection contains entities; fromDBCollection contains abstract
    if abstract_raw_doc and "entity" in abstract_raw_doc:
        entityList=abstract_raw_doc["entity"]
    else:
        entityList=[]
    for duid,entity in entityList.items():
        if duid in duid_set:
            entityKey = "entity." + duid
            DBCollection.update_one({"docId":pmid},{ "$unset": { entityKey: 1}})

#merge the sites into one big site
def merge_sites_in_conjunction(pmid,DBCollection,TextDBCollection):

    abstract_raw_text = TextDBCollection.find_one({"docId":pmid})
    if abstract_raw_text:
        abstract_text= abstract_raw_text["text"]
    else:
        abstract_text=None

    abstract_raw_doc = DBCollection.find_one({"docId":pmid})

    if abstract_raw_doc and "entity" in abstract_raw_doc:
        entityList=abstract_raw_doc["entity"]
        entityDuidList = abstract_raw_doc["entity"].keys()
    else:
        entityList=[]
        entityDuidList=[]


    site_list=[]
    for duid,entity in entityList.items():

        if entity["entityType"] == "SpecificSite" or entity["entityType"] == "SiteSequence" or entity["entityType"] == "Site" or entity["entityType"] == "AminoAcid":
            # Throw out all single digit single letter,
            # Plus double digits for P,H,S,C
            site_list.append((entity["entityText"],entity["duid"],entity["charStart"],entity["charEnd"],entity["entityType"],entity["source"],entity["sentenceIndex"]))

    delete_duid_set=[]
    add_site_list=[]
    duid_for_original_site=[i[1] for i in site_list]
    entityDict_add_list=[]

    if abstract_text:

        for iter_i in range(50):
            #print("Iteration ",iter_i)
            site_list+=add_site_list
            site_list=[si for si in site_list if si[1] not in delete_duid_set]
            #print('site_list:',site_list)

            if len(site_list)<2:
                break
            add_site_list=[]
            for si in range(len(site_list)):
                for sii in range(si+1,len(site_list)):
                    found_merge=False
                    if (site_list[si][2]>=site_list[sii][2] and site_list[si][2]<=site_list[sii][3]) or \
                            (site_list[sii][2]>=site_list[si][2] and site_list[sii][2]<=site_list[si][3]):
                        continue

                    if site_list[si][2]>site_list[sii][3]:
                        text_between=abstract_text[site_list[sii][3]+1:site_list[si][2]]
                    else:
                        text_between=abstract_text[site_list[si][3]+1:site_list[sii][2]]

                    if (text_between.strip()==',' or text_between.strip()=='and' or text_between.strip()==', and' or text_between.strip()=='') and site_list[si][6]==site_list[sii][6]:
                        delete_duid_set.append(site_list[si][1])

                        delete_duid_set.append(site_list[sii][1])
                        merged_text=abstract_text[min(site_list[si][2],site_list[sii][2]):max(site_list[si][3],site_list[sii][3])+1]

                        duid = get_duid(entityDuidList)
                        entityDuidList.append(duid)
                        entityDict = OrderedDict()
                        entityDict["duid"] = duid
                        entityDict["entityType"] = site_list[si][4]
                        entityDict["charEnd"] = max(site_list[si][3],site_list[sii][3])
                        entityDict["source"] = site_list[si][5]
                        entityDict["charStart"] = min(site_list[si][2],site_list[sii][2])
                        entityDict["entityId"] = []
                        entityDict["sentenceIndex"] = site_list[si][6]
                        entityDict["entityText"] = merged_text
                        #print("entity1:",site_list[si])
                        #print("entity2:",site_list[sii])
                        #print("entityDict:",entityDict)

                        add_site_list.append((entityDict["entityText"],entityDict["duid"],entityDict["charStart"],entityDict["charEnd"],entityDict["entityType"],entityDict["source"],entityDict["sentenceIndex"]))

                        entityDict_add_list.append(entityDict)
                        found_merge=True
                        break

                if found_merge:
                    break
            if not found_merge:
                break

    #print('pmid:',pmid)
    #print('entityDict_add_list:',entityDict_add_list)
    for entity_i in entityDict_add_list:
        if 'duid' in entity_i.keys() and entity_i['duid'] in [si[1] for si in site_list]:
            entityKey = "entity." + entity_i['duid']
            #print("Entity will be added:",entity_i)
            DBCollection.update_one({"docId":pmid},{ "$set": { entityKey: entity_i}})


    for duid,entity in entityList.items():
        if duid in delete_duid_set:
            entityKey = "entity." + duid
            #print("Entity duid will be deleted:",entity)
            DBCollection.update_one({"docId":pmid},{ "$unset": { entityKey: 1}})
    return None

def extract_pmidDoc(pmid,DBCollection, TextDBCollection):
    abstract_raw_doc = DBCollection.find_one({"docId":pmid}) # toDBCollection contains entities; fromDBCollection contains abstract
    if abstract_raw_doc and "entity" in abstract_raw_doc:
        remove_singleLetterSingleOrDoubleDigitSites_thatAreProteins(pmid,abstract_raw_doc["entity"],DBCollection,TextDBCollection)
        merge_sites_in_conjunction(pmid,DBCollection,TextDBCollection)
        # if listOfEntityToRemove:
        #     numOfEntitiesRemoved = remove_entity(pmid,listOfEntityToRemove,DBCollection)
        #     print(numOfEntitiesRemoved)


def run_remove_site(pmidFile, DBCollection, TextDBCollection):
    pmidList = pd.read_csv(pmidFile,header=None).iloc[:,0].tolist()
    for index,pmid in enumerate(pmidList):
        #print(index,":",pmid)
        extract_pmidDoc(str(pmid),DBCollection,TextDBCollection)

if __name__ == "__main__":

    pmidFile = sys.argv[1]
    db = sys.argv[2]
    col = sys.argv[3]
    text = sys.argv[4]
    '''
    pmidFile = 'glygen_testset.txt'
    db = 'glygen_test_new'
    col = 'entities'
    text = 'text'
    '''
    #--- create database instances---
    # Environment variables
    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_name = os.environ.get("DBNAME",db) # change database name for your own dbName
    CollectionName = os.environ.get("COLLECTION",col)
    TextCollectionName = os.environ.get("COLLECTION",text)
    # Database URI
    MONGODB_URI = 'mongodb://'+mongodb_host+':'+mongodb_port+'/'
    # Database object
    client = MongoClient(MONGODB_URI)
    opts = CodecOptions(document_class=SON)
    # Database
    dbName = client[db_name] # glygen
    # Collection
    DBCollection = dbName[CollectionName].with_options(codec_options=opts)
    TextDBCollection = dbName[TextCollectionName].with_options(codec_options=opts)
    # Run actual functions
    run_remove_site(pmidFile, DBCollection, TextDBCollection)