#!/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
from bson.son import SON
from bson.codec_options import CodecOptions
import shortuuid
import aminoAcidSiteDetectionTool_v4 as aaDetector

def extract_pmidDoc(pmid,fromDBCollection,toDBCollection):
    abstract_raw_doc = toDBCollection.find_one({"docId":pmid}) # toDBCollection contains entities; fromDBCollection contains abstract
    if abstract_raw_doc and "entity" in abstract_raw_doc:
        entityDuidList = abstract_raw_doc["entity"].keys()
        aminoAcidSiteEntityList = aaDetector.run_forEachPmid(pmid,fromDBCollection,toDBCollection)
        #print('Detected sites: ',aminoAcidSiteEntityList)
        if aminoAcidSiteEntityList:
            entityList = aminoAcidSiteEntityList["entity"]
            numOfEntitiesAdded = add_entity(pmid,entityDuidList,entityList,toDBCollection)
            #print(numOfEntitiesAdded)

def add_entity(pmid,entityDuidList,entityList,toDBCollection):
    count = 0
    for entity in entityList:
	    #print(entity)
        duid = get_duid(entityDuidList)
        entityDict = OrderedDict()
        entityDict["duid"] = duid
        entityDict["entityType"] = entity["entityType"]
        entityDict["charEnd"] = entity["charEnd"]
        entityDict["source"] = entity["source"]
        entityDict["charStart"] = entity["charStart"]
        entityDict["entityId"] = []
        entityDict["sentenceIndex"] = entity["sentenceIndex"]
        entityDict["entityText"] = entity["entityText"]

        # entityGroup[duid] = entityDict
        entityKey = "entity." + duid
        # print entityDict
        # print
        toDBCollection.update_one({"docId":pmid},{ "$set": { entityKey: entityDict}})
        count+=1
    # print " - "*10
    return count

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



def run_detect_site_entity(pmidFile,fromDBCollection,toDBCollection):
    pmidList = pd.read_csv(pmidFile,header=None).iloc[:,0].tolist() # : for all rows, 0 for col1
    # print(pmidList)
    for index,pmid in enumerate(pmidList):
        #print(index,":",pmid)
        extract_pmidDoc(str(pmid),fromDBCollection,toDBCollection)
        # break

if __name__ == "__main__":
    pmidFile = sys.argv[1]
    dbF = sys.argv[2]
    dbT = sys.argv[3]
    colF = sys.argv[4]
    colT = sys.argv[5]


    #--- 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_from = os.environ.get("DBNAME_FROM",dbF) # change database name for your own dbName
    db_name_to = os.environ.get("DBNAME_TO",dbT) # change database name for your own dbName

    fromCollectionName = os.environ.get("COLLECTION_FROM",colF)
    toCollectionName = os.environ.get("COLLECTION_TO",colT)
    # Database URI
    MONGODB_URI = 'mongodb://'+mongodb_host+':'+mongodb_port+'/'

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

    # Database
    dbNameFrom = client[db_name_from] # medline
    dbNameTo = client[db_name_to] # New DB: glygen

    # Collection
    fromDBCollection = dbNameFrom[fromCollectionName].with_options(codec_options=opts)
    toDBCollection = dbNameTo[toCollectionName].with_options(codec_options=opts)


    run_detect_site_entity(pmidFile,fromDBCollection,toDBCollection)