generate_data.py

"""
generate_data.py
----------------
Creates synthetic data for the churn-observability demo and writes it into a
local DuckDB database file.

It produces two raw tables:
  - raw_customers : one row per customer (the customer dimension + churn label)
  - raw_events    : many rows per customer (logins, feature use, support tickets)

dbt later turns these into a single `customer_features` table.

Day 1 generates HEALTHY data only. Failure injections (schema / volume /
freshness / drift) are added in a later step.

Run it directly:  python generate_data.py
Or via Make:      make generate
"""

import os
import datetime as dt

import numpy as np
import pandas as pd
import duckdb

# Same seed every run -> reproducible data. Important so the model and the
# saved drift "reference snapshot" stay stable between runs.
RNG = np.random.default_rng(42)

N_CUSTOMERS = 2000
PLANS = ["basic", "standard", "premium"]
PLAN_WEIGHTS = [0.5, 0.35, 0.15]
PLAN_PRICE = {"basic": 10.0, "standard": 25.0, "premium": 60.0}
REGIONS = ["us-east", "us-west", "eu", "apac"]


def make_customers() -> pd.DataFrame:
    """Build the customer dimension and a churn label derived from behaviour."""
    customer_id = np.arange(1, N_CUSTOMERS + 1)

    # Signed up somewhere in the last ~2 years.
    tenure_days = RNG.integers(15, 730, size=N_CUSTOMERS)
    today = dt.date.today()
    signup_date = [today - dt.timedelta(days=int(d)) for d in tenure_days]

    plan_type = RNG.choice(PLANS, size=N_CUSTOMERS, p=PLAN_WEIGHTS)
    monthly_charges = np.array([PLAN_PRICE[p] for p in plan_type]) + RNG.normal(0, 2, N_CUSTOMERS)
    monthly_charges = monthly_charges.round(2)
    region = RNG.choice(REGIONS, size=N_CUSTOMERS)

    # Per-customer behaviour "truth" used both to create events and to decide churn.
    n_logins = RNG.poisson(15, N_CUSTOMERS) + 1          # at least 1 login
    n_features = RNG.poisson(8, N_CUSTOMERS)
    n_tickets = RNG.poisson(1.0, N_CUSTOMERS)            # many customers have 0
    sentiment_mean = np.clip(RNG.normal(0.2, 0.4, N_CUSTOMERS), -1, 1)
    session_mean = np.clip(RNG.normal(20, 8, N_CUSTOMERS), 2, None)

    # Churn is more likely with: short tenure, many tickets, negative sentiment,
    # few logins. This gives the model a real signal to learn.
    logit = (
        -1.0
        - 0.0020 * tenure_days
        + 0.30 * n_tickets
        - 0.80 * sentiment_mean
        - 0.04 * n_logins
    )
    churn_prob = 1.0 / (1.0 + np.exp(-logit))
    is_churned = (RNG.random(N_CUSTOMERS) < churn_prob).astype(int)

    customers = pd.DataFrame(
        {
            "customer_id": customer_id,
            "signup_date": signup_date,
            "plan_type": plan_type,
            "monthly_charges": monthly_charges,
            "region": region,
            "is_churned": is_churned,
            # behaviour columns kept only to drive event generation, dropped below
            "gen_n_logins": n_logins,
            "gen_n_features": n_features,
            "gen_n_tickets": n_tickets,
            "gen_sentiment_mean": sentiment_mean,
            "gen_session_mean": session_mean,
        }
    )
    return customers


def make_events(customers: pd.DataFrame) -> pd.DataFrame:
    """Expand each customer into many activity events over the last 30 days."""
    rows = []
    event_id = 1
    today = dt.date.today()

    for c in customers.itertuples(index=False):
        # logins + feature use -> these carry session_minutes
        for _ in range(int(c.gen_n_logins) + int(c.gen_n_features)):
            etype = "login" if RNG.random() < 0.65 else "feature_use"
            rows.append(
                {
                    "event_id": event_id,
                    "customer_id": c.customer_id,
                    "event_date": today - dt.timedelta(days=int(RNG.integers(0, 30))),
                    "event_type": etype,
                    "session_minutes": round(float(np.clip(RNG.normal(c.gen_session_mean, 5), 1, None)), 1),
                    "sentiment": None,
                }
            )
            event_id += 1

        # support tickets -> these carry a sentiment score (the "model-derived" feature)
        for _ in range(int(c.gen_n_tickets)):
            rows.append(
                {
                    "event_id": event_id,
                    "customer_id": c.customer_id,
                    "event_date": today - dt.timedelta(days=int(RNG.integers(0, 30))),
                    "event_type": "support_ticket",
                    "session_minutes": None,
                    "sentiment": round(float(np.clip(RNG.normal(c.gen_sentiment_mean, 0.2), -1, 1)), 3),
                }
            )
            event_id += 1

    return pd.DataFrame(rows)


def write_to_duckdb(customers: pd.DataFrame, events: pd.DataFrame) -> str:
    path = os.environ.get("DUCKDB_PATH", "data/churn.duckdb")
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)

    # loaded_at = "when this data landed". dbt uses it for freshness checks and
    # compares it against UTC "now", so write naive-UTC (not local) time —
    # otherwise a non-UTC machine shows a phantom age and freshness can misfire.
    now = pd.Timestamp.now(tz="UTC").tz_localize(None)
    customers = customers.drop(columns=[c for c in customers.columns if c.startswith("gen_")]).copy()
    customers["loaded_at"] = now
    events = events.copy()
    events["loaded_at"] = now

    con = duckdb.connect(path)
    con.register("customers_df", customers)
    con.register("events_df", events)
    con.execute("CREATE OR REPLACE TABLE raw_customers AS SELECT * FROM customers_df")
    con.execute("CREATE OR REPLACE TABLE raw_events AS SELECT * FROM events_df")
    n_cust = con.execute("SELECT count(*) FROM raw_customers").fetchone()[0]
    n_evt = con.execute("SELECT count(*) FROM raw_events").fetchone()[0]
    churn_rate = con.execute("SELECT avg(is_churned) FROM raw_customers").fetchone()[0]
    con.close()

    print(f"Wrote {n_cust:,} customers and {n_evt:,} events to {path}")
    print(f"Churn rate: {churn_rate:.1%}")
    return path


def main():
    customers = make_customers()
    events = make_events(customers)
    write_to_duckdb(customers, events)


if __name__ == "__main__":
    main()