#!/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", 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)
    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()

        aspsc_new1 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        acronymEntityList = aspsc_new1.add_acronym_protein(pmid)
        if acronymEntityList:
            entityList = acronymEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc_new1
        aspsc_new2 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        extraEntityList = aspsc_new2.add_extra_protein_entity(pmid)
        if extraEntityList:
            entityList = extraEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)

        del aspsc_new2
        aspsc_new3 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        aspsc_new3.extend_protein_names(pmid)

        del aspsc_new3
        aspsc_new4 = add_sugar_protein_site_class(dbF, dbT, colF, colT, textT, textT)
        ctermEntityList = aspsc_new4.add_c_term_protein(pmid)
        if ctermEntityList:
            entityList = ctermEntityList
            numOfEntitiesAdded = add_entity(pmid, entityDuidList, entityList, toDBCollection)
        del aspsc_new4


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"]
        if "entityId" in entity:
            entityDict["entityId"] = entity["entityId"]
        else:
            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_add_protein_ancronym(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
    # print(pmidList)
    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()

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 = 'unicarb.txt'
    # dbF = 'pubtator'
    # dbT = 'unicarb'
    # colF = 'medline.aligned'
    # colT = 'entities'

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