# %% Generate a population with a known treatment effect
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.formula.api as smf
import rdrobust as rd
import rddensity as density

# A seed keeps the sample unchanged when we run the code again.
rng = np.random.default_rng(2026)
n = 20000
original_score = rng.uniform(-2, 2, n)
ability = rng.normal(0, 1, n)
outcome_noise = rng.normal(0, 1, n)

# The score determines eligibility. The true treatment effect is 2.
treatment_before = (original_score >= 0).astype(int)
outcome_without_treatment = 10 + 0.5 * original_score + 2 * ability + outcome_noise
outcome_before = outcome_without_treatment + 2 * treatment_before

# %% Let some people change their recorded score
# People just below zero with above-average ability move across the cutoff.
# Reflecting their score puts them just above zero without creating a mass at zero.
moves = (original_score >= -0.4) & (original_score < 0) & (ability > 0)
recorded_score = original_score.copy()
recorded_score[moves] = -original_score[moves]

# Ability and the untreated outcome stay unchanged. Treatment follows the new score.
treatment_after = (recorded_score >= 0).astype(int)
outcome_after = outcome_without_treatment + 2 * treatment_after

sample = pd.DataFrame({
    "original_score": original_score,
    "recorded_score": recorded_score,
    "ability": ability,
    "moves": moves,
    "treatment_before": treatment_before,
    "treatment_after": treatment_after,
    "outcome_before": outcome_before,
    "outcome_after": outcome_after,
    "outcome_without_treatment": outcome_without_treatment
})
print("People who change their score:", int(sample["moves"].sum()))

# %% Plot the score distribution before manipulation
# Zero is a bin boundary, so no bin crosses the treatment cutoff.
bin_edges = np.linspace(-2, 2, 41)
fig, ax = plt.subplots(figsize=(10, 5))
ax.hist(sample["original_score"], bins=bin_edges, density=True,
        color="#0072B2", edgecolor="white")
ax.axvline(0, color="black", linestyle="--", linewidth=2)
ax.set(xlabel="Original score", ylabel="Density", xlim=(-2, 2), ylim=(0, 0.46))
fig.tight_layout()
fig.savefig("resources/images/week-07/density-before.png", dpi=200)
plt.close(fig)

# %% Plot the score distribution after manipulation, on the same axes
fig, ax = plt.subplots(figsize=(10, 5))
ax.hist(sample["recorded_score"], bins=bin_edges, density=True,
        color="#D55E00", edgecolor="white")
ax.axvline(0, color="black", linestyle="--", linewidth=2)
ax.set(xlabel="Recorded score", ylabel="Density", xlim=(-2, 2), ylim=(0, 0.46))
fig.tight_layout()
fig.savefig("resources/images/week-07/density-after.png", dpi=200)
plt.close(fig)

# %% Calculate bin means for ability and outcomes
sample["original_bin"] = pd.cut(sample["original_score"], bin_edges, right=False)
sample["recorded_bin"] = pd.cut(sample["recorded_score"], bin_edges, right=False)
before_means = sample.groupby("original_bin", observed=True)[
    ["original_score", "ability", "outcome_before", "outcome_without_treatment"]
].mean()
after_means = sample.groupby("recorded_bin", observed=True)[
    ["recorded_score", "ability", "outcome_after", "outcome_without_treatment"]
].mean()

# %% Show how the composition changes
fig, ax = plt.subplots(figsize=(10, 5))
ax.scatter(before_means["original_score"], before_means["ability"],
           color="#0072B2", marker="o", label="Before manipulation")
ax.scatter(after_means["recorded_score"], after_means["ability"],
           color="#D55E00", marker="^", label="After manipulation")
ax.axvline(0, color="black", linestyle="--", linewidth=2)
ax.axhline(0, color="black", linewidth=0.7)
ax.set(xlabel="Score", ylabel="Mean ability within each bin", xlim=(-1, 1))
ax.legend()
fig.tight_layout()
fig.savefig("resources/images/week-07/ability-sorting.png", dpi=200)
plt.close(fig)

# %% Show the consequence for observed outcomes
fig, ax = plt.subplots(figsize=(10, 5))
ax.scatter(before_means["original_score"], before_means["outcome_before"],
           color="#0072B2", marker="o", label="Before manipulation")
ax.scatter(after_means["recorded_score"], after_means["outcome_after"],
           color="#D55E00", marker="^", label="After manipulation")
ax.axvline(0, color="black", linestyle="--", linewidth=2)
ax.set(xlabel="Score", ylabel="Mean observed outcome within each bin", xlim=(-1, 1))
ax.legend()
fig.tight_layout()
fig.savefig("resources/images/week-07/outcome-sorting.png", dpi=200)
plt.close(fig)

# %% Estimate a sharp RD with familiar regression code
# Centre the score at the cutoff. Here the cutoff is already zero.
# score_right allows the slope to differ above the cutoff.
sample["score_right_before"] = sample["original_score"] * sample["treatment_before"]
sample["score_right_after"] = sample["recorded_score"] * sample["treatment_after"]
local_before = sample.loc[sample["original_score"].abs() < 0.4].copy()
local_after = sample.loc[sample["recorded_score"].abs() < 0.4].copy()

ols_before = smf.ols(
    "outcome_before ~ treatment_before + original_score + score_right_before",
    data=local_before
).fit(cov_type="HC1")
print(ols_before.summary())

ols_after = smf.ols(
    "outcome_after ~ treatment_after + recorded_score + score_right_after",
    data=local_after
).fit(cov_type="HC1")
print(ols_after.summary())

# %% Give observations closer to zero more weight
local_before["weight"] = 1 - local_before["original_score"].abs() / 0.4
weighted_before = smf.wls(
    "outcome_before ~ treatment_before + original_score + score_right_before",
    data=local_before, weights=local_before["weight"]
).fit(cov_type="HC1")
print(weighted_before.summary())

# %% Use rdrobust for bandwidth selection and bias-corrected inference
rd_before = rd.rdrobust(y=sample["outcome_before"], x=sample["original_score"],
                       c=0, p=1, kernel="triangular", bwselect="mserd")
print(rd_before)
rd_after = rd.rdrobust(y=sample["outcome_after"], x=sample["recorded_score"],
                      c=0, p=1, kernel="triangular", bwselect="mserd")
print(rd_after)

# %% Test for a density discontinuity
# This is the Cattaneo-Jansson-Ma implementation of density-discontinuity testing.
density_before = density.rddensity(X=sample["original_score"], c=0)
print(density_before)
density_after = density.rddensity(X=sample["recorded_score"], c=0)
print(density_after)

# The default printout in some versions only shows the estimation settings.
# Extract the default jackknife test statistic and p-value explicitly.
print("Before: density-test statistic =", density_before.test["t_jk"])
print("Before: density-test p-value =", density_before.test["p_jk"])
print("After: density-test statistic =", density_after.test["t_jk"])
print("After: density-test p-value =", density_after.test["p_jk"])