##########################################################
# Segmentation Analysis in Python - Chestnut Ridge       #
# Cluster Analysis for Segmentation                      #
##########################################################

# This script reproduces the full Chestnut Ridge segmentation example in Python:
#   Step 1. Visualize average order size by zip code (map)
#   Step 2. Hierarchical clustering and elbow plot to choose the number of segments
#   Step 3. K-means clustering with six segments
#   Step 4. Profile the segments on the bases (attitudes and behaviors)
#   Step 5. Profile the segments on the descriptors (demographics)
#   Step 6. Evaluate segment attractiveness (sales and profit)
# Run it from top to bottom. Each step prints its result or shows a plot.
# The "# %%" lines split the script into cells you can run one at a time
# in VS Code, Spyder, or PyCharm.

# %%
## Install Packages (if needed - only installs packages that are missing)
## Each entry is {name used with import: name used with pip install}
import sys, subprocess, importlib.util
packages = {"pandas": "pandas", "matplotlib": "matplotlib", "scipy": "scipy", "sklearn": "scikit-learn", "geopandas": "geopandas", "pygris": "pygris"}
new_packages = [pip_name for import_name, pip_name in packages.items() if importlib.util.find_spec(import_name) is None]
if new_packages: subprocess.check_call([sys.executable, "-m", "pip", "install", *new_packages])

## Load Packages
import pandas as pd                               # data manipulation
import matplotlib.pyplot as plt                   # plots and maps
from scipy.cluster.hierarchy import linkage       # hierarchical clustering
from sklearn.cluster import KMeans                # k-means clustering
import geopandas as gpd                           # maps of zip code boundaries
from pygris import zctas, states                  # zip code and state boundaries
from tkinter import Tk, filedialog                # file-choosing dialog box

## Functions to choose a file to open, or name a file to save, with a dialog box
## (the Python version of file.choose() in R - no need to set the working directory)
def file_choose(title="Choose a file"):
    root = Tk()
    root.withdraw()                     # hide the empty Tk window
    root.attributes("-topmost", True)   # bring the dialog to the front
    path = filedialog.askopenfilename(title=title, filetypes=[("CSV files", "*.csv"), ("All files", "*.*")])
    root.destroy()
    if not path:
        raise SystemExit("No file was chosen.")
    return path

def file_save(title="Save the file as", initialfile=""):
    root = Tk()
    root.withdraw()
    root.attributes("-topmost", True)
    path = filedialog.asksaveasfilename(title=title, initialfile=initialfile, defaultextension=".csv",
                                        filetypes=[("CSV files", "*.csv")])
    root.destroy()
    if not path:
        raise SystemExit("No file name was chosen.")
    return path

## Set Seed
# Segmentation algorithms start at a random point. Setting the seed ensures
# your results are the same as the results in this example.
SEED = 1

# Import Data
# Read zip codes as text so zip codes that start with 0 (e.g., 08053) keep their leading zero
seg = pd.read_csv(file_choose("Choose retail_segmentation.csv"), dtype={"zip_code": str})  ## Choose retail_segmentation.csv file
seg["zip_code"] = seg["zip_code"].str.zfill(5)

# Take a first look at the data
print(seg.describe().round(2).T)


# %%
#####################################################
# Step 1: Average Order Size by Zip Code (Map)      #
#####################################################

# Average order size for the customers in each zip code
zip_summary = (seg.groupby("zip_code")
                  .agg(avg_order_size=("avg_order_size", "mean"),
                       customers=("Cust_No", "count"))
                  .reset_index())

# Download zip code (ZCTA) and state boundaries from the US Census Bureau
# (the first download may take a few minutes; files are cached for later use)
zip_shapes = zctas(cb=True, year=2020, starts_with=list(zip_summary["zip_code"]), cache=True)
zip_map = zip_shapes.merge(zip_summary, left_on="ZCTA5CE20", right_on="zip_code")
state_map = states(cb=True, year=2020, cache=True)
state_map = state_map[state_map["STUSPS"].isin(["NJ", "PA", "DE", "MD", "DC", "VA", "WV"])].to_crs(zip_map.crs)
xmin, ymin, xmax, ymax = zip_map.total_bounds

# Locations of the two brick-and-mortar stores (a point inside each store's zip code)
stores = pd.DataFrame({"store": ["Northern Virginia (22102)", "Western New Jersey (08053)"],
                       "zip_code": ["22102", "08053"]})
store_points = (zip_map.set_index("zip_code").loc[stores["zip_code"]]
                .representative_point())
stores["lng"] = store_points.x.values
stores["lat"] = store_points.y.values

# Map of average order size by zip code
fig, ax = plt.subplots(figsize=(10, 7))
state_map.plot(ax=ax, color="0.95", edgecolor="0.6")
zip_map.plot(ax=ax, column="avg_order_size", cmap="RdYlBu_r", legend=True,
             legend_kwds={"label": "Avg. Order Size ($)", "shrink": 0.6})
ax.scatter(stores["lng"], stores["lat"], marker="^", s=150, color="black", edgecolor="white")
for _, store in stores.iterrows():
    ax.annotate(store["store"], (store["lng"], store["lat"] + 0.15),
                ha="center", fontsize=10, fontweight="bold")
ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)
ax.set_title("Average Order Size by Zip Code\nTriangles mark the two stores")
ax.set_axis_off()
plt.show()


