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

def calculate_similarity(seq1, seq2):
    # 计算两个会话序列的相似度，使用交集/并集的比例
    set1, set2 = set(seq1), set(seq2)
    intersection = len(set1.intersection(set2))
    union = len(set1.union(set2))
    return intersection / union if union > 0 else 0

def get_top_ten_similar_sessions(current_seq, all_seqs):
    # 获取与当前会话相似度最高的十个会话
    similarities = []
    for seq in all_seqs:
        if seq != current_seq:  # 避免与自身比较
            similarity = calculate_similarity(current_seq, seq)
            similarities.append((seq, similarity))
    similarities.sort(key=lambda x: x[1], reverse=True)
    return [seq for seq, _ in similarities[:10]]

# 解析参数部分保持不变
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='Tmall', help='dataset name: diginetica/yoochoose/sample')
opt = parser.parse_args()
print(opt)

# 数据预处理部分修改
with open('tmall_data.csv', 'w') as tmall_data:
    with open('dataset15.csv', 'r') as tmall_file:
        header = tmall_file.readline()
        tmall_data.write(header)
        lines = tmall_file.readlines()
        # 只处理前百分之一的数据
        num_lines = len(lines) // 500
        for line in lines[:num_lines]:
            data = line[:-1].split('\t')
            if int(data[2]) > 120000:
                break
            tmall_data.write(line)

print("-- Starting @ %ss" % datetime.datetime.now())
with open('tmall_data.csv', "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'])

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

# 过滤部分保持不变
for s in list(sess_clicks):
    if len(sess_clicks[s]) < 2:
        del sess_clicks[s]
        del sess_date[s]

iid_counts = {}
for s in sess_clicks:
    seq = sess_clicks[s]
    for iid in seq:
        if iid in iid_counts:
            iid_counts[iid] += 1
        else:
            iid_counts[iid] = 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] >= 5, curseq))
    if len(filseq) < 2 or len(filseq) > 40:
        del sess_clicks[s]
        del sess_date[s]
    else:
        sess_clicks[s] = filseq

# 分割训练和测试集部分保持不变
dates = list(sess_date.items())
maxdate = dates[0][1]
for _, date in dates:
    if maxdate < date:
        maxdate = date

splitdate = maxdate - 100
print('Splitting date', splitdate)
tra_sess = list(filter(lambda x: x[1] < splitdate, dates))
tes_sess = list(filter(lambda x: x[1] > splitdate, dates))

tra_sess = sorted(tra_sess, key=operator.itemgetter(1))
tes_sess = sorted(tes_sess, key=operator.itemgetter(1))
print(len(tra_sess))
print(len(tes_sess))
print(tra_sess[:3])
print(tes_sess[:3])
print("-- Splitting train set and test set @ %ss" % datetime.datetime.now())

item_dict = {}

def process_seqs(iseqs, idates):
    # 先处理子序列
    out_seqs = []
    out_dates = []
    labs = []
    ids = []
    for id, seq, date in zip(range(len(iseqs)), iseqs, idates):
        for i in range(1, len(seq)):
            sub_seq = seq[:-i]
            tar = seq[-i]
            labs += [tar]
            out_seqs += [sub_seq]
            out_dates += [date]
            ids += [id]
    return out_seqs, out_dates, labs, ids

def obtian_tra():
    train_ids = []
    train_seqs = []
    train_dates = []
    item_ctr = 1
    all_train_seqs = []
    
    # 先转换物品ID
    for s, date in tra_sess:
        seq = sess_clicks[s]
        outseq = []
        for i in seq:
            if i in item_dict:
                outseq += [item_dict[i]]
            else:
                outseq += [item_ctr]
                item_dict[i] = item_ctr
                item_ctr += 1
        train_ids += [s]
        train_dates += [date]
        train_seqs += [outseq]
        all_train_seqs.append(outseq)
    
    # 处理子序列
    tr_seqs, tr_dates, tr_labs, tr_ids = process_seqs(train_seqs, train_dates)
    
    # 计算上下文相似度
    train_contexts = []
    for seq in tr_seqs:
        top_ten = get_top_ten_similar_sessions(seq, all_train_seqs)
        train_contexts.append(top_ten)
    
    print('item_ctr:', item_ctr)
    return tr_ids, tr_dates, tr_seqs, tr_labs, train_contexts

def obtian_tes():
    test_ids = []
    test_seqs = []
    test_dates = []
    all_train_seqs = []
    
    # 获取训练集序列用于上下文
    for s, date in tra_sess:
        seq = sess_clicks[s]
        outseq = []
        for i in seq:
            if i in item_dict:
                outseq += [item_dict[i]]
        all_train_seqs.append(outseq)
    
    # 处理测试集
    for s, date in tes_sess:
        seq = sess_clicks[s]
        outseq = []
        for i in seq:
            if i in item_dict:
                outseq += [item_dict[i]]
        test_ids += [s]
        test_dates += [date]
        test_seqs += [outseq]
    
    # 处理子序列
    te_seqs, te_dates, te_labs, te_ids = process_seqs(test_seqs, test_dates)
    
    # 计算上下文相似度
    test_contexts = []
    for seq in te_seqs:
        top_ten = get_top_ten_similar_sessions(seq, all_train_seqs)
        test_contexts.append(top_ten)
    
    return te_ids, te_dates, te_seqs, te_labs, test_contexts

# 获取训练和测试数据
tr_ids, tr_dates, tr_seqs, tr_labs, tr_contexts = obtian_tra()
te_ids, te_dates, te_seqs, te_labs, te_contexts = obtian_tes()

tra = (tr_seqs, tr_labs, tr_contexts)
tes = (te_seqs, te_labs, te_contexts)

# 输出结果
print('train_test')
print(len(tr_seqs))
print(len(te_seqs))
print(tr_seqs[:3], tr_dates[:3], tr_labs[:3])
print(te_seqs[:3], te_dates[:3], te_labs[:3])

all = 0
for seq in tr_seqs:
    all += len(seq)
for seq in te_seqs:
    all += len(seq)
print('avg length: ', all * 1.0 / (len(tr_seqs) + len(te_seqs)))

# 保存结果
if not os.path.exists('Tmall'):
    os.makedirs('Tmall')
pickle.dump(tra, open('Tmall/train.txt', 'wb'))
pickle.dump(tes, open('Tmall/test.txt', 'wb'))
pickle.dump(tr_seqs, open('Tmall/all_train_seq.txt', 'wb'))