import argparse
import time
import csv
import pickle
import operator
import datetime
import os

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='Nowplaying', help='dataset name: diginetica/yoochoose/sample')
opt = parser.parse_args()
print(opt)

dataset = 'nowplaying.csv'


print("-- Starting @ %ss" % datetime.datetime.now())
with open(dataset, "r") as f:
    reader = csv.DictReader(f, delimiter='\t')
    sess_clicks = {}
    sess_date = {}
    ctr = 0
    curid = -1
    curdate = None
    for data in reader:
        sessid = int(data['SessionId'])
        if curdate and not curid == sessid:
            date = curdate
            sess_date[curid] = date
        curid = sessid

        item = int(data['ItemId'])
        curdate = float(data['Time'])
        # print(item, curdate)
        if sessid in sess_clicks:
            sess_clicks[sessid] += [(item, curdate)] ## Support timestamp
        else:
            sess_clicks[sessid] = [(item, curdate)] ## support timestamp
        ctr += 1
    date = float(data['Time'])
    sess_date[curid] = date
print('ctr:', ctr)
print("-- Reading data @ %ss" % datetime.datetime.now())

# Filter out length 1 sessions
for s in list(sess_clicks):
    if len(sess_clicks[s]) == 1:
        del sess_clicks[s]
        del sess_date[s]

# Count number of times each item appears
iid_counts = {}
for s in sess_clicks:
    seq = sess_clicks[s]
    for iid in seq:
        if iid[0] in iid_counts:
            iid_counts[iid[0]] += 1
        else:
            iid_counts[iid[0]] = 1

sorted_counts = sorted(iid_counts.items(), key=operator.itemgetter(1))

length = len(sess_clicks)
for s in list(sess_clicks):
    curseq = sess_clicks[s]
    filseq = list(filter(lambda i: iid_counts[i[0]] >= 5, curseq))
    if len(filseq) < 2 or len(filseq) > 30:
        del sess_clicks[s]
        del sess_date[s]
    else:
        sess_clicks[s] = filseq

# Split out test set based on dates
dates = list(sess_date.items())
maxdate = dates[0][1]

for _, date in dates:
    if maxdate < date:
        maxdate = date

# Two months for test
splitdate = maxdate - 60 * 86400

print('Splitting date', splitdate)      # Yoochoose: ('Split date', 1411930799.0)
tra_sess = filter(lambda x: x[1] < splitdate, dates)
tes_sess = filter(lambda x: x[1] > splitdate, dates)

# Sort sessions by date
tra_sess = sorted(tra_sess, key=operator.itemgetter(1))     # [(session_id, timestamp), (), ]
tes_sess = sorted(tes_sess, key=operator.itemgetter(1))     # [(session_id, timestamp), (), ]
print(len(tra_sess))    # 186670    # 7966257
print(len(tes_sess))    # 15979     # 15324
print(tra_sess[:3])
print(tes_sess[:3])
print("-- Splitting train set and test set @ %ss" % datetime.datetime.now())

# Choosing item count >=5 gives approximately the same number of items as reported in paper
item_dict = {}
# Convert training sessions to sequences and renumber items to start from 1
def obtian_tra():
    train_ids = []
    train_seqs = []
    train_times = []
    train_dates = []
    item_ctr = 1
    for s, date in tra_sess:
        seq = sess_clicks[s]
        outseq  = []
        outtime = []
        for i in seq:
            outtime += [i[1]]
            if i in item_dict:
                outseq += [item_dict[i[0]]]
            else:
                outseq += [item_ctr]
                item_dict[i[0]] = item_ctr
                item_ctr += 1
        if len(outseq) < 2:  # Doesn't occur
            continue
        train_ids   += [s]
        train_dates += [date]
        train_seqs  += [outseq]
        train_times += [outtime] 
    print('item_ctr')
    print(item_ctr)     # 43098, 37484
    return train_ids, train_dates, train_seqs, train_times


# Convert test sessions to sequences, ignoring items that do not appear in training set
def obtian_tes():
    test_ids   = []
    test_seqs  = []
    test_times = []
    test_dates = []
    for s, date in tes_sess:
        seq = sess_clicks[s]
        outseq  = []
        outtime = []
        for i in seq:
            if i in item_dict:
                outseq  += [item_dict[i[0]]]
                outtime += [i[1]]
        if len(outseq) < 2:
            continue
        test_ids   += [s]
        test_dates += [date]
        test_seqs  += [outseq]
        test_times += [outtime]
    return test_ids, test_dates, test_seqs, test_times


tra_ids, tra_dates, tra_seqs, tra_times = obtian_tra()
tes_ids, tes_dates, tes_seqs, tes_times  = obtian_tes()


def process_seqs(iseqs, idates, itimes):
    out_seqs  = []
    out_dates = []
    out_times = []
    labs = []
    ids = []
    for id, seq, date, times in zip(range(len(iseqs)), iseqs, idates, itimes):
        for i in range(1, len(seq)):
            tar = seq[-i]
            labs      += [tar]
            out_seqs  += [seq[:-i]]
            out_times += [times[:-i]] 
            out_dates += [date]
            ids += [id]
    return out_seqs, out_dates, out_times, labs, ids


tr_seqs, tr_dates, tr_times, tr_labs, tr_ids = process_seqs(tra_seqs, tra_dates, tra_times)
te_seqs, te_dates, te_times, te_labs, te_ids = process_seqs(tes_seqs, tes_dates, tes_times)
tra = (tr_seqs, tr_times, tr_labs)
tes = (te_seqs, te_times, te_labs)
print('train_test')
print(len(tr_seqs))
print(len(te_seqs))
print(tr_seqs[:3], tr_times[:3], tr_dates[:3], tr_labs[:3])
print(te_seqs[:3], tr_times[:3], te_dates[:3], te_labs[:3])
all = 0

for seq in tra_seqs:
    all += len(seq)
for seq in tes_seqs:
    all += len(seq)
print('avg length: ', all/(len(tra_seqs) + len(tes_seqs) * 1.0))
print('all:', all)

if not os.path.exists('nowplaying'):
    os.makedirs('nowplaying')
pickle.dump(tra, open('nowplaying/train.txt', 'wb'))
pickle.dump(tes, open('nowplaying/test.txt', 'wb'))
pickle.dump(tra_seqs, open('Nowplaying/all_train_seq.txt', 'wb'))
