##############################################################
# Chapter 16 - Using Marketing Experiments to Optimize the   #
#              Marketing Mix                                 #
# Chestnut Ridge Email Campaign Causal Random Forest         #
# (Python)                                                   #
##############################################################

# This script shows another way to estimate the effect of the Chestnut Ridge email
# campaign in Python. It uses the same data as the propensity score matching example
# (propensity_score_matching.py), but estimates the effect with a causal random forest.
# NOTE: This approach is NOT covered in the book's examples. It is provided as an
# optional extension for comparison with the propensity score matching results.
#   Step 1. Install packages (if needed), load packages, and set seed
#   Step 2. Read in the email marketing campaign data
#   Step 3. Train a causal random forest
#   Step 4. Estimate the average treatment effect of the email
# Run it from top to bottom. Each step prints its result.
# 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 = {"numpy": "numpy", "pandas": "pandas", "sklearn": "scikit-learn", "econml": "econml"}
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 and Set Seed
from tkinter import Tk, filedialog

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from econml.dml import CausalForestDML

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          #
##############################################################

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

## Number of customers
print(len(data))


# %%
##############################################################
# Step 3: Train a causal random forest                       #
##############################################################

## X: customer characteristics (the same covariates used for matching)
## W: treatment (1 = received the email, 0 = no email)
## Y: outcome (purchase amount)
x_vars = ["revenue", "number_of_orders", "number_of_orders2", "recency_days"]
X = data[x_vars]
W = data["email"]
Y = data["purchase_amt"]

## Like causal_forest() in R's grf package, the forest first predicts purchase amount (Y)
## and the chance of getting the email (W) from X, then grows 2,000 "honest" trees on
## what those predictions leave unexplained. This can take a minute or two.
email_forest = CausalForestDML(
    model_y=RandomForestRegressor(min_samples_leaf=5, random_state=1),
    model_t=RandomForestClassifier(min_samples_leaf=5, random_state=1),
    discrete_treatment=True,
    n_estimators=2000,
    min_samples_leaf=5,
    cv=5,
    random_state=1)
email_forest.fit(Y, W, X=X)


# %%
##############################################################
# Step 4: Estimate the average treatment effect of the email #
##############################################################

## Average treatment effect (ATE) on the full sample, with its standard error
## (the same doubly robust estimate as average_treatment_effect(target.sample = "all") in R)
ate = pd.Series({"estimate": email_forest.ate_[0], "std.err": email_forest.ate_stderr_[0]})
print(ate.round(4))

## 95% confidence interval for the ATE
print(email_forest.ate__inference().summary_frame(alpha=0.05).round(4))