
import os
import json
import pymongo
import datetime
import numpy as np

"""
[u'ViewCount',
 u'CreationDate',
 u'answersMin',
 u'Tags',
 u'acceptedAnswer',
 u'AnswerCount',
 u'CommentCount',
 u'answersMarks',
 u'AcceptedAnswerId',
 u'Score',
 u'commentsMarks',
 u'PostTypeId',
 u'OwnerUserId',
 u'_id',
 u'FavoriteCount',
 u'Id']
"""

def answer_covariates(answer_db_entry, date0, datef):
    arrivals_id = answer_db_entry["Id"]
    arrival_time = (answer_db_entry["CreationDate"] - date0).total_seconds()

    parent_question = answer_db_entry['ParentId']
    if "OwnerUserId" in answer_db_entry.keys():
        owner_id = answer_db_entry["OwnerUserId"]
    else:
        owner_id = -2
    score = answer_db_entry["Score"]

    arrival_and_covariates = [arrivals_id,
                              arrival_time,
                              parent_question,
                              owner_id,
                              score]

    return np.asarray(arrival_and_covariates),answer_db_entry["CreationDate"]

def question_covariates(arrival_db_entry, date0, datef):
    arrivals_id = arrival_db_entry["Id"]

    if (arrival_db_entry["AcceptedAnswerId"] is not None) and (arrival_db_entry["acceptedAnswer"]["date"] is not None):
        accepted_answer_id = arrival_db_entry["AcceptedAnswerId"]
        arrival_time = (arrival_db_entry["CreationDate"] - date0).total_seconds()
        if arrival_db_entry["acceptedAnswer"]["date"]  < datef:
            departure_time = (arrival_db_entry["acceptedAnswer"]["date"] - date0).total_seconds()
        else:
            departure_time = (arrival_db_entry["CreationDate"] - datef).total_seconds()
    else:
        arrival_time = (arrival_db_entry["CreationDate"] - date0).total_seconds()
        departure_time = (arrival_db_entry["CreationDate"] - datef).total_seconds()
        accepted_answer_id = -1

    if "OwnerUserId" in arrival_db_entry.keys():
        owner_id = arrival_db_entry["OwnerUserId"]
    else:
        owner_id = -1

    score = arrival_db_entry["Score"]
    if arrival_db_entry["FavoriteCount"] is not None:
        favorite_count = arrival_db_entry["FavoriteCount"]
    else:
        favorite_count = 0
    view_count = float(arrival_db_entry["ViewCount"])
    tags = []
    for t in arrival_db_entry["Tags"]:
        try:
            tags.append(TAG_LOOKUP[t])
        except:
            tags.append(500)

    tags.extend([0 for i in range(5 - len(tags))])#makes sure every tags covariate is of the same size

    arrival_and_covariates = [arrivals_id,
                              accepted_answer_id,
                              arrival_time,
                              owner_id,
                              score,
                              favorite_count,
                              view_count]
    arrival_and_covariates.extend(tags)
    arrival_and_covariates.append(departure_time)

    return np.asarray(arrival_and_covariates),arrival_db_entry["CreationDate"]

def check_last_file(data_dir,db,collection_list_names,date0,covariate_string="arrivals_covariates"):
    """

    :param collection_list_names:
    :param date0:
    :param data_dir:
    :param db:
    :param covariate_string:
    :return:
    """
    all_covariate_files = [int(filename.split("_")[-2]) for filename in os.listdir(data_dir) if covariate_string in filename]
    if len(all_covariate_files) > 0:
        # take latest file
        latest_file_num = max(all_covariate_files)
        latest_file = data_dir + covariate_string + "_{0}_.txt".format(latest_file_num)
        ALL_COVARIATES = np.loadtxt(latest_file)
        # take last id from latest file and check collections in order to find this index
        last_id = int(ALL_COVARIATES[:, 0][-1])
        latest_collection_index = 0
        for collection_name in collection_list_names:
            latest_collection = db[collection_name]
            arrivals_db = latest_collection.find_one({"Id": last_id})
            if arrivals_db is not None:
                latest_date = arrivals_db["CreationDate"]
                break
            latest_collection_index+=1
        return latest_collection, latest_collection_index, latest_file_num, latest_date
    else:
        latest_collection_index = 0
        latest_collection = db[collection_list_names[latest_collection_index]]
        latest_file_num = 0

        return latest_collection,latest_collection_index,latest_file_num, date0

def save_arrivals_covariates_step(db,collection_list,date0,datef,data_dir,MAXIMUM_ARRIVALS_PER_FILE,metadata):
    latest_collection, latest_collection_index, latest_file_num, latest_date = check_last_file(data_dir, db, collection_list_names,date0)

    print "latest_date ",latest_date
    print "latest_collection_index ",latest_collection_index
    block_file_name = data_dir + "arrivals_covariates_{0}_.txt".format(latest_file_num+1)
    ALL_COVARIATES = []
    while len(ALL_COVARIATES) < MAXIMUM_ARRIVALS_PER_FILE:
        num_covariates_left = MAXIMUM_ARRIVALS_PER_FILE - len(ALL_COVARIATES)
        FILE_BLOCK_CURSOR = latest_collection.find({'PostTypeId':1,'CreationDate':{"$gt":latest_date,
                                                                                   "$lt":datef}}).sort([("CreationDate",1)]).limit(num_covariates_left)
        if FILE_BLOCK_CURSOR.count() == 0:
            latest_collection_index +=1
            if len(collection_list) == latest_collection_index:
                break
            else:
                latest_collection = db[collection_list[latest_collection_index]]
        NEW_COVARIATES = []
        for a in FILE_BLOCK_CURSOR:
            latest_covariate,latest_date = question_covariates(a, date0, datef)
            NEW_COVARIATES.append(latest_covariate)
        ALL_COVARIATES.extend(NEW_COVARIATES)
    if len(ALL_COVARIATES) > 0:
        try:
            np.savetxt(block_file_name,ALL_COVARIATES)
        except:
            f = open(block_file_name,"w")
            for a in ALL_COVARIATES:
                num_str = map(str, a)
                if len(num_str) == 13:
                    try:
                        f.write(" ".join(num_str) + "\n")
                    except:
                        pass
            f.close()
            metadata["arrivals_files"][latest_file_num+1] = len(ALL_COVARIATES)
        total_arrivals = metadata["total_arrivals"]
        metadata["total_arrivals"] = total_arrivals + len(ALL_COVARIATES)
        return False, metadata
    else:
        return True, metadata

