import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO

# The simulation of the lesson, step by step: 10,000 people whose income grows
# with age; about 30% of the incomes are then hidden by three different
# mechanisms. Dots: a random 1,500 of the people (blue = income recorded,
# orange = income missing). Lines: the true mean and the mean of what we see.
rng = np.random.default_rng(0)
n = 10_000
age = rng.uniform(20, 65, n)
income = np.exp(0.5 + 0.03 * age + rng.normal(0, 0.4, n))     # R$ thousand per month
people = pd.DataFrame({'age': age, 'income': income})
young = people['age'] < 35
rich = people['income'] > people['income'].quantile(2 / 3)
p_missing = {'MCAR': np.full(n, 0.3),
             'MAR': np.where(young, 0.6, 0.15),
             'MNAR': np.where(rich, 0.6, 0.15)}
stories = {'MCAR': 'a coin decides, for everyone',
           'MAR': 'the young skip it more (age is seen)',
           'MNAR': 'the rich skip it more (income is not)'}

BLUE, ORANGE, GRAY = '#58a6ff', '#f0883e', 'gray'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 3, figsize=(10, 3.6), sharey=True)
show = np.random.default_rng(1).choice(n, 1500, replace=False)
true_mean = people['income'].mean()

for ax, (mech, p) in zip(axes, p_missing.items()):
    missing = rng.random(n) < p
    seen = people['income'].mask(missing)
    m = missing[show]
    ax.scatter(age[show][~m], income[show][~m], s=5, color=BLUE, alpha=0.5, lw=0, rasterized=True,
               label='recorded')
    ax.scatter(age[show][m], income[show][m], s=5, color=ORANGE, alpha=0.7, lw=0, rasterized=True,
               label='missing')
    ax.axhline(true_mean, color=GRAY, ls='--', lw=1.2)
    ax.axhline(seen.mean(), color=BLUE, lw=1.8)
    ax.text(21, 21.2, f'true mean {true_mean:.2f}\nmean of what we see {seen.mean():.2f}',
            fontsize=8, va='top')
    ax.set_title(f'{mech}: {stories[mech]}', fontsize=9)
    ax.set_xlabel('age', fontsize=8.5)
    ax.tick_params(labelsize=8)
    ax.set_ylim(0, 22)
axes[0].set_ylabel('income (R$ thousand / month)', fontsize=8.5)
axes[0].legend(fontsize=7.5, frameon=False, loc='center left', markerscale=2.5, bbox_to_anchor=(0, 0.62))
fig.suptitle('About 30% missing in every panel; only the mechanism changes, and with it the mean', fontsize=10)

plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()