import warnings
warnings.filterwarnings('ignore') # suppress warnings

# Data setup (Groceries Dataset)
import pandas as pd
import kagglehub, os
import numpy as np
np.random.seed(42)
path = kagglehub.dataset_download('heeraldedhia/groceries-dataset')
files = os.listdir(path)
csv_file = [f for f in files if f.endswith('.csv')][0]
df = pd.read_csv(os.path.join(path, csv_file))
print(df.shape)
print(df.head(3))

# See unique columns and types
print(df.dtypes)
print(df.columns)

# Check for missing values
print(df.isnull().sum())

# Drop any rows with missing data (if any)
df = df.dropna()
print(df.shape)

# How many unique shoppers and unique items are there?
print('Unique shoppers:', df['Member_number'].nunique())
print('Unique items:', df['itemDescription'].nunique())

# Busiest shopping dates
print(df['Date'].value_counts().head(5))

# Create baskets: one set per shopper
basket_df = df.groupby(['Member_number'])['itemDescription'].apply(list).reset_index()
basket_df['basket_size'] = basket_df['itemDescription'].apply(len)
print(basket_df.head(3))

# How big are most baskets?
print(basket_df['basket_size'].describe())

# What are the top 10 most common items?
all_items = df['itemDescription']
print(all_items.value_counts().head(10))

from mlxtend.preprocessing import TransactionEncoder

# Convert baskets to dummy matrix for apriori
te = TransactionEncoder()
te_ary = te.fit(basket_df['itemDescription']).transform(basket_df['itemDescription'])
basket_matrix = pd.DataFrame(te_ary, columns=te.columns_)
print(basket_matrix.shape)
print(basket_matrix.head(3))

from mlxtend.frequent_patterns import apriori, association_rules

# Find frequent itemsets
frequent_itemsets = apriori(basket_matrix, min_support=0.02, use_colnames=True)
rules = association_rules(frequent_itemsets, metric='lift', min_threshold=1.2)
print(rules[['antecedents', 'consequents', 'support', 'confidence', 'lift']].head(5))

# Visualize top 5 rules by confidence
import matplotlib.pyplot as plt

top_rules = rules.sort_values('confidence', ascending=False).head(5)
plt.barh(range(5), top_rules['confidence'], tick_label=[str(list(a)[0]) + '' + str(list(c)[0]) for a, c in zip(top_rules['antecedents'], top_rules['consequents'])])
plt.xlabel('Confidence')
plt.title('Top 5 Association Rules')
plt.show()

# Add label for whole milk
basket_matrix['buys_whole_milk'] = basket_matrix['whole milk']

X = basket_matrix.drop('buys_whole_milk', axis=1)
y = basket_matrix['buys_whole_milk']

# Split into train and test sets
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)
print(X_train.shape, X_test.shape)

# Train a simple classifier (Logistic Regression)
from sklearn.linear_model import LogisticRegression

lr = LogisticRegression(max_iter=200)
lr.fit(X_train, y_train)
print('Test accuracy:', lr.score(X_test, y_test))

# Pick top 10 frequent items for clustering
cols = all_items.value_counts().head(10).index.tolist()
sample_basket = basket_matrix[cols]
from sklearn.cluster import KMeans
kmeans = KMeans(n_clusters=3, random_state=42)
clusters = kmeans.fit_predict(sample_basket)
sample_basket['cluster'] = clusters
print(sample_basket.groupby('cluster').mean())

# Try it: predict another item's purchase
item = input('Enter an item name to predict (e.g., soda): ')
if item in basket_matrix.columns:
    from sklearn.linear_model import LogisticRegression
    y2 = basket_matrix[item]
    X2 = basket_matrix.drop([item, 'buys_whole_milk'], axis=1, errors='ignore')
    X2_train, X2_test, y2_train, y2_test = train_test_split(X2, y2, test_size=0.25, random_state=42)
    lr2 = LogisticRegression(max_iter=100)
    lr2.fit(X2_train, y2_train)
    print(f'Accuracy in predicting {item}:', lr2.score(X2_test, y2_test))
else:
    print('Item not found. Try a different name.')
    

# Explore: how many clusters work best?
for k in range(2, 6):
    kmeans_k = KMeans(n_clusters=k, random_state=42).fit(sample_basket[cols])
    print(f'k={k}: Inertia={kmeans_k.inertia_:.2f}')