##############################################################
# Chapter 16 - Using Marketing Experiments to Optimize the   #
#              Marketing Mix                                 #
# Chestnut Ridge Email Campaign Propensity Score Matching    #
# (Python)                                                   #
##############################################################

# This script reproduces the full Chestnut Ridge email campaign propensity score matching example in Python:
#   Step 1. Install packages (if needed), load packages, and set seed
#   Step 2. Read in the email marketing campaign data
#   Step 3. Look at pre-matched data
#   Step 4. Difference of covariates (pre-matching means)
#   Step 5. Determine optimal caliper size
#   Step 6. Execute the matching algorithm
#   Step 7. Tobit model with the matched data
#   Step 8. Predict the expected value of purchase_amt
#   Step 9. Compare the expected value for the email and no email groups (matched)
#   Step 10. Run the Tobit model with the unmatched data
#   Step 11. Look at the propensity scores before and after matching
#   Step 12. Create the revenue change variables
#   Step 13. Compare predicted and actual revenue gains (Figure 16.2)
#   Step 14. Estimated gain per customer from the email campaign
# 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.

# %%
##############################################################
# Step 1: Install packages (if needed), load packages, and   #
#         set seed                                           #
##############################################################

## 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", "numpy": "numpy", "scipy": "scipy", "statsmodels": "statsmodels", "matplotlib": "matplotlib"}
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
from tkinter import Tk, filedialog

import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
from statsmodels.base.model import GenericLikelihoodModel
from scipy.stats import norm
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

np.random.seed(1)

## Function to choose a file 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


# %%
##############################################################
# Step 2: Read in the email marketing campaign data          #
##############################################################

psmatch = pd.read_csv(file_choose("Choose retail_psmatch.csv"))  ## Choose the file retail_psmatch.csv


# %%
##############################################################
# Step 3: Look at pre-matched data                           #
##############################################################

print(psmatch.groupby("email").agg(n_customers=("customer", "size"),
                                   mean_purchase_amt=("purchase_amt", "mean"))
      .round(2))


# %%
##############################################################
# Step 4: Difference of covariates (pre-matching means)      #
##############################################################

psmatch_keep = ["email", "revenue", "number_of_orders", "number_of_orders2",
                "recency_days"]
psmatch_cov = psmatch[psmatch_keep]
print(psmatch_cov.groupby("email").mean().round(2))


# %%
##############################################################
# Step 5: Determine optimal caliper size                     #
##############################################################

ps_match = smf.logit("email ~ revenue + number_of_orders + number_of_orders2"
                     " + recency_days", data=psmatch).fit(disp=0)
ps_match_df = pd.DataFrame({"pr_score": ps_match.predict(),
                            "email": psmatch["email"]})
print(0.2 * ps_match_df["pr_score"].std())


# %%
##############################################################
# Step 6: Execute the matching algorithm                     #
##############################################################

## Remove any rows with missing values
psmatch_nomiss = psmatch[["customer", "purchase_amt"] + psmatch_keep].dropna()
psmatch_nomiss = psmatch_nomiss.reset_index(drop=True)

## Estimate the propensity score (probability of receiving the email)
ps_model = smf.logit("email ~ revenue + number_of_orders + number_of_orders2"
                     " + recency_days", data=psmatch_nomiss).fit(disp=0)
psmatch_nomiss["distance"] = ps_model.predict()

## Nearest neighbor 1-to-1 matching with a caliper (no replacement)
## This follows the same steps as matchit() in R's MatchIt package:
##  - Email customers are matched in order, largest propensity score first
##  - Each is matched to the closest No Email customer not already matched
##  - No match is made if the closest No Email customer is outside the caliper
## Note: matchit() applies the caliper in standard deviations of the
## propensity score (its std.caliper = TRUE default), so the same is done here.
caliper = 0.023
caliper_width = caliper * psmatch_nomiss["distance"].std()

