import pandas as pd
import numpy as np
import warnings
warnings.filterwarnings('ignore')

# Beginner Example 1: Load market basket data for bundle analysis
np.random.seed(42)
products = ['Bread','Milk','Butter','Eggs','Apples','Chicken','Rice','Cheese']
transaction_ids = np.repeat(np.arange(1,301),3)
product_choices = np.random.choice(products,len(transaction_ids))
basket_df = pd.DataFrame({'TransactionID':transaction_ids,'Product':product_choices})
print(basket_df.shape)
print(basket_df.head(5))

# Beginner Example 2: Count unique product combinations per transaction
combinations_per_basket = basket_df.groupby('TransactionID')['Product'].apply(set)
print(combinations_per_basket.head(5))

# Beginner Example 3: Find transaction counts for each product
product_transaction_counts = basket_df.groupby('Product')['TransactionID'].nunique()
print(product_transaction_counts.sort_values(ascending=False))

# Beginner Example 4: Create a product co-occurrence matrix
from collections import Counter

pairs_counter = Counter()
for items in combinations_per_basket:
    for prod1 in items:
        for prod2 in items:
            if prod1 < prod2:
                pairs_counter[(prod1, prod2)] += 1
sorted_pairs = sorted(pairs_counter.items(), key=lambda x: -x[1])
print('Top 5 most common product pairs:')
for pair, count in sorted_pairs[:5]:
    print(f'{pair}: {count} baskets')

# Beginner Example 5: Visualizing simple bundle frequency
import matplotlib.pyplot as plt

top_pairs = sorted_pairs[:5]
labels = [f'{a} & {b}' for (a,b),_ in top_pairs]
counts = [count for _, count in top_pairs]
plt.bar(labels, counts, color='skyblue')
plt.ylabel('Count of Baskets')
plt.title('Most Frequent Product Bundles')
plt.show()

# Intermediate Example 1: Use the real Online Retail Transactions Dataset
url = 'https://archive.ics.uci.edu/ml/machine-learning-databases/00502/online_retail_II.xlsx'
retail_df = pd.read_excel(url, sheet_name='Year 2010-2011')
retail_df['InvoiceDate'] = pd.to_datetime(retail_df['InvoiceDate'])
print(retail_df.shape)
print(retail_df.head(3))

# Intermediate Example 2: Sample data for faster analysis
sample_df = retail_df.sample(n=10000, random_state=42)
print(sample_df.head(3))

# Intermediate Example 3: Find top 10 products by number of transactions
top_products = sample_df.groupby('Description')['Invoice'].nunique().sort_values(ascending=False).head(10)
print(top_products)

# Intermediate Example 4: Calculate revenue by product pair
# First, aggregate all products per invoice
invoice_products = sample_df.groupby('Invoice')['Description'].apply(set)
pair_revenue = Counter()
for idx, group in sample_df.groupby('Invoice'):
    prods = set(group['Description'])
    total_revenue = (group['Quantity'] * group['Price']).sum()
    for prod1 in prods:
        for prod2 in prods:
            if prod1 < prod2:
                pair_revenue[(prod1, prod2)] += total_revenue
top_revenue_pairs = sorted(pair_revenue.items(), key=lambda x: -x[1])[:5]
print('Top 5 product pairs by combined invoice revenue:')
for pair, revenue in top_revenue_pairs:
    print(pair, '->', round(revenue,2))

# Intermediate Example 5: Cross-sell ratio for a lead product ('WHITE HANGING HEART T-LIGHT HOLDER')
lead = 'WHITE HANGING HEART T-LIGHT HOLDER'
with_lead = invoice_products[invoice_products.apply(lambda x: lead in x)]
cross_counts = Counter()
for prods in with_lead:
    for prod in prods:
        if prod != lead:
            cross_counts[prod] += 1
cross_sells = cross_counts.most_common(5)
print(f'Top 5 products purchased with {lead}:')
for prod, count in cross_sells:
    print(f'{prod}: {count} baskets')

