import argparse
import numpy as np
RATING_FILE_NAME=dict({'movie':'ratings.dat','book':'BX-Book-Ratings.csv','news':'ratings.txt'})
SEP=dict({'movie':'::','book':';','news':'\t'})
THRESHOLD=dict({'movie':4,'book':0,'news':0})

def read_item_index_to_entity_id_file():
    file='../data/'+DATASET+'/item_index2entity_id_rehashed.txt'
    print('reading item index to entity id file:'+file+'...')
    i=0
    for line in open(file,encoding='utf-8').readlines():

        item_index=line.strip().split('\t')[0]
        # print(15,item_index)
        satori_id=line.strip().split('\t')[1]
        # print(17,satori_id)
        item_index_old2new[item_index]=i
        entity_id2index[satori_id]=i
        i+=1


def convert_rating():
    file='../data/'+DATASET+'/'+RATING_FILE_NAME[DATASET]
    # print(24)
    # print(item_index_old2new.values())#商品数
    item_set=set(item_index_old2new.values())
    user_pos_ratings=dict()
    user_neg_ratings=dict()
    for line in open(file,encoding='utf-8').readlines()[1:]:
        # print(line)
        array=line.strip().split(SEP[DATASET])
        if DATASET=='book':
            array=list(map(lambda x: x[1:-1],array))
        item_index_old=array[1]#商品索引
        if item_index_old not in item_index_old2new: #the item is not in the final item set
            continue
        item_index=item_index_old2new[item_index_old]#新的商品索引
        user_index_old=int(array[0])#用户index
        rating=float(array[2])#评分
        #评分大于4作为正用户
        if rating>=THRESHOLD[DATASET]:
            if user_index_old not in user_pos_ratings:
                user_pos_ratings[user_index_old]=set()
            user_pos_ratings[user_index_old].add(item_index)
        else:
            if user_index_old not in user_neg_ratings:
                user_neg_ratings[user_index_old]=set()
            user_neg_ratings[user_index_old].add(item_index)

    writer=open('../data/'+DATASET+'/ratings_final.txt','w',encoding='utf-8')
    user_cnt=0
    user_index_old2new=dict()
    for user_index_old,pos_item_set in user_pos_ratings.items():
        if user_index_old not in user_index_old2new:
            user_index_old2new[user_index_old]=user_cnt
            user_cnt+=1
        user_index=user_index_old2new[user_index_old]
        for item in pos_item_set:
            writer.write('%d\t%d\t1\n' %(user_index,item))
        unwatched_set=item_set-pos_item_set
        if user_index_old in user_neg_ratings:
            unwatched_set-=user_neg_ratings[user_index_old]
        for item in np.random.choice(list(unwatched_set),size=len(pos_item_set),replace=False):
            writer.write('%d\t%d\t0\n'%(user_index,item))
    writer.close()
    print('number of users: %d'%user_cnt)
    print('number of items: %d'%len(item_set))

def convert_kg():
    print('converting kg file..')
    entity_cnt=len(entity_id2index)
    relation_cnt=0
    writer=open('../data/'+DATASET+'/kg_final.txt','w',encoding='utf-8')
    files=[]
    if DATASET=='movie':
        files.append(open('../data/'+DATASET+'/kg_part1_rehashed.txt',encoding='utf-8'))
        files.append(open('../data/'+DATASET+'/kg_part2_rehashed.txt',encoding='utf-8'))
    else:
        files.append(open('../data/'+DATASET+'/kg_rehashed.txt',encoding='utf-8'))

    for file in files:
        for line in file:
            array=line.strip().split('\t')
            head_old=array[0]
            relation_old=array[1]
            # print(relation_old)
            tail_old=array[2]

            if head_old not in entity_id2index:
                entity_id2index[head_old]=entity_cnt
                entity_cnt+=1
            head=entity_id2index[head_old]
            if tail_old not in entity_id2index:
                entity_id2index[tail_old]=entity_cnt
                entity_cnt+=1
            tail=entity_id2index[tail_old]
            if relation_old not in relation_id2index:
                relation_id2index[relation_old]=relation_cnt
                relation_cnt+=1
            relation=relation_id2index[relation_old]
            print(relation_id2index)
            print('-----')
            print(relation)
            writer.write('%d\t%d\t%d\n'%(head,relation,tail))
    writer.close()
    print('number of entities (containing items):%d'%entity_cnt)
    print('number of relations:%d'%relation_cnt)

if __name__ == '__main__':
    '''
    number of users: 6036
    number of items: 2445
    number of entities 182011
    number of relations:12
    '''
    np.random.seed(555)
    parser=argparse.ArgumentParser()
    parser.add_argument('-d','--dataset',type=str,default='movie',help='which dataset to preprocess')
    args=parser.parse_args()
    DATASET=args.dataset
    entity_id2index=dict()
    relation_id2index=dict()
    item_index_old2new=dict()
    read_item_index_to_entity_id_file()
    convert_rating()
    convert_kg()