ps = psmatch_nomiss["distance"].to_numpy()
treated = np.where(psmatch_nomiss["email"] == 1)[0]
controls = np.where(psmatch_nomiss["email"] == 0)[0]
treated = treated[np.argsort(-ps[treated], kind="stable")]
available = np.ones(len(controls), dtype=bool)

pairs = []
for t in treated:
    gap = np.abs(ps[controls] - ps[t])
    gap[~available] = np.inf
    best = np.argmin(gap)
    if gap[best] <= caliper_width:
        available[best] = False
        pairs.append((t, controls[best]))

## Keep the matched customers (subclass identifies each matched pair)
subclass = {}
for pair_id, (t, c) in enumerate(pairs, start=1):
    subclass[t] = pair_id
    subclass[c] = pair_id
matched_data = psmatch_nomiss.loc[sorted(subclass)].copy()
matched_data["weights"] = 1.0
matched_data["subclass"] = [subclass[i] for i in matched_data.index]
print(matched_data.shape)

## Look at post-matching means
matched_data_cov = matched_data[psmatch_keep]
print(matched_data_cov.groupby("email").mean().round(2))


# %%
##############################################################
# Step 7: Tobit model with the matched data                  #
##############################################################

## A Tobit model (left-censored at 0), estimated by maximum likelihood.
## The last parameter is log(sigma), reported as Log(scale) as in R.
class Tobit(GenericLikelihoodModel):
    def loglikeobs(self, params):
        beta, sigma = params[:-1], np.exp(params[-1])
        xb = self.exog @ beta
        y = self.endog
        return np.where(y <= 0,
                        norm.logcdf(-xb / sigma),
                        norm.logpdf((y - xb) / sigma) - np.log(sigma))

def fit_tobit(data):
    y = data["purchase_amt"]
    X = sm.add_constant(data[["email", "revenue", "number_of_orders",
                              "number_of_orders2", "recency_days"]])
    ols = sm.OLS(y, X).fit()
    start = np.append(ols.params, np.log(np.sqrt(ols.scale)))
    model = Tobit(y, X, extra_params_names=["Log(scale)"])
    result = model.fit(start_params=start, method="bfgs", maxiter=2000, disp=0)
    result = model.fit(start_params=result.params, method="newton", disp=0)
    print(f"Observations: Total = {len(y)}, Left-censored = {(y <= 0).sum()},"
          f" Uncensored = {(y > 0).sum()}")
    return result, X

tobit_treat, X_matched = fit_tobit(matched_data)
print(tobit_treat.summary())
print(f"Scale: {np.exp(tobit_treat.params[-1]):.1f}")


# %%
##############################################################
# Step 8: Predict the expected value of purchase_amt         #
##############################################################

def lam(x):
    return norm.pdf(x) / norm.cdf(x)

mu = X_matched @ tobit_treat.params[:-1]
sigma = np.exp(tobit_treat.params[-1])
p0 = norm.cdf(mu / sigma)
ey0 = mu + sigma * lam(mu / sigma)
matched_data["pred_purchase_amt"] = p0 * ey0


# %%
##############################################################
# Step 9: Compare the expected value for the email and no    #
#         email groups (matched)                             #
##############################################################

print(matched_data.groupby("email")["pred_purchase_amt"].mean().round(2))


# %%
##############################################################
# Step 10: Run the Tobit model with the unmatched data       #
##############################################################

tobit_treat_nomatch, X_all = fit_tobit(psmatch)
print(tobit_treat_nomatch.summary())
print(f"Scale: {np.exp(tobit_treat_nomatch.params[-1]):.1f}")

## Predict the expected value for the unmatched data
mu = X_all @ tobit_treat_nomatch.params[:-1]
sigma = np.exp(tobit_treat_nomatch.params[-1])
p0 = norm.cdf(mu / sigma)
ey0 = mu + sigma * lam(mu / sigma)
psmatch["pred_purchase_amt"] = p0 * ey0

## Compare the expected value for the email and no email groups (no match)
print(psmatch.groupby("email")["pred_purchase_amt"].mean().round(2))


# %%
##############################################################
# Step 11: Look at the propensity scores before and after    #
#          matching                                          #
##############################################################