# Intermediate Example 6: Visualize cross-sell counts for lead product
labels = [prod for (prod,_) in cross_sells]
counts = [count for (_,count) in cross_sells]
plt.bar(labels, counts, color='orange')
plt.title('Top Cross-Sell Items with Lead Product')
plt.ylabel('Basket Count')
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.show()

# Advanced Example 1: Create a pivot table of top cross-sold products per main item
pivot_counts = {}
for lead in top_products.index[:5]:
    cross_counts = Counter()
    with_lead = invoice_products[invoice_products.apply(lambda x: lead in x)]
    for prods in with_lead:
        for prod in prods:
            if prod != lead:
                cross_counts[prod] += 1
    pivot_counts[lead] = cross_counts.most_common(3)
for lead, pairs in pivot_counts.items():
    print(f'{lead}:')
    for prod, count in pairs:
        print(f'   {prod}: {count}')

# Advanced Example 2: Compute lift metric for product pairs
def compute_lift(pair, basket_sets, prods_list):
    prod_a, prod_b = pair
    num_baskets = len(basket_sets)
    a_count = sum([prod_a in items for items in basket_sets])
    b_count = sum([prod_b in items for items in basket_sets])
    ab_count = sum([(prod_a in items) and (prod_b in items) for items in basket_sets])
    if a_count==0 or b_count==0:
        return 0
    return (ab_count / num_baskets) / ((a_count / num_baskets) * (b_count / num_baskets))
all_pairs = [pair for (pair,_) in sorted_pairs[:10]]
lifts = [(pair, compute_lift(pair, combinations_per_basket, products)) for pair in all_pairs]
lifts.sort(key=lambda x: -x[1])
print('Top product bundles by lift:')
for pair, lift in lifts[:5]:
    print(f'{pair}: lift={round(lift,2)}')

# Advanced Example 3: Association rule mining using mlxtend (if installed)
try:
    from mlxtend.frequent_patterns import apriori, association_rules
    # Prepare transaction-basket format
    basket_sets = basket_df.groupby('TransactionID')['Product'].apply(list)
    from mlxtend.preprocessing import TransactionEncoder
    te = TransactionEncoder()
    te_ary = te.fit(basket_sets.tolist()).transform(basket_sets.tolist())
    basket_ready = pd.DataFrame(te_ary, columns=te.columns_)
    frequent_itemsets = apriori(basket_ready, min_support=0.05, use_colnames=True)
    rules = association_rules(frequent_itemsets, metric='lift', min_threshold=1)
    print(rules[['antecedents','consequents','support','confidence','lift']].head())
except Exception as e:
    print('mlxtend not installed. Please install with: pip install mlxtend')

# Error Handling Example 1: Check for missing products in baskets
missing_products = basket_df['Product'].isnull().sum()
print(f'Missing product entries: {missing_products}')

# Error Handling Example 2: Detect empty baskets after filtering
baskets_with_no_products = combinations_per_basket.apply(lambda s: len(s)==0).sum()
print(f'Number of transactions with no products: {baskets_with_no_products}')

# Error Handling Example 3: Incorrect product grouping bug example
wrong_group = basket_df.groupby('Product')['TransactionID'].count()
correct_group = basket_df.groupby('Product')['TransactionID'].nunique()
print('Count:', wrong_group.max(), 'Unique:', correct_group.max())

# End-to-end Example: Find and recommend a high-revenue cross-sell bundle
lead_item = top_products.index[0]
with_lead = invoice_products[invoice_products.apply(lambda x: lead_item in x)]
cross_counter = Counter()
for prods in with_lead:
    for prod in prods:
        if prod != lead_item:
            cross_counter[prod] += 1
top_partner, max_count = cross_counter.most_common(1)[0]
recommend_pair = (lead_item, top_partner)
bundled_total = 0.0
for idx, group in sample_df.groupby('Invoice'):
    items = set(group['Description'])
    if recommend_pair[0] in items and recommend_pair[1] in items:
        bundled_total += (group['Quantity'] * group['Price']).sum()
print(f'Recommend bundle: {lead_item} + {top_partner}')
print(f'Estimated bundle revenue in sample: {bundled_total:.2f}')