"""
Reproduce: "What is a customer actually worth?"
-----------------------------------------------
End-to-end script for the customer-value investigation on
Data and Other Dangerous Things.

Rebuilds the synthetic dataset (25k customers), runs descriptive stats,
a pairwise correlation matrix, simple vs multiple regression, and VIF
diagnostics for multicollinearity.

Run:
    pip install numpy pandas scikit-learn statsmodels
    python customer-value.py

Read-only: do not modify this file in place — copy it first.
"""
from __future__ import annotations

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant

RNG = np.random.default_rng(7)
N = 25_000

# ---------------------------------------------------------------- 1. Data
regions = ["north", "south", "east", "west", "central"]
age            = RNG.integers(18, 75, N)
tenure_months  = np.clip((age - 18) * 12 * RNG.uniform(0.1, 0.9, N), 1, 600).astype(int)
salary         = RNG.lognormal(10.5, 0.4, N) * (1 + 0.005 * (age - 30))
product_count  = np.clip(RNG.poisson(1 + tenure_months / 60, N), 1, 8)
app_logins     = RNG.poisson(15 - 0.15 * (age - 30).clip(0), N).clip(0)
email_engage   = np.clip(RNG.beta(2, 5, N) + 0.002 * (50 - age), 0, 1)
deposit        = salary * RNG.uniform(0.05, 3.0, N) * (1 + tenure_months / 240)
region         = RNG.choice(regions, N)

# customer_value = f(behaviour, tenure, deposits) + noise
noise = RNG.normal(0, 40, N)
customer_value = (
    120
    + 0.0009 * salary
    + 3.2   * tenure_months / 12
    + 18    * product_count
    + 1.5   * app_logins
    + 65    * email_engage
    + 0.00007 * deposit
    + noise
)

df = pd.DataFrame({
    "age": age, "salary": salary.round(0), "region": region,
    "tenure_months": tenure_months, "product_count": product_count,
    "app_logins": app_logins, "email_engagement": email_engage.round(3),
    "deposit_balance": deposit.round(0), "customer_value": customer_value.round(2),
})

# ---------------------------------------------------------------- 2. Descriptives
print("descriptive stats:")
print(df.describe().round(2))

# ---------------------------------------------------------------- 3. Correlation
num_cols = ["age", "salary", "tenure_months", "product_count",
            "app_logins", "email_engagement", "deposit_balance", "customer_value"]
print("\npairwise correlations (pearson):")
print(df[num_cols].corr().round(3))

# ---------------------------------------------------------------- 4. Split
X = df.drop(columns=["customer_value"])
y = df["customer_value"]
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.20, random_state=42,
)

# ---------------------------------------------------------------- 5. Preproc
# StandardScaler on numeric, OneHotEncoder on region.
numeric = ["age", "salary", "tenure_months", "product_count",
           "app_logins", "email_engagement", "deposit_balance"]
preproc = ColumnTransformer([
    ("num", StandardScaler(), numeric),
    ("cat", OneHotEncoder(handle_unknown="ignore", drop="first"), ["region"]),
])

# ---------------------------------------------------------------- 6. Models
simple = Pipeline([
    ("prep", ColumnTransformer([("num", StandardScaler(), ["salary"])])),
    ("lr", LinearRegression()),
]).fit(X_train[["salary"]], y_train)

full = Pipeline([("prep", preproc), ("lr", LinearRegression())]).fit(X_train, y_train)

print(f"\nR² — salary only (test):     {r2_score(y_test, simple.predict(X_test[['salary']])):.3f}")
print(f"R² — all features (test):    {r2_score(y_test, full.predict(X_test)):.3f}")

# ---------------------------------------------------------------- 7. VIF
vif_frame = add_constant(df[numeric])
vif = pd.Series(
    [variance_inflation_factor(vif_frame.values, i) for i in range(vif_frame.shape[1])],
    index=vif_frame.columns,
)
print("\nvariance inflation factors:")
print(vif.round(2).to_string())