colors = {0: "#2a78d6", 1: "#eb6834"}
labels = {0: "No Email", 1: "Email"}
bins = np.arange(0, 1.025, 0.025)
comma = FuncFormatter(lambda x, pos: f"{x:,.0f}")

fig, axes = plt.subplots(2, 1, sharex=True, figsize=(10, 7))
for ax, data, title in zip(axes, [psmatch_nomiss, matched_data],
                           ["Before Matching", "After Matching"]):
    for group in [0, 1]:
        ax.hist(data.loc[data["email"] == group, "distance"], bins=bins,
                alpha=0.6, color=colors[group], edgecolor="white",
                label=labels[group])
    ax.set_title(title, loc="right", fontsize=10)
    ax.set_ylabel("Number of Customers")
    ax.yaxis.set_major_formatter(comma)
    ax.grid(color="#e5e5e5", linewidth=0.5)
    ax.spines[["top", "right"]].set_visible(False)
axes[0].legend(loc="upper center", ncol=2, frameon=False)
axes[-1].set_xlabel("Propensity Score (Probability of Receiving the Email)")
fig.suptitle("Propensity Scores Before and After Matching", x=0.02, ha="left")
fig.tight_layout()
plt.show()


# %%
##############################################################
# Step 12: Create the revenue change variables               #
##############################################################

## Compare the purchase amount in the campaign month (January 2020) to the
## average monthly spending in the prior two years (revenue / 24 months)
matched_data["revenue_change"] = (matched_data["purchase_amt"]
                                  - matched_data["revenue"] / 24)
matched_data["pred_revenue_change"] = (matched_data["pred_purchase_amt"]
                                       - matched_data["revenue"] / 24)


# %%
##############################################################
# Step 13: Compare predicted and actual revenue gains        #
#          (Figure 16.2)                                     #
##############################################################

measures = ["pred_purchase_amt", "pred_revenue_change", "purchase_amt",
            "revenue_change"]
measure_labels = ["Avg. Pred Purchase Amt", "Avg. Pred Revenue Change",
                  "Avg. Purchase Amt", "Avg. Revenue Change"]
gains = matched_data.groupby("email")[measures].mean()

## Put the measures in rows and the email groups in columns
gains_table = gains.T
gains_table.columns = ["email_0", "email_1"]
gains_table["difference"] = gains_table["email_1"] - gains_table["email_0"]
gains_table.index = measure_labels
print(gains_table.round(2))

## Plot the comparison
x = np.arange(len(measures))
width = 0.35
fig, ax = plt.subplots(figsize=(10, 6))
for offset, group in [(-width / 2, 0), (width / 2, 1)]:
    bars = ax.bar(x + offset, gains.loc[group], width, color=colors[group],
                  label=labels[group])
    ax.bar_label(bars, labels=[f"${v:,.2f}" for v in gains.loc[group]],
                 padding=3, fontsize=9)
ax.set_xticks(x, measure_labels)
ax.set_ylabel("Average per Customer")
ax.yaxis.set_major_formatter(FuncFormatter(lambda v, pos: f"${v:,.0f}"))
ax.set_title("Comparison of Predicted and Actual Revenue Gains", loc="left")
ax.grid(axis="y", color="#e5e5e5", linewidth=0.5)
ax.set_axisbelow(True)
ax.spines[["top", "right"]].set_visible(False)
ax.legend(loc="upper center", ncol=2, frameon=False)
fig.tight_layout()
plt.show()


# %%
##############################################################
# Step 14: Estimated gain per customer from the email        #
#          campaign                                          #
##############################################################

## Difference in average purchase amount (Email - No Email)
print(gains_table.loc["Avg. Purchase Amt", "difference"])
print(gains_table.loc["Avg. Pred Purchase Amt", "difference"])

## Difference-in-differences (Email - No Email change vs. prior monthly average)
print(gains_table.loc["Avg. Revenue Change", "difference"])
print(gains_table.loc["Avg. Pred Revenue Change", "difference"])