workload.py

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd


@dataclass(slots=True)
class WorkloadSummary:
    rows: int
    customers: int
    working_set_mb: float
    total_revenue: float
    total_margin: float
    repeat_customer_rate: float
    top_customers: list[dict[str, float | int | str]]
    ltv_bands: list[dict[str, float | int | str]]

    def as_dict(self) -> dict[str, object]:
        return {
            "rows": self.rows,
            "customers": self.customers,
            "working_set_mb": self.working_set_mb,
            "total_revenue": self.total_revenue,
            "total_margin": self.total_margin,
            "repeat_customer_rate": self.repeat_customer_rate,
            "top_customers": self.top_customers,
            "ltv_bands": self.ltv_bands,
        }


def build_customer_ltv_report(
    *,
    rows: int = 250_000,
    customer_count: int = 25_000,
    top_n: int = 10,
    seed: int = 7,
) -> WorkloadSummary:
    """Simulate a customer analytics job that leans heavily on NumPy and pandas."""
    rng = np.random.default_rng(seed)

    customer_ids = rng.integers(1, customer_count + 1, size=rows, dtype=np.int32)
    order_values = rng.lognormal(mean=4.2, sigma=0.55, size=rows).astype(np.float64)
    quantities = rng.integers(1, 6, size=rows, dtype=np.int16)
    discount_rate = rng.uniform(0.0, 0.35, size=rows).astype(np.float32)
    acquisition_cost = rng.gamma(shape=2.5, scale=7.5, size=rows).astype(np.float32)
    support_cost = rng.gamma(shape=1.4, scale=2.8, size=rows).astype(np.float32)
    refund_flag = rng.random(size=rows) < 0.05
    fulfilled_fast = rng.random(size=rows) < 0.62
    channels = np.take(np.array(["ads", "email", "seo", "partner"]), rng.integers(0, 4, size=rows))
    countries = np.take(np.array(["US", "CA", "GB", "DE", "AU", "JP"]), rng.integers(0, 6, size=rows))

    gross_revenue = order_values * quantities
    discount_amount = gross_revenue * discount_rate
    refunded_revenue = np.where(refund_flag, gross_revenue * 0.85, 0.0)
    shipping_cost = np.where(fulfilled_fast, 9.5, 4.0).astype(np.float32)
    net_revenue = gross_revenue - discount_amount - refunded_revenue
    contribution_margin = net_revenue - acquisition_cost - support_cost - shipping_cost

    frame = pd.DataFrame(
        {
            "customer_id": customer_ids,
            "channel": pd.Categorical(channels),
            "country": pd.Categorical(countries),
            "gross_revenue": gross_revenue,
            "discount_amount": discount_amount,
            "refunded_revenue": refunded_revenue,
            "net_revenue": net_revenue,
            "contribution_margin": contribution_margin,
            "fulfilled_fast": fulfilled_fast,
        }
    )

    customer_rollup = (
        frame.groupby("customer_id", sort=False, observed=True)
        .agg(
            orders=("gross_revenue", "size"),
            gross_revenue=("gross_revenue", "sum"),
            net_revenue=("net_revenue", "sum"),
            contribution_margin=("contribution_margin", "sum"),
            fast_fulfillment_rate=("fulfilled_fast", "mean"),
        )
        .reset_index()
    )

    customer_rollup["ltv_band"] = pd.qcut(
        customer_rollup["net_revenue"].rank(method="first"),
        q=5,
        labels=["emerging", "growing", "core", "vip", "elite"],
    )
    customer_rollup["ltv_band"] = customer_rollup["ltv_band"].astype("string")

    band_rollup = (
        customer_rollup.groupby("ltv_band", observed=False)
        .agg(
            customers=("customer_id", "size"),
            avg_net_revenue=("net_revenue", "mean"),
            avg_margin=("contribution_margin", "mean"),
            avg_orders=("orders", "mean"),
        )
        .reset_index()
    )

    top_customers_frame = customer_rollup.nlargest(top_n, "net_revenue")[
        ["customer_id", "net_revenue", "contribution_margin", "orders", "ltv_band"]
    ].round({"net_revenue": 2, "contribution_margin": 2})
    ltv_bands_frame = band_rollup.round({"avg_net_revenue": 2, "avg_margin": 2, "avg_orders": 2})

    top_customers = [
        {
            "customer_id": int(record["customer_id"]),
            "net_revenue": float(record["net_revenue"]),
            "contribution_margin": float(record["contribution_margin"]),
            "orders": int(record["orders"]),
            "ltv_band": str(record["ltv_band"]),
        }
        for record in top_customers_frame.to_dict(orient="records")
    ]
    ltv_bands = [
        {
            "ltv_band": str(record["ltv_band"]),
            "customers": int(record["customers"]),
            "avg_net_revenue": float(record["avg_net_revenue"]),
            "avg_margin": float(record["avg_margin"]),
            "avg_orders": float(record["avg_orders"]),
        }
        for record in ltv_bands_frame.to_dict(orient="records")
    ]

    working_set_bytes = int(frame.memory_usage(deep=True).sum() + customer_rollup.memory_usage(deep=True).sum())
    repeat_customer_rate = float((customer_rollup["orders"] > 1).mean())

    return WorkloadSummary(
        rows=rows,
        customers=int(customer_rollup.shape[0]),
        working_set_mb=round(working_set_bytes / (1024 * 1024), 2),
        total_revenue=round(float(frame["net_revenue"].sum()), 2),
        total_margin=round(float(frame["contribution_margin"].sum()), 2),
        repeat_customer_rate=round(repeat_customer_rate, 4),
        top_customers=top_customers,
        ltv_bands=ltv_bands,
    )