import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline

np.random.seed(42)

# -----------------------------
# Simulation settings
# -----------------------------
N_CUSTOMERS = 50_000
N_MONTHS = 12
DISCOUNT_RATE_MONTHLY = 0.01

# -----------------------------
# Customer-level heterogeneity
# -----------------------------
customers = pd.DataFrame({
    "customer_id": np.arange(N_CUSTOMERS),
})

# Cohort month is the first active month. This lets us compare customers by age since acquisition.
customers["cohort_month"] = np.random.choice(np.arange(1, 7), size=N_CUSTOMERS, p=[0.18, 0.17, 0.17, 0.16, 0.16, 0.16])

# Latent customer quality drives spend, conversion, retention, and targeting probability.
customers["customer_quality"] = np.random.normal(0, 1, N_CUSTOMERS)

# Public-dataset-like behavioral features. These are not real public data; they mimic common transaction features.
customers["historical_frequency"] = np.random.poisson(lam=np.exp(1.1 + 0.35 * customers["customer_quality"]))
customers["historical_monetary"] = np.random.lognormal(mean=3.2 + 0.35 * customers["customer_quality"], sigma=0.55)
customers["recency"] = np.clip(np.random.normal(45 - 8 * customers["customer_quality"], 18, N_CUSTOMERS), 1, 120)

# Simple customer segments based on historical monetary value.
customers["customer_segment"] = pd.qcut(
    customers["historical_monetary"],
    q=3,
    labels=["low_value", "mid_value", "high_value"]
)

# Baseline conversion propensity: what would happen without treatment.
logit_base = (
    -3.4
    + 0.35 * customers["customer_quality"]
    + 0.025 * customers["historical_frequency"]
    - 0.006 * customers["recency"]
)
customers["baseline_propensity"] = 1 / (1 + np.exp(-logit_base))

# Baseline monthly value conditional on conversion.
customers["baseline_order_value"] = 20 + 0.55 * customers["historical_monetary"] + np.random.normal(0, 8, N_CUSTOMERS)
customers["baseline_order_value"] = customers["baseline_order_value"].clip(lower=5)

customers.head()

# -----------------------------
# Expand to customer-month panel
# -----------------------------
panel = customers.loc[customers.index.repeat(N_MONTHS)].copy()
panel["calendar_month"] = np.tile(np.arange(1, N_MONTHS + 1), N_CUSTOMERS)
panel["months_since_cohort"] = panel["calendar_month"] - panel["cohort_month"] + 1

# Keep only months where the customer has already entered the population.
panel = panel[panel["months_since_cohort"] >= 1].copy()

# Treatment assignment is intentionally biased.
# Higher-quality and higher-value customers are more likely to receive the offer.
treat_logit = (
    -1.5
    + 0.75 * panel["customer_quality"]
    + 0.012 * panel["historical_frequency"]
    + 0.007 * panel["historical_monetary"]
    - 0.004 * panel["recency"]
)
panel["treatment_probability"] = 1 / (1 + np.exp(-treat_logit))
panel["treatment"] = np.random.binomial(1, panel["treatment_probability"])

# True treatment effects are known because this is simulated.
# Treatment modestly increases conversion and slightly improves retention.
panel["true_conversion_lift"] = 0.025
panel["true_retention_lift"] = 0.010

# Base churn probability decreases for better customers.
base_churn = (
    0.18
    - 0.025 * panel["customer_quality"]
    - 0.0007 * panel["historical_monetary"]
    + 0.005 * panel["months_since_cohort"]
)

panel["monthly_churn_probability"] = (
    base_churn
    - panel["true_retention_lift"] * panel["treatment"]
).clip(0.02, 0.45)

# Simulate active status. Once a customer churns, they stay inactive.
panel = panel.sort_values(["customer_id", "calendar_month"]).reset_index(drop=True)
panel["active_flag"] = 0
panel["churn_flag"] = 0

for customer_id, idx in panel.groupby("customer_id").groups.items():
    active = True
    for i in idx:
        if active:
            panel.at[i, "active_flag"] = 1
            churn = np.random.binomial(1, panel.at[i, "monthly_churn_probability"])
            panel.at[i, "churn_flag"] = churn
            if churn == 1:
                active = False
        else:
            panel.at[i, "active_flag"] = 0
            panel.at[i, "churn_flag"] = 0

