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 json
from create_textTable_forDB import run_create_text_table
from create_entityTable_forDB import run_create_entity_table
from update_entity_in_database import update_entity_info
from generate_aminoAcidSiteEntities import run_detect_site_entity
from remove_sites_afterFurtherAnalysis_fromGlygenDB import run_remove_site
from generate_sugar_protein_site import run_generate_sugar_protein_site
from add_protein_ancronym import run_add_protein_ancronym
from add_sugar_protein_site_class import add_sugar_protein_site_class
from f_terms_EDG import add_columns, process_edg_output2, merge_rows, substract_file, extrac_certain_columns_from_postprocessed_file
from site_substrate_fusion import site_substrate_fusion
from process_edg_relations_pmids import run_process_edg_relations,exstract_subset_edgFile
from np_generator import run_np_generate_file
from add_oger_entity import run_update_oger_file
from convert_uniprot_site import convert_uniprot_id_and_site_info
from file_process import overwrite_previous_file, append_to_previous_file

pmidFile = os.environ.get("PMID_FILE","glygen_test.txt")
edgFileAll = os.environ.get("EDG_FILE","glygen_set_all.tsv")
edgRuleFile = os.environ.get("EDG_RULE_FILE","edg_all_rules_101820.txt")
edgServer= os.environ.get("EDG_SERVER","0.0.0.0")
edgServerPort= os.environ.get("EDG_SERVER_PORT","8903")
#available model: append, overwrite
#but for glygen_set, glygen_large, glygen_full_length, only append mode is supported
processMode= os.environ.get("PROCESS_MODE","append")
processNumber= os.environ.get("PROC_NUM",10)

#--- 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_text_from = os.environ.get("DBNAME_FROM_TEXT",'medline_current')
db_entity_from = os.environ.get("DBNAME_FROM_ENTITY",'pubtator')
db_name_to = os.environ.get("DBNAME_TO",'glygen_test') # change database name for your own dbName

fromCollectionText = os.environ.get("COLLECTION_FROM_TEXT",'text')
fromCollectionEntity = os.environ.get("COLLECTION_FROM_ENTITY",'medline.aligned')
toCollectionText = os.environ.get("COLLECTION_TO_TEXT",'text')
toCollectionEntity = os.environ.get("COLLECTION_TO_ENTITY",'entities')
# Database URI
MONGODB_URI = 'mongodb://'+mongodb_host+':'+mongodb_port+'/'

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

# Database
dbTextFrom = client[db_text_from]
dbEntityFrom = client[db_entity_from]
dbNameTo = client[db_name_to] # New DB: glygen

# Collection
TextfromDBCollection = dbTextFrom[fromCollectionText].with_options(codec_options=opts)
EntityfromDBCollection = dbEntityFrom[fromCollectionEntity].with_options(codec_options=opts)
toDBCollectionText = dbNameTo[toCollectionText].with_options(codec_options=opts)
toDBCollectionEntity = dbNameTo[toCollectionEntity].with_options(codec_options=opts)

#step 0 if the PROCESS_MODE is replace, drop the collection first
if processMode=='overwrite' and db_name_to=='glygen_set_new':
    toDBCollectionEntity.drop()
    toDBCollectionText.drop()
elif processMode=='append':
    print("PROCESS_MODE is append!")
else:
    sys.exit("Overwriting of glygen_set and glygen_larege is not supported!")
#step 0 generate EDG file
pmidListInMongo = toDBCollectionEntity.distinct('docId')
run_process_edg_relations(pmidFile, TextfromDBCollection,edgRuleFile,edgServer,edgServerPort,edgFileAll)
edgFile=exstract_subset_edgFile(pmidFile,edgFileAll,list(pmidListInMongo),db_name_to)

#step 0 generate np file
run_np_generate_file(pmidFile,edgServer,edgServerPort)

#step 0 update oger entity file
run_update_oger_file()

#step 1
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 1/11: Create DB text table!")
run_create_text_table(pmidFile,TextfromDBCollection,toDBCollectionText)

#step 2
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 2/11: Create DB entity table!")
run_create_entity_table(pmidFile,EntityfromDBCollection,toDBCollectionEntity)

#step 3
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 3/11: Update entities in DB tables!")
UEI=update_entity_info(db_entity_from,db_name_to,fromCollectionEntity,toCollectionEntity,fromCollectionText,toCollectionText)

pmidList = pd.read_csv(pmidFile,header=None).iloc[:,0].tolist()
pmidList=list(set(pmidList))
for pi in pmidList:
    UEI.update_acronym_normalization(str(pi))
for pi in pmidList:
    UEI.extend_protein_normalization(str(pi))
for pi in pmidList:
    UEI.remove_protein_is_part_of_sugar(str(pi))
for pi in pmidList:
    UEI.remove_site_is_part_of_sugar(str(pi))
for pi in pmidList:
    UEI.remove_protein_is_part_of_site(str(pi))
for pi in pmidList:
    UEI.update_specific_protein_names(str(pi))

#step 4
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 4/11: Site detection!")
run_detect_site_entity(pmidFile,EntityfromDBCollection,toDBCollectionEntity)

#step 5
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 5/11: Site detection refinement!")
run_remove_site(pmidFile, toDBCollectionEntity, toDBCollectionText)

#step 6
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 6/11: Detect suger, protein!")
run_generate_sugar_protein_site(pmidFile, db_entity_from, db_name_to, fromCollectionEntity, toCollectionEntity, toCollectionText, edgFile, int(processNumber))

