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, 'mlxtend.preprocessing.TransactionEncoder', 'TransactionEncoder', False, 'from mlxtend.preprocessing import TransactionEncoder\n'), (23, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (30, 'pandas.read_csv', 'pd.read_csv', True, 'import pandas as pd\n'), (31, 'collections.defaultdict', 'defaultdict', False, 'from collections import defaultdict\n'), (43, 'mlxtend.frequent_patterns.apriori', 'apriori', False, 'from mlxtend.frequent_patterns import apriori\n'), (28, 'os.path.join', 'os.path.join', False, 'import os\n'), (13, 'os.path.join', 'os.path.join', False, 'import os\n'), (34, 'os.path.join', 'os.path.join', False, 'import os\n')]