# Conversion probability is active-dependent and treatment-dependent.
panel["conversion_probability"] = (
    panel["baseline_propensity"]
    + panel["true_conversion_lift"] * panel["treatment"]
).clip(0, 0.80)

panel["conversion"] = np.random.binomial(1, panel["conversion_probability"] * panel["active_flag"])

# Monthly value is realized only when active and converted.
noise = np.random.normal(0, 10, len(panel))
panel["monthly_value"] = (
    panel["conversion"]
    * (panel["baseline_order_value"] + 5 * panel["treatment"] + noise)
).clip(lower=0)

panel["discount_factor"] = 1 / ((1 + DISCOUNT_RATE_MONTHLY) ** (panel["months_since_cohort"] - 1))
panel["discounted_monthly_value"] = panel["monthly_value"] * panel["discount_factor"]

panel.head()

cohort_ltv = (
    panel.groupby(["cohort_month", "months_since_cohort"], as_index=False)
    .agg(
        customers=("customer_id", "nunique"),
        monthly_value=("monthly_value", "mean"),
        discounted_monthly_value=("discounted_monthly_value", "mean")
    )
    .sort_values(["cohort_month", "months_since_cohort"])
)

cohort_ltv["cumulative_ltv"] = cohort_ltv.groupby("cohort_month")["monthly_value"].cumsum()
cohort_ltv["cumulative_discounted_ltv"] = cohort_ltv.groupby("cohort_month")["discounted_monthly_value"].cumsum()

cohort_ltv.head(12)

plt.figure(figsize=(9, 5))
for cohort, d in cohort_ltv.groupby("cohort_month"):
    plt.plot(d["months_since_cohort"], d["cumulative_ltv"], marker="o", label=f"Cohort {cohort}")

plt.title("Cumulative LTV by Cohort")
plt.xlabel("Months Since Cohort Start")
plt.ylabel("Average Cumulative Value per Customer")
plt.legend(title="Cohort Month")
plt.grid(True, alpha=0.3)
plt.show()

retention = (
    panel.groupby(["cohort_month", "months_since_cohort"], as_index=False)
    .agg(retention_rate=("active_flag", "mean"))
)

retention.head(12)

plt.figure(figsize=(9, 5))
for cohort, d in retention.groupby("cohort_month"):
    plt.plot(d["months_since_cohort"], d["retention_rate"], marker="o", label=f"Cohort {cohort}")

plt.title("Retention Curves by Cohort")
plt.xlabel("Months Since Cohort Start")
plt.ylabel("Share of Customers Active")
plt.ylim(0, 1.05)
plt.legend(title="Cohort Month")
plt.grid(True, alpha=0.3)
plt.show()

# Estimate empirical survival by customer segment and customer age.
segment_survival = (
    panel.groupby(["customer_segment", "months_since_cohort"], as_index=False, observed=True)
    .agg(empirical_survival=("active_flag", "mean"))
)

segment_value = (
    panel[panel["active_flag"] == 1]
    .groupby(["customer_segment", "months_since_cohort"], as_index=False, observed=True)
    .agg(avg_value_if_active=("monthly_value", "mean"))
)

survival_ltv = segment_survival.merge(
    segment_value,
    on=["customer_segment", "months_since_cohort"],
    how="left"
)

survival_ltv["avg_value_if_active"] = survival_ltv["avg_value_if_active"].fillna(0)
survival_ltv["discount_factor"] = 1 / ((1 + DISCOUNT_RATE_MONTHLY) ** (survival_ltv["months_since_cohort"] - 1))
survival_ltv["survival_adjusted_value"] = (
    survival_ltv["empirical_survival"]
    * survival_ltv["avg_value_if_active"]
    * survival_ltv["discount_factor"]
)

survival_ltv["survival_adjusted_ltv"] = (
    survival_ltv.groupby("customer_segment", observed=True)["survival_adjusted_value"].cumsum()
)

survival_ltv.head(12)

plt.figure(figsize=(9, 5))
for segment, d in survival_ltv.groupby("customer_segment", observed=True):
    plt.plot(d["months_since_cohort"], d["survival_adjusted_ltv"], marker="o", label=str(segment))

plt.title("Survival-Adjusted LTV by Customer Segment")
plt.xlabel("Months Since Cohort Start")
plt.ylabel("Average Survival-Adjusted LTV")
plt.legend(title="Segment")
plt.grid(True, alpha=0.3)
plt.show()

# Customer-level LTV over the observed horizon.
customer_ltv = (
    panel.groupby("customer_id", as_index=False)
    .agg(
        observed_ltv=("discounted_monthly_value", "sum"),
        ever_treated=("treatment", "max"),
        avg_treatment_probability=("treatment_probability", "mean"),
        historical_frequency=("historical_frequency", "first"),
        historical_monetary=("historical_monetary", "first"),
        recency=("recency", "first"),
        customer_quality=("customer_quality", "first"),
        baseline_propensity=("baseline_propensity", "first"),
        customer_segment=("customer_segment", "first")
    )
)

naive_treated_ltv = customer_ltv.loc[customer_ltv["ever_treated"] == 1, "observed_ltv"].mean()
naive_untreated_ltv = customer_ltv.loc[customer_ltv["ever_treated"] == 0, "observed_ltv"].mean()
naive_difference = naive_treated_ltv - naive_untreated_ltv

naive_treated_ltv, naive_untreated_ltv, naive_difference

# Regression adjustment using pre-treatment features.
features = ["ever_treated", "historical_frequency", "historical_monetary", "recency", "baseline_propensity", "customer_segment"]
X = customer_ltv[features]
y = customer_ltv["observed_ltv"]

preprocess = ColumnTransformer(
    transformers=[
        ("segment", OneHotEncoder(drop="first"), ["customer_segment"]),
    ],
    remainder="passthrough"
)

regression_adjustment = Pipeline(steps=[
    ("preprocess", preprocess),
    ("model", LinearRegression())
])

regression_adjustment.fit(X, y)

# Estimate adjusted treatment effect by predicting each customer twice:
# once as treated and once as untreated, while holding other features fixed.
X_treated = X.copy()
X_untreated = X.copy()
X_treated["ever_treated"] = 1
X_untreated["ever_treated"] = 0

pred_treated = regression_adjustment.predict(X_treated)
pred_untreated = regression_adjustment.predict(X_untreated)

adjusted_difference = np.mean(pred_treated - pred_untreated)
adjusted_difference

# Approximate the true simulated incremental LTV by replaying the known treatment effects.
# This uses the known conversion and retention lift used in the simulation.
# It is not something available in real observational data; it is included here as a benchmark.

true_incremental_by_month = panel.copy()

# Expected incremental conversion value among active treated observations.
true_incremental_by_month["incremental_conversion_value"] = (
    true_incremental_by_month["active_flag"]
    * true_incremental_by_month["treatment"]
    * true_incremental_by_month["true_conversion_lift"]
    * true_incremental_by_month["baseline_order_value"]
    * true_incremental_by_month["discount_factor"]
)

# Approximate incremental retention value:
# if treatment lowers churn, it preserves some future customer value.
# This simple approximation values the retained month using expected baseline value.
true_incremental_by_month["incremental_retention_value"] = (
    true_incremental_by_month["treatment"]
    * true_incremental_by_month["true_retention_lift"]
    * true_incremental_by_month["baseline_propensity"]
    * true_incremental_by_month["baseline_order_value"]
    * true_incremental_by_month["discount_factor"]
)

true_incremental_total = (
    true_incremental_by_month["incremental_conversion_value"].sum()
    + true_incremental_by_month["incremental_retention_value"].sum()
)

true_incremental_per_treated_customer = true_incremental_total / customer_ltv["ever_treated"].sum()

comparison = pd.DataFrame({
    "estimate": ["Naive observed difference", "Model-adjusted difference", "True simulated incremental LTV"],
    "ltv_difference": [naive_difference, adjusted_difference, true_incremental_per_treated_customer]
})

comparison

plt.figure(figsize=(8, 5))
plt.bar(comparison["estimate"], comparison["ltv_difference"])
plt.title("Observed vs Adjusted vs True Incremental LTV")
plt.ylabel("Average LTV Difference")
plt.xticks(rotation=20, ha="right")
plt.grid(True, axis="y", alpha=0.3)
plt.show()