#step 7
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 7/11: Add protein ancronym!")
run_add_protein_ancronym(pmidFile, db_entity_from, db_name_to, fromCollectionEntity, toCollectionEntity, toCollectionText, edgFile,int(processNumber))

#step 8
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 8/11: Update entity information!")
UEI=update_entity_info(db_entity_from,db_name_to,fromCollectionEntity,toCollectionEntity,fromCollectionText,toCollectionText)

pmidList = pd.read_csv(pmidFile,header=None).iloc[:,0].tolist()
pmidList=list(set(pmidList))
for pi in pmidList:
    UEI.update_acronym_normalization(str(pi))
for pi in pmidList:
    UEI.extend_protein_normalization(str(pi))
for pi in pmidList:
    UEI.remove_protein_is_part_of_sugar(str(pi))
for pi in pmidList:
    UEI.remove_site_is_part_of_sugar(str(pi))
for pi in pmidList:
    UEI.remove_protein_is_part_of_site(str(pi))
for pi in pmidList:
    UEI.update_specific_protein_names(str(pi))

#step 9
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 9/11: Remove unwanted entities!")
run_remove_site(pmidFile, toDBCollectionEntity, toDBCollectionText)

#step 10
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 10/11: Post-processing!")
ASPSC=add_sugar_protein_site_class(db_entity_from,db_name_to,fromCollectionEntity,toCollectionEntity,fromCollectionText,toCollectionText)
tsv_input=edgFile

pmidList = pd.read_csv(pmidFile,header=None).iloc[:,0].tolist()
pmidList=[str(i) for i in pmidList]

tsv_input_split=tsv_input.split('.')
tsv_output=tsv_input_split[0]+'_postprocessing_before_merge_tmp.tsv'
tsv_output2=tsv_input_split[0]+'_filtered_out_before_merge_tmp.tsv'
tsv_output_cterm=tsv_input_split[0]+'_cterm_before_merge_tmp.tsv'
tsv_centain_colums=tsv_input_split[0]+'_output_tmp.tsv'

tsv_input_new=tsv_input_split[0]+'_new_tmp.tsv'
#add some columns on the files
add_columns(tsv_input,tsv_input_new,['PSource','NProtein','NID','SiteSource','SiteName'])
process_edg_output2(tsv_input_new,tsv_output,tsv_output2,tsv_output_cterm,ASPSC,pmidList)

output_file=tsv_input_split[0]+'_postprocessing_tmp.tsv'
output_file2=tsv_input_split[0]+'_filtered_out_tmp.tsv'
output_file_cterm=tsv_input_split[0]+'_cterm_tmp.tsv'
output_file_other=tsv_input_split[0]+'_other_tmp.tsv'
final_file=tsv_input_split[0]+'_final_tmp.tsv'

merge_rows(tsv_output,output_file)
merge_rows(tsv_output2,output_file2)
merge_rows(tsv_output_cterm,output_file_cterm)
substract_file(output_file,output_file_cterm,output_file_other)

extrac_certain_columns_from_postprocessed_file(output_file,final_file)

#step 11
currenttime = time.asctime(time.localtime(time.time()))
print("["+currenttime+"]:","Step 11/11: Generate output files!")
task_name = edgFile.split('.')[0]
db_name = db_name_to

file_input_site_fusion1=task_name+'_cterm_tmp.tsv'
file_input_site_fusion2=task_name+'_postprocessing_tmp.tsv'
output_file_site_fusion=task_name+'_postprocessing_site_fusion_tmp.tsv'
final_file_site_fusion =task_name+'_site_fusion_final_tmp.tsv'

site_substrate_fusion(file_input_site_fusion1,file_input_site_fusion2,output_file_site_fusion,db_name,pmidFile)
extrac_certain_columns_from_postprocessed_file(output_file_site_fusion,final_file_site_fusion)

#convert the output to the 5-column output
file_input_5_column=task_name+'_site_fusion_final_tmp.tsv'
output_file_5_column=task_name+'_site_fusion_output_tmp.tsv'
output_file_contain_empty_site_5_column=task_name+'_site_fusion_output_contain_empty_site_tmp.tsv'
columns_list_5_column=['pmid','uniprotkb_ac','protein','amino_acid','site_position']
convert_uniprot_id_and_site_info(file_input_5_column,output_file_5_column,output_file_contain_empty_site_5_column,columns_list_5_column,'')

#update the output based on the PROCESS_MODE
if task_name in ['glygen_set','glygen_large','glygen_pmc','glygen_set_new']:
    post_processed_file=task_name+'_postprocessing.tsv'
    post_processed_site_fusion_file=task_name+'_postprocessing_site_fusion.tsv'
    final_5_column_file=task_name+'_site_fusion_output.tsv'
    final_5_column_empty_site_file=task_name+'_site_fusion_output_contain_empty_site.tsv'

    if processMode=='overwrite':
        overwrite_previous_file(post_processed_file,file_input_site_fusion2)
        overwrite_previous_file(post_processed_site_fusion_file,output_file_site_fusion)
        overwrite_previous_file(final_5_column_file,output_file_5_column)
        overwrite_previous_file(final_5_column_empty_site_file,output_file_contain_empty_site_5_column)
    elif processMode=='append':

        append_to_previous_file(post_processed_file,file_input_site_fusion2)
        append_to_previous_file(post_processed_site_fusion_file,output_file_site_fusion)
        append_to_previous_file(final_5_column_file,output_file_5_column)
        append_to_previous_file(final_5_column_empty_site_file,output_file_contain_empty_site_5_column)