def save_departures_covariates_step(db,collection_list,date0,datef,data_dir,MAXIMUM_ARRIVALS_PER_FILE,metadata):
    latest_collection, latest_collection_index, latest_file_num, latest_date = check_last_file(data_dir, db, collection_list_names,date0,"departures_covariates")
    block_file_name = data_dir + "departures_covariates_{0}_.txt".format(latest_file_num+1)
    ALL_COVARIATES = []
    while len(ALL_COVARIATES) < MAXIMUM_ARRIVALS_PER_FILE:
        num_covariates_left = MAXIMUM_ARRIVALS_PER_FILE - len(ALL_COVARIATES)
        FILE_BLOCK_CURSOR = latest_collection.find({'PostTypeId':2,'CreationDate':{"$gt":latest_date,
                                                                                   "$lt":datef}}).sort([("CreationDate",1)]).limit(num_covariates_left)
        if FILE_BLOCK_CURSOR.count() == 0:
            latest_collection_index +=1
            if len(collection_list) == latest_collection_index:
                break
            else:
                latest_collection = db[collection_list[latest_collection_index]]
        NEW_COVARIATES = []
        for a in FILE_BLOCK_CURSOR:
            latest_covariate,latest_date = answer_covariates(a, date0, datef)
            NEW_COVARIATES.append(latest_covariate)
        ALL_COVARIATES.extend(NEW_COVARIATES)
    if len(ALL_COVARIATES) > 0:
        np.savetxt(block_file_name,ALL_COVARIATES)
        metadata["departures_files"][latest_file_num+1] = len(ALL_COVARIATES)
        total_departures = metadata["total_departures"]
        metadata["total_departures"] = total_arrivals + len(ALL_COVARIATES)
        return False,metadata
    else:
        return True,metadata

if __name__=="__main__":

    #data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/Service_Data/bitcoin_exchange/"
    #data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Stackoverflow_3Months/"
    #data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/bitcoin_exchange/"
    data_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Stackoverflow/"


    #================================
    # PICK DATA BASE AND DATES
    #================================
    client = pymongo.MongoClient("129.26.78.40:27017")
    #client = pymongo.MongoClient()

    db = client["Stackoverflow"]
    #db = client['BitCoinExchange']
    """
    id_index = 1
    TAG_LOOKUP = {}
    for tag in db["Tags"].find():
        TAG_LOOKUP[tag["TagName"]] = id_index
        id_index += 1
    json.dump(TAG_LOOKUP, open(data_dir + "tags_lookup.json", "w"))
    """
    TAG_LOOKUP = json.loads(open("/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/RecurrentNeuralPointProcess/data/Stackoverflow/best_tags_lookup.json","r").read())

    collection_list_names = [a for a in db.collection_names() if "Posts" in a]
    collection_list_names.sort()

    batch_index = 0
    for year in range(2011,2019):
        for month in range(1,13):
            for week in range(4):
                date0 = datetime.datetime(year,month, (week*7)+1)
                try:
                    datef = datetime.datetime(year,month, ((week+1)*7)+1)
                except:
                    datef = datetime.datetime(year, month, 28)
                print date0
                print datef
                #===================================
                # METADATA INFO
                #===================================
                columns_string = "arrivals_id,accepted_answer_id,arrival_time,owner_id,score,favorite_count,view_count,tag0,tags1,tags2,tags3,tags4,departure_time".split(
                    ",")
                columns_string2 = "arrivals_id,arrival_time,parent_question,owner_id,score".split(",")
                arrivals_header = dict(zip(range(len(columns_string)), columns_string))
                departures_header = dict(zip(range(len(columns_string)), columns_string2))
                total_arrivals = 0.
                total_departures = 0.
                metadata = {"arrivals_files": {},
                            "arrivals_file_header": arrivals_header,
                            "departures_files": {},
                            "departures_file_header": departures_header,
                            "real_temporal_support": [
                                (date0.strftime("%Y-%m-%d %H:%M:%S"), datef.strftime("%Y-%m-%d %H:%M:%S"))],
                            "temporal_support": [(0., (datef - date0).total_seconds())],
                            "total_arrivals":0.,
                            "total_departures":0.}
                #--------CREATE DIR-------------------------------
                current_dir = data_dir + str(batch_index) + "/"
                os.makedirs(current_dir)
                #=================================
                #  SAVE FILES
                #=================================
                MAXIMUM_ARRIVALS_PER_FILE = 50000
                while True:
                    ready,metadata = save_arrivals_covariates_step(db,collection_list_names,date0,datef,current_dir,MAXIMUM_ARRIVALS_PER_FILE,metadata)
                    if ready:
                        break

                while True:
                    ready,metadata = save_departures_covariates_step(db,collection_list_names,date0,datef,current_dir,MAXIMUM_ARRIVALS_PER_FILE,metadata)
                    if ready:
                        break

                json.dump(metadata,open(current_dir+"metadata.json","w"))
                batch_index+=1