# %%
#####################################################
# Step 2: Determine the Number of Customer Segments #
#####################################################

# The bases: Past Purchase Behaviors and Marketing Efforts
bases = ["avg_order_size", "avg_order_freq", "crossbuy", "multichannel",
         "per_sale", "tenure", "avg_mktg_cnt", "return_rate"]

# Run hierarchical clustering with bases variables
# Standardizing each basis (subtract the mean, divide by the standard deviation)
# puts all bases on the same scale
bases_scaled = (seg[bases] - seg[bases].mean()) / seg[bases].std()
seg_hclust = linkage(bases_scaled, method="complete", metric="euclidean")

# Elbow plot for first 10 segments
# The third column of the linkage result holds the distance (height) of each merge
elbow = pd.DataFrame({"segments": range(1, 11),
                      "height": sorted(seg_hclust[:, 2], reverse=True)[:10]})

fig, ax = plt.subplots(figsize=(9, 6))
ax.plot(elbow["segments"], elbow["height"], color="blue", marker="o", markerfacecolor="black")
ax.set_xticks(range(1, 11))
ax.set_title("Elbow Plot from the Hierarchical Clustering Algorithm")
ax.set_xlabel("Number of Segments")
ax.set_ylabel("Distance (Height)")
plt.show()

# The kink in the curve around six segments suggests a six-segment solution.


# %%
#####################################################
# Step 3: K-Means Clustering with Six Segments      #
#####################################################

# Run k-means with 6 segments
# n_init=100 runs k-means from 100 random starting points and keeps the best solution;
# tol=0 lets each run continue until no customer changes segment
seg_kmeans = KMeans(n_clusters=6, n_init=100, tol=0, max_iter=1000, random_state=SEED)
seg_kmeans.fit(seg[bases])

# Add segment number back to original data
segmentation_result = seg.copy()
segmentation_result["segment"] = seg_kmeans.labels_ + 1   # number segments 1-6 instead of 0-5

# K-means numbers the segments arbitrarily. To match the segment numbers used in the
# book (and in the R version of this example), rank the segments from smallest to largest
# average order size, then give them the book's numbers in that order.
order_size_rank = (segmentation_result.groupby("segment")["avg_order_size"].mean()
                   .rank().astype(int))
book_numbers = {1: 4, 2: 2, 3: 6, 4: 1, 5: 3, 6: 5}   # rank -> book segment number
segmentation_result["segment"] = (segmentation_result["segment"]
                                  .map(order_size_rank).map(book_numbers))

# Number of customers in each segment
print(segmentation_result["segment"].value_counts().sort_index())

# Export data to a CSV file (used as the input data for the Chapter 4 example)
segmentation_result.to_csv(file_save("Save as segmentation_result.csv", "segmentation_result.csv"),
                           index=False)  ## Name file segmentation_result.csv


# %%
#####################################################
# Step 4: Profile the Segments on the Bases         #
#####################################################

# Sales per year = average order size x average order frequency x (1 - return rate)
# Profit per year = sales per year x margin (52%) - catalogs x cost per catalog ($0.75)
segmentation_result["sales_per_year"] = (segmentation_result["avg_order_size"]
                                         * segmentation_result["avg_order_freq"]
                                         * (1 - segmentation_result["return_rate"]))
segmentation_result["profit_per_year"] = (segmentation_result["sales_per_year"] * 0.52
                                          - segmentation_result["avg_mktg_cnt"] * 0.75)

# Average values of the bases by segment, plus segment size, total sales, and total profit
bases_profile = segmentation_result.groupby("segment")[bases].mean()
bases_profile["segment_size"] = segmentation_result.groupby("segment").size()
bases_profile["total_sales"] = segmentation_result.groupby("segment")["sales_per_year"].sum()
bases_profile["total_profit"] = segmentation_result.groupby("segment")["profit_per_year"].sum()

# Show segments as columns, as in the book
print("\nAverage Values of Bases by Segment")
print(bases_profile.T.round(2).to_string())


# %%
#####################################################
# Step 5: Profile the Segments on the Descriptors   #
#####################################################

descriptors = ["age", "household_size", "income", "loyalty_card", "married", "own_home"]

# Average values of the descriptors by segment, plus segment size
descriptors_profile = segmentation_result.groupby("segment")[descriptors].mean()
descriptors_profile["segment_size"] = segmentation_result.groupby("segment").size()

print("\nAverage Values of Descriptors by Segment")
print(descriptors_profile.T.round(2).to_string())


# %%
#####################################################
# Step 6: Segment Attractiveness                    #
#####################################################

# Revenue and profit from each segment per year
attractiveness = bases_profile[["segment_size", "total_sales", "total_profit"]].copy()
attractiveness["share_of_customers"] = (attractiveness["segment_size"]
                                        / attractiveness["segment_size"].sum())
attractiveness["sales_per_customer"] = (attractiveness["total_sales"]
                                        / attractiveness["segment_size"])
attractiveness["profit_per_customer"] = (attractiveness["total_profit"]
                                         / attractiveness["segment_size"])
attractiveness = attractiveness[["segment_size", "share_of_customers", "total_sales",
                                 "sales_per_customer", "total_profit", "profit_per_customer"]]

print("\nRevenue and Profit from Each Segment per Year")
print(attractiveness.round(2).to_string())