# Experiment 9: Apriori Algorithm for Association Rule Mining
# Requirements: pip install pandas mlxtend

# ==========================================
# STEP 1: Import Libraries & Create Dataset
# ==========================================
import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

transactions = [
    ['Milk', 'Bread', 'Butter'],
    ['Bread', 'Butter'],
    ['Milk', 'Bread'],
    ['Milk', 'Bread', 'Butter', 'Eggs'],
    ['Bread', 'Eggs'],
    ['Milk', 'Bread', 'Eggs'],
    ['Milk', 'Butter'],
    ['Bread', 'Butter', 'Eggs'],
    ['Milk', 'Bread', 'Butter'],
    ['Bread', 'Milk']
]

# Convert transactions into one-hot encoded format
te = TransactionEncoder()
te_array = te.fit(transactions).transform(transactions)
df = pd.DataFrame(te_array, columns=te.columns_)

print("--- Transaction Dataset ---")
print(df.astype(int))

# ==========================================
# STEP 2: Generate Frequent Itemsets
# ==========================================
min_support = 0.30

frequent_itemsets = apriori(
    df,
    min_support=min_support,
    use_colnames=True
)

frequent_itemsets['itemset_size'] = frequent_itemsets[
    'itemsets'
].apply(len)

print("\n--- Frequent Itemsets ---")
print(frequent_itemsets.sort_values(
    by=['itemset_size', 'support'],
    ascending=[True, False]
))

# ==========================================
# STEP 3: Generate Association Rules
# ==========================================
min_confidence = 0.60

rules = association_rules(
    frequent_itemsets,
    metric="confidence",
    min_threshold=min_confidence
)

rules = rules.sort_values(
    by=['lift', 'confidence'],
    ascending=False
)

print("\n--- Association Rules ---")
print(rules[
    ['antecedents', 'consequents',
     'support', 'confidence', 'lift']
])

# ==========================================
# STEP 4: Display Strong Rules
# ==========================================
strong_rules = rules[
    (rules['confidence'] >= min_confidence) &
    (rules['lift'] > 1)
]

print("\n--- Strong Positive Association Rules ---")
print(strong_rules[
    ['antecedents', 'consequents',
     'support', 'confidence', 'lift']
])
