import pandas as pd
import os
from tqdm import tqdm
from collections import defaultdict
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori

dataPath = "data/static"
itemSetList = []


def loadDataSet():
    with open(os.path.join(dataPath, "aprioriData.csv"), 'r') as f:
        for line in f.readlines():
            line = line.replace('\n', '')
            cates = line.split(' ')
            itemSetList.append(list(map(int, cates)))


def myApriori():
    te = TransactionEncoder()
    te_ary = te.fit(itemSetList).transform(itemSetList)
    df = pd.DataFrame(te_ary, columns=te.columns_)
    return df


def dataInit():
    if os.path.exists(os.path.join(dataPath, "aprioriData.csv")):
        return
    df = pd.read_csv("data/static/static.csv")
    user_category = defaultdict(set)
    for idx, row in tqdm(df.iterrows(), total=df.shape[0], desc="category data generate"):
        user_category[row['USER_ID']].add(row['CATEGORY_ID'])
    with open(os.path.join(dataPath, "aprioriData.csv"), 'w+') as f:
        for k, v in tqdm(user_category.items()):
            f.write(' '.join(sorted(list(map(str, v))))+'\n')


if __name__ == '__main__':
    dataInit()
    loadDataSet()
    df = myApriori()
    frequent_itemsets = apriori(df, min_support=0.0035, use_colnames=True)
    frequent_itemsets['length'] = frequent_itemsets['itemsets'].apply(lambda x: len(x))
    print(frequent_itemsets[(frequent_itemsets['length'] >= 2)])

[((21, 9, 21, 29), 'mlxtend.preprocessing.TransactionEncoder', 'TransactionEncoder', ({}, {}), '()', False, 'from mlxtend.preprocessing import TransactionEncoder\n'), ((23, 9, 23, 50), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((30, 9, 30, 46), 'pandas.read_csv', 'pd.read_csv', ({(30, 21, 30, 45): '"""data/static/static.csv"""'}, {}), "('data/static/static.csv')", True, 'import pandas as pd\n'), ((31, 20, 31, 36), 'collections.defaultdict', 'defaultdict', ({(31, 32, 31, 35): 'set'}, {}), '(set)', False, 'from collections import defaultdict\n'), ((43, 24, 43, 74), 'mlxtend.frequent_patterns.apriori', 'apriori', (), '', False, 'from mlxtend.frequent_patterns import apriori\n'), ((28, 22, 28, 63), 'os.path.join', 'os.path.join', ({(28, 35, 28, 43): 'dataPath', (28, 45, 28, 62): '"""aprioriData.csv"""'}, {}), "(dataPath, 'aprioriData.csv')", False, 'import os\n'), ((13, 14, 13, 55), 'os.path.join', 'os.path.join', ({(13, 27, 13, 35): 'dataPath', (13, 37, 13, 54): '"""aprioriData.csv"""'}, {}), "(dataPath, 'aprioriData.csv')", False, 'import os\n'), ((34, 14, 34, 55), 'os.path.join', 'os.path.join', ({(34, 27, 34, 35): 'dataPath', (34, 37, 34, 54): '"""aprioriData.csv"""'}, {}), "(dataPath, 'aprioriData.csv')", False, 'import os\n')]