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

# A model learns the past (step 6 of the advanced lab: same seeds, same numbers). 4,000 past loan
# decisions: the score (ability to repay, 0 to 100) has the same distribution in groups A and B,
# but the officers approved A from 50 up and B only from 65 up. A depth-3 tree trained to imitate
# them is scored on 2,000 new applicants. Left: with the group as a column, the share approved at
# each score (bands of 5 points). Right: the same with the group column removed and only the zone
# where people live, which 80% of B and 20% of A share, left in its place.
rng = np.random.default_rng(0)
def applicants(n):
    group_B = rng.integers(0, 2, n)
    return pd.DataFrame({
        'score': rng.integers(0, 101, n),
        'group_B': group_B,
        'zone': (rng.random(n) < np.where(group_B == 1, 0.8, 0.2)).astype(int),
    })

past = applicants(4000)
past['approved'] = (past['score'] >= np.where(past['group_B'] == 1, 65, 50)).astype(int)
new = applicants(2000)
for name, cols in [('with', ['score', 'group_B']), ('proxy', ['score', 'zone'])]:
    m = DecisionTreeClassifier(max_depth=3, random_state=0).fit(past[cols], past['approved'])
    new[name] = m.predict(new[cols])

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.9), sharey=True)
band = new['score'] // 5 * 5 + 2.5
titles = {'with': 'the group is a column', 'proxy': 'group column removed, zone kept'}
for ax, name in zip(axes, ['with', 'proxy']):
    ax.axvspan(50, 65, color='gray', alpha=0.12, lw=0)
    for g, col, lab in [(0, BLUE, 'group A'), (1, ORANGE, 'group B')]:
        r = new[new['group_B'] == g].groupby(band[new['group_B'] == g])[name].mean()
        ax.plot(r.index, r.values, 'o-', color=col, lw=2, ms=4, label=lab)
    s = new[new['score'].between(50, 64)].groupby('group_B')[name].mean()
    ax.text(24, 0.55, f'gray band, scores 50 to 64:\napproved A {s[0]:.2f} · B {s[1]:.2f}', ha='center', va='center', fontsize=8.5)
    ax.set_title(titles[name], fontsize=10)
    ax.set_xlabel('repayment score (bands of 5 points)')
    ax.set_xlim(0, 100)
    ax.spines[['top', 'right']].set_visible(False)
axes[0].set_ylabel('share of new applicants approved')
axes[0].legend(fontsize=8, frameon=False, loc='upper left')
plt.tight_layout()

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

plt.close()