def loadDataSet():
    dataset = [[1,3,4],[2,3,5],[1,2,3,5],[2,5]]
    return dataset

#dataSet = loadDataSet()

def createC1(dataSet):
    '''
    生成第一個候選集合C1
    參數(shù)：
    dataset：原始數(shù)據(jù)集
    返回：
    frozenset形式的候選集合C1
    '''
    C1 = []
    for transaction in dataSet:
        for item in transaction:
            if not {item} in C1:
                C1.append({item})
    C1.sort()
    return list(map(frozenset,C1))
#print(C1)


#計算支持度，并保留支持度大于最小支持度的項集
def scanD(D,Ck,minSupport):
    '''
    生成滿足最小支持度的頻繁項集L1
    參數(shù)：
    D:原始數(shù)據(jù)集
    Ck：候選項集
    minsupport：最小支持度
    返回：
    retlist：頻繁項集
    supportdata：候選項集的支持度
    '''
    ssCnt = {}
    for tid in D:
        for can in Ck:
            if can.issubset(tid):   #判斷can是否是tid的子集
                if can not in ssCnt.keys():
                    ssCnt[can]=1
                else:
                    ssCnt[can]+=1
    numItems = float(len(D))
    retList = []   # 頻繁項集
    supportData = {}  #候選項集Ck的支持度字典（key：候選項，value：支持度）
    for key in ssCnt:
        support = ssCnt[key]/numItems  #支持度
        supportData[key] = support
        if support >= minSupport:
            retList.append(key)
    return retList,supportData



#print("這是L1",L1,supportData)


'''
當(dāng)集合中項的個數(shù)大于0時：
    構(gòu)建一個k項集組成的列表
    檢查數(shù)據(jù)確保每個項集都是頻繁的
    保留頻繁項集并構(gòu)建（k+1）項集組成列表
'''

#項集擴展
def aprioriGen(Lk,k):
    Ck = []
    lenLk = len(Lk)
    for i in range(lenLk):
        for j in range(i+1,lenLk):
            #前k-2個項相同時，將兩個集合合并
            L1 = list(Lk[i])[:k-2]
            L1.sort()
            L2 = list(Lk[j])[:k-2]
            L2.sort()
            if L1==L2:
                Ck.append(Lk[i]|Lk[j])
    return Ck



def apriori(D,minSupport=0.5):
    C1 = createC1(D)
    L1,supportData = scanD(D,C1,minSupport)
    L = [L1]
    k = 2
    while (len(L[k-2])>0):
        Ck = aprioriGen(L[k-2],k)
        Lk,supK = scanD(D,Ck,minSupport)
        supportData.update(supK)
        L.append(Lk)
        k+=1
    return L,supportData



#print(L,supportData)



def generate_big_rules(L, supportData, min_conf):
    """
    從頻繁項集產(chǎn)生關(guān)聯(lián)規(guī)則.
    參數(shù):
        L: 候選項集列表
        supportData:頻繁項集及其支持度
    返回:
        big_rule_list: 關(guān)聯(lián)規(guī)則
    """
    big_rule_list = []
    sub_set_list = []
    for i in range(0, len(L)):
        for freq_set in L[i]:   #遍歷第i個頻繁項集列表
            for sub_set in sub_set_list:
                if sub_set.issubset(freq_set):
                    conf = supportData[freq_set] / supportData[freq_set - sub_set]  #置信度等于支持度相除
                    big_rule = (freq_set - sub_set, sub_set, conf)
                    if conf >= min_conf and big_rule not in big_rule_list:
                        # print freq_set-sub_set, " => ", sub_set, "conf: ", conf
                        big_rule_list.append(big_rule)
            sub_set_list.append(freq_set)
    return big_rule_list

# big_rules_list = generate_big_rules(L,supportData,min_conf=0.4)
# print(big_rules_list)

if __name__ == "__main__":
    """
    Test
    """
    dataSet = loadDataSet()
    C1 = createC1(dataSet)
    L,supportData = apriori(dataSet,minSupport=0.5)
    #L1,supportData = scanD(dataSet,C1,0.5)
    big_rules_list = generate_big_rules(L, supportData, min_conf=0.4)
    for Lk in L:
        if len(list(Lk)) == 0:
            break
        print("="*50)
        print("frequent " + str(len(list(Lk)[0])) + "-itemsets\t\tsupport")
        print("="*50)
        for freq_set in Lk:
            print(freq_set, supportData[freq_set])
    print()
    print("Big Rules")
    for item in big_rules_list:
        print(item[0], "=>", item[1], "conf: ", item[2])