# -*- coding: utf-8 -*-
"""
Causal modelling of NPS — reproducible demonstration on SYNTHETIC data.

NOTHING HERE COMES FROM A CLIENT. Every row is generated by the simulator
below. The real engagement this illustrates is under NDA, so the data, the
company and its results are not reproduced anywhere. What is reproduced is
the *method*, on data whose ground truth is known because we wrote it.

That last point is the reason a synthetic demonstration is worth publishing:
with real data nobody can check whether a causal estimate is right. Here the
true effect is a constant in the code, so the estimator can be graded.

Run:  python nps-causal-demo.py
Out:  ../assets/case-nps-causal.png
"""

import numpy as np
import pandas as pd
import statsmodels.api as sm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

SEMENTE = 20260828
EFEITO_VERDADEIRO = 0.80      # points of NPS caused by the intervention
N = 6000

# Paleta do site (theme.scss)
PAPEL, TINTA, ACENTO, REGUA = "#f7f5f0", "#1c1b18", "#7d2e2e", "#d9d4c8"
CINZA = "#8a857a"


def gerar(n=N, semente=SEMENTE):
    """Customers, a success-team intervention, and an NPS score.

    The confounding is deliberate: bigger, longer-tenured accounts are both
    more likely to receive the proactive call AND happier to begin with.
    That is what makes the naive comparison wrong.
    """
    rng = np.random.default_rng(semente)

    porte = rng.normal(0, 1, n)                       # account size (standardised)
    tempo_casa = rng.gamma(2.0, 6.0, n)               # tenure, months
    tickets = rng.poisson(2.0 + 0.8 * np.clip(porte, 0, None), n)

    # Quem recebe a ligacao proativa nao e sorteado: depende de porte e tempo.
    log_odds = -0.7 + 0.9 * porte + 0.02 * tempo_casa - 0.15 * tickets
    p = 1 / (1 + np.exp(-log_odds))
    intervencao = rng.binomial(1, p)

    # NPS 0-10. O efeito causal da intervencao e EFEITO_VERDADEIRO, e nada mais.
    nps = (
        6.4
        + 0.85 * porte                                # confundidor
        + 0.020 * tempo_casa                          # confundidor
        - 0.30 * tickets                              # confundidor
        + EFEITO_VERDADEIRO * intervencao             # <-- o que queremos recuperar
        + rng.normal(0, 1.4, n)
    )
    nps = np.clip(nps, 0, 10)

    return pd.DataFrame(
        dict(porte=porte, tempo_casa=tempo_casa, tickets=tickets,
             intervencao=intervencao, nps=nps, p_verdadeiro=p)
    )


def ingenuo(df):
    """Difference in means: what a dashboard reports."""
    a = df.loc[df.intervencao == 1, "nps"]
    b = df.loc[df.intervencao == 0, "nps"]
    dif = a.mean() - b.mean()
    ep = np.sqrt(a.var(ddof=1) / len(a) + b.var(ddof=1) / len(b))
    return dif, 1.96 * ep


def ajustado(df):
    """Backdoor adjustment: regress on treatment plus the confounders."""
    X = sm.add_constant(df[["intervencao", "porte", "tempo_casa", "tickets"]])
    m = sm.OLS(df["nps"], X).fit()
    return m.params["intervencao"], 1.96 * m.bse["intervencao"]


def ipw(df):
    """Inverse-probability weighting on an estimated propensity score."""
    X = sm.add_constant(df[["porte", "tempo_casa", "tickets"]])
    ps = sm.Logit(df["intervencao"], X).fit(disp=0).predict(X)
    ps = np.clip(ps, 0.02, 0.98)
    w = np.where(df.intervencao == 1, 1 / ps, 1 / (1 - ps))
    Xw = sm.add_constant(df[["intervencao"]])
    m = sm.WLS(df["nps"], Xw, weights=w).fit()
    return m.params["intervencao"], 1.96 * m.bse["intervencao"]


def grafico(estimativas, caminho):
    fig, ax = plt.subplots(figsize=(8.6, 4.2), dpi=170)
    fig.patch.set_facecolor(PAPEL)
    ax.set_facecolor(PAPEL)

    rotulos = [e[0] for e in estimativas]
    y = np.arange(len(rotulos))[::-1]

    for (rot, val, erro), yy in zip(estimativas, y):
        cor = CINZA if "Correlation" in rot else ACENTO
        ax.errorbar(val, yy, xerr=erro, fmt="o", ms=7, lw=2, capsize=5,
                    color=cor, ecolor=cor, zorder=3)
        ax.text(val, yy + 0.24, f"{val:.2f}", ha="center", va="bottom",
                fontsize=10, color=cor, fontweight="600")

    ax.axvline(EFEITO_VERDADEIRO, color=TINTA, ls="--", lw=1.2, zorder=2)
    ax.text(EFEITO_VERDADEIRO, len(rotulos) - 0.45,
            f"  true effect = {EFEITO_VERDADEIRO:.2f}",
            color=TINTA, fontsize=10, va="center")

    ax.set_yticks(y)
    ax.set_yticklabels(rotulos, fontsize=11, color=TINTA)
    ax.set_xlabel("Estimated effect on NPS (points)", fontsize=10, color=TINTA)
    ax.tick_params(colors=TINTA, labelsize=9)
    for lado in ("top", "right", "left"):
        ax.spines[lado].set_visible(False)
    ax.spines["bottom"].set_color(REGUA)
    ax.grid(axis="x", color=REGUA, lw=0.7, alpha=0.7, zorder=0)
    ax.set_ylim(-0.7, len(rotulos) - 0.25)
    ax.set_title("Synthetic data — the true effect is known by construction",
                 fontsize=11, color=CINZA, loc="left", pad=12)

    fig.tight_layout()
    fig.savefig(caminho, facecolor=PAPEL)
    print("grafico:", caminho)


if __name__ == "__main__":
    df = gerar()

    d_ing, e_ing = ingenuo(df)
    d_aj, e_aj = ajustado(df)
    d_ipw, e_ipw = ipw(df)

    print(f"n = {len(df)}, tratados = {int(df.intervencao.sum())}")
    print(f"efeito verdadeiro          {EFEITO_VERDADEIRO:.3f}")
    print(f"correlacao (dif. de medias) {d_ing:.3f} +/- {e_ing:.3f}   vies = {d_ing - EFEITO_VERDADEIRO:+.3f}")
    print(f"ajuste por regressao        {d_aj:.3f} +/- {e_aj:.3f}   vies = {d_aj - EFEITO_VERDADEIRO:+.3f}")
    print(f"ponderacao (IPW)            {d_ipw:.3f} +/- {e_ipw:.3f}   vies = {d_ipw - EFEITO_VERDADEIRO:+.3f}")

    grafico(
        [("Correlation (raw difference)", d_ing, e_ing),
         ("Backdoor adjustment", d_aj, e_aj),
         ("Propensity weighting (IPW)", d_ipw, e_ipw)],
        "../assets/case-nps-causal.png",
    )