import pandas as pd
import matplotlib.pyplot as plt
clean = pd.read_csv('online_retail_clean.csv', parse_dates=['InvoiceDate'])
clean['CustomerID'] = clean['CustomerID'].astype(int)
clean.shape

clean['InvoiceMonth'] = clean['InvoiceDate'].values.astype('datetime64[M]')
clean[['InvoiceDate', 'InvoiceMonth']].head(3)

first_purchase = clean.groupby('CustomerID')['InvoiceMonth'].min()
first_purchase = first_purchase.rename('CohortMonth')
first_purchase.head()

merged = clean.merge(first_purchase, on='CustomerID')
merged[['CustomerID', 'InvoiceMonth', 'CohortMonth']].head(3)

year_diff = merged['InvoiceMonth'].dt.year - merged['CohortMonth'].dt.year
month_diff = merged['InvoiceMonth'].dt.month - merged['CohortMonth'].dt.month
merged['CohortIndex'] = year_diff * 12 + month_diff
merged['CohortIndex'].describe()

cohort_data = merged.groupby(['CohortMonth', 'CohortIndex'])['CustomerID'].nunique().reset_index()
cohort_data.head()

cohort_pivot = cohort_data.pivot(index='CohortMonth', columns='CohortIndex', values='CustomerID')
cohort_pivot.shape
cohort_pivot.iloc[:5, :6]

cohort_sizes = cohort_pivot.iloc[:, 0]
cohort_sizes

retention = cohort_pivot.divide(cohort_sizes, axis=0)
retention.round(3).iloc[:6, :6]

month1_retention = retention[1].dropna().sort_index()
month1_retention.round(3)
month1_retention.mean().round(3)

plt.figure(figsize=(11, 7))
plt.imshow(retention, cmap='YlGnBu', aspect='auto', vmin=0, vmax=1)
plt.colorbar(label='Real Retention Rate')
plt.yticks(range(len(retention.index)), [d.strftime('%Y-%m') for d in retention.index])
plt.xlabel('Real Months Since First Purchase')
plt.title('Real Customer Retention Heatmap by Cohort')
plt.tight_layout()
plt.savefig('cohort_retention_heatmap.png', dpi=120)
plt.close()

plt.figure(figsize=(9, 6))
for cohort in retention.index[:6]:
    plt.plot(retention.columns, retention.loc[cohort], marker='o', label=cohort.strftime('%Y-%m'))
plt.xlabel('Real Months Since First Purchase')
plt.ylabel('Real Retention Rate')
plt.title('Real Retention Curves for the First Six Cohorts')
plt.legend()
plt.tight_layout()
plt.savefig('cohort_retention_curves.png', dpi=120)
plt.close()

revenue_data = merged.groupby(['CohortMonth', 'CohortIndex'])['Revenue'].sum().reset_index()
revenue_pivot = revenue_data.pivot(index='CohortMonth', columns='CohortIndex', values='Revenue')
revenue_pivot.iloc[:5, :6].round(0)

avg_revenue_per_active = revenue_pivot / cohort_pivot
avg_revenue_per_active.iloc[:5, :6].round(2)

best_cohort = month1_retention.idxmax()
worst_cohort = month1_retention.idxmin()
print(f'Best real cohort: {best_cohort.strftime("%Y-%m")} at {round(month1_retention[best_cohort]*100,1)}% month-1 retention')
print(f'Weakest real cohort: {worst_cohort.strftime("%Y-%m")} at {round(month1_retention[worst_cohort]*100,1)}% month-1 retention')

early_cohorts = month1_retention.iloc[:len(month1_retention)//2]
late_cohorts = month1_retention.iloc[len(month1_retention)//2:]
round(early_cohorts.mean(), 3), round(late_cohorts.mean(), 3)

month6_retention = retention[6].dropna()
month6_retention.round(3)
month6_retention.mean().round(3)

customer_months = cohort_pivot.sum(axis=1)
customer_months.sort_values(ascending=False).round(0)

manual_check = merged[(merged['CohortMonth'] == cohort_pivot.index[0]) & (merged['CohortIndex'] == 1)]['CustomerID'].nunique()
manual_check == cohort_pivot.iloc[0, 1]

retention.round(4).to_csv('cohort_retention_table.csv')
reloaded = pd.read_csv('cohort_retention_table.csv', index_col=0)
reloaded.shape == retention.shape

months_active = merged.groupby('CustomerID')['CohortIndex'].apply(lambda s: (s == 0).sum() == len(s))
one_time_pct = months_active.mean()
round(one_time_pct * 100, 1)

plt.figure(figsize=(9, 5))
plt.bar([d.strftime('%Y-%m') for d in cohort_sizes.index], cohort_sizes.values, color='steelblue')
plt.xticks(rotation=45)
plt.ylabel('Real New Customers')
plt.title('Real New Customer Acquisition by Cohort Month')
plt.tight_layout()
plt.savefig('cohort_acquisition_sizes.png', dpi=120)
plt.close()

orders_per_customer = clean.groupby('CustomerID')['InvoiceNo'].nunique()
repeat_rate = (orders_per_customer > 1).mean()
round(repeat_rate * 100, 1)

repeat_customers = orders_per_customer[orders_per_customer > 1].index
first_two = clean[clean['CustomerID'].isin(repeat_customers)].sort_values('InvoiceDate')
def gap_to_second_visit(dates):
    uniq = sorted(dates.unique())
    return (uniq[1] - uniq[0]) / pd.Timedelta(days=1) if len(uniq) > 1 else None
gap_days = first_two.groupby('CustomerID')['InvoiceDate'].apply(gap_to_second_visit).dropna()
gap_days.median().round(1)

month0_month1_customers = merged[(merged['CohortIndex'] == 1)]['CustomerID'].unique()
total_revenue_per_customer = clean.groupby('CustomerID')['Revenue'].sum()
returners_revenue = total_revenue_per_customer.loc[total_revenue_per_customer.index.isin(month0_month1_customers)].mean()
non_returners_revenue = total_revenue_per_customer.loc[~total_revenue_per_customer.index.isin(month0_month1_customers)].mean()
round(returners_revenue, 2), round(non_returners_revenue, 2)

matured = retention[retention[6].notna()]
matured.shape[0]
matured[[0, 1, 3, 6]].round(3)

largest_cohort_month = cohort_sizes.idxmax()
largest_cohort_customers = first_purchase[first_purchase == largest_cohort_month].index
clean[clean['CustomerID'].isin(largest_cohort_customers)]['Country'].value_counts().head(5)

gap = round((returners_revenue - non_returners_revenue) / non_returners_revenue * 100, 1)
print(f'Customers retained through month one are worth {gap}% more in real lifetime revenue on average, making early-lifecycle retention campaigns a genuinely high-value real investment.')

total_from_pivot = cohort_pivot.iloc[:, 0].sum()
total_unique_customers = clean['CustomerID'].nunique()
total_from_pivot == total_unique_customers

print(f'Analyzed {len(cohort_pivot)} real monthly cohorts covering {total_unique_customers} real unique customers.')