#!/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
import multiprocessing
import multiprocessing.pool
from random import randint
from collections import OrderedDict
from bson.son import SON
from bson.codec_options import CodecOptions
import shortuuid
#import add_sugar_protein_site as spsDetector
from add_sugar_protein_site_class import add_sugar_protein_site_class


def extract_pmidDoc(tuple_input):

    pmid, dbF, dbT, colF, colT, textT, edg=tuple_input
    edg_file = edg

    # --- 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_ENTITY", 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_ENTITY", colF)
    toCollectionName = os.environ.get("COLLECTION_TO_ENTITY", 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)

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

    if abstract_raw_doc and "entity" in abstract_raw_doc:
        entityDuidList = abstract_raw_doc["entity"].keys()
        aspsc = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)

        sugarEntityList = aspsc.generate_sugar_entity_list_from_mongoDB_v2(pmid)
        proteinEntityList = aspsc.add_oger_protein_complex(pmid)
        siteEntityList = aspsc.generate_site_entity_list_from_EDG_file(pmid, edg_file)

        if sugarEntityList:
            entityList = sugarEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        if proteinEntityList:
            entityList = proteinEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)
            #print('numOfEntitiesAdded proteinEntityList', numOfEntitiesAdded)

        if siteEntityList:
            entityList = siteEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc

    abstract_raw_doc = toDBCollection.find_one({"docId": pmid})
    if abstract_raw_doc and "entity" in abstract_raw_doc:
        aspsc1 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        proteinEntityList, update_entityList = aspsc1.add_oger_protein(pmid)
        if proteinEntityList:
            entityList = proteinEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)
        if update_entityList:
            numOfEntitiesAdded = update_entity(pmid, entityDuidList, update_entityList, toDBCollection)
        del aspsc1

        aspsc2 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        proteinEntityList = aspsc2.add_oger_complex(pmid)
        if proteinEntityList:
            entityList = proteinEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc2

        aspsc10 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        proteinEntityList = aspsc10.normalize_subunit_chain_complex(pmid)
        if proteinEntityList:
            entityList = proteinEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc10

        aspsc20 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        proteinEntityList = aspsc20.normalize_subunit_chain_complex_protein(pmid)
        if proteinEntityList:
            entityList = proteinEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc20

        aspsc_new0 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        newEntityList0 = aspsc_new0.generate_protein_entity_list_from_EDG_file(pmid, edg_file)
        if newEntityList0:
            entityList = newEntityList0
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc_new0

        aspsc_new = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        newEntityList = aspsc_new.add_new_protein_name(pmid)
        if newEntityList:
            entityList = newEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc_new


def add_entity(pmid, entityDuidList, entityList, toDBCollection):
    count = 0
    for entity in entityList:
        # print(entity)

        duid = get_duid(entityDuidList)
        entityDict = OrderedDict()
        if "duid" not in entity:
            entityDict["duid"] = duid
        else:
            entityDict["duid"] = entity["duid"]
        entityDict["entityType"] = entity["entityType"]
        entityDict["charEnd"] = entity["charEnd"]
        entityDict["source"] = entity["source"]
        entityDict["charStart"] = entity["charStart"]
        if "entityId" in entity:
            entityDict["entityId"] = entity["entityId"]
            if entity["entityType"]=='Complex':
                entityDict["entityId"]=[]
        else:
            entityDict["entityId"] = []
        entityDict["sentenceIndex"] = entity["sentenceIndex"]
        entityDict["entityText"] = entity["entityText"]

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


def update_entity(pmid, entityDuidList, entityList, toDBCollection):
    count = 0
    entityDoc = toDBCollection.find_one({"docId": str(pmid)})
    for entity_i in entityList:
        # print(entity)
        if entityDoc:

            for duid, entity in entityDoc["entity"].items():

                if entity_i[0] == duid:
                    entity["entityId"] = entity_i[1]
                    entityKey = "entity." + duid
                    toDBCollection.update_one({"docId": pmid}, {"$set": {entityKey: entity}})
                    count += 1
                    #print('Update the pubtator id', entity)
                    break

    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_generate_sugar_protein_site(pmidFile, dbF, dbT, colF, colT, textT, edg,proc_num):

    pmidList = pd.read_csv(pmidFile, header=None).iloc[:, 0].tolist()  # : for all rows, 0 for col1

    pool = MyPool(processes=proc_num)

    #results = [pool.apply(extract_pmidDoc, args=(str(pmid), dbF, dbT, colF, colT, textT, edg)) for pmid in pmidList]
    results = pool.map(extract_pmidDoc, [(str(pmid), dbF, dbT, colF, colT, textT, edg) for pmid in pmidList])


    pool.close()
    '''
    for index, pmid in enumerate(pmidList):
        print(index, ":", pmid)
        extract_pmidDoc((str(pmid), dbF, dbT, colF, colT, textT, edg))
        # break
    '''

class NoDaemonProcess(multiprocessing.Process):
    # make 'daemon' attribute always return False
    def _get_daemon(self):
        return False
    def _set_daemon(self, value):
        pass
    daemon = property(_get_daemon, _set_daemon)
class MyPool(multiprocessing.pool.Pool):
    Process = NoDaemonProcess

def sleepawhile(t):
    print("Sleeping %i seconds..." % t)
    time.sleep(t)
    return t

def work(num_procs):
    print("Creating %i (daemon) workers and jobs in child." % num_procs)
    pool = multiprocessing.Pool(num_procs)

    result = pool.map(sleepawhile,
                      [randint(1, 5) for x in range(num_procs)])

    # The following is not really needed, since the (daemon) workers of the
    # child's pool are killed when the child is terminated, but it's good
    # practice to cleanup after ourselves anyway.
    pool.close()
    pool.join()
    return result

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

    pmidFile = 'glygen_testset.txt'
    dbF = 'pubtator'
    dbT = 'glygen_test_new'
    colF = 'medline.aligned'
    colT = 'entities'
    textT = 'text'
    edg = 'glygen_testset_new.tsv'
    '''


    run_generate_sugar_protein_site(pmidFile, dbF, dbT, colF, colT, textT, edg)