from io import StringIO

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.base import clone
from sklearn.compose import make_column_transformer
from sklearn.datasets import fetch_openml
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# The lab's monitoring, steps 7, 8 and 10, drawn per quarter
raw = fetch_openml(data_id=1461, as_frame=True).frame
names = ['age', 'job', 'marital', 'education', 'default', 'balance', 'housing', 'loan',
         'contact', 'day', 'month', 'duration', 'campaign', 'pdays', 'previous', 'poutcome']
calls = raw.drop(columns='Class').set_axis(names, axis=1)
calls['y'] = (raw['Class'] == '2').astype(int)
month = calls['month'].astype(str)
year = 2008 + (month.ne(month.shift()) & month.eq('jan')).cumsum()
calls['date'] = pd.to_datetime(year.astype(str) + '-' + month, format='%Y-%b').dt.to_period('M')
CAT = ['job', 'marital', 'education', 'default', 'housing', 'loan', 'contact', 'poutcome']
NUM = ['age', 'balance', 'campaign', 'pdays', 'previous']
FEATURES = CAT + NUM
for c in CAT:
    calls[c] = calls[c].astype(str)
model = make_pipeline(make_column_transformer((OneHotEncoder(handle_unknown='ignore'), CAT),
                                              (StandardScaler(), NUM)),
                      LogisticRegression(max_iter=2000))
history = calls[calls['date'] < '2009-05']
fit_part, val_part = history[history['date'] < '2009-02'], history[history['date'] >= '2009-02']
val_auc = roc_auc_score(val_part['y'], clone(model).fit(fit_part[FEATURES], fit_part['y'])
                        .predict_proba(val_part[FEATURES])[:, 1])
model.fit(history[FEATURES], history['y'])


def psi(ref, cur, bins=10, eps=1e-4):
    edges = np.unique(np.quantile(ref, np.linspace(0, 1, bins + 1)))
    edges[0], edges[-1] = -np.inf, np.inf
    e = np.clip(np.histogram(ref, edges)[0] / len(ref), eps, None)
    a = np.clip(np.histogram(cur, edges)[0] / len(cur), eps, None)
    return float(np.sum((a - e) * np.log(a / e)))


production = calls[calls['date'] >= '2009-05'].copy()
production['batch'] = production['date'].dt.qyear.astype(str) + 'Q' + production['date'].dt.quarter.astype(str)
ref_score = model.predict_proba(history[FEATURES])[:, 1]
rows = []
for batch, g in production.groupby('batch'):
    start = g['date'].min()
    recent = calls[(calls['date'] < start) & (calls['date'] >= start - 6)]
    s_old = model.predict_proba(g[FEATURES])[:, 1]
    s_new = clone(model).fit(recent[FEATURES], recent['y']).predict_proba(g[FEATURES])[:, 1]
    rows.append({'batch': batch, 'rate': g['y'].mean(), 'pred static': s_old.mean(),
                 'pred retrained': s_new.mean(), 'score PSI': psi(ref_score, s_old),
                 'AUC static': roc_auc_score(g['y'], s_old), 'AUC retrained': roc_auc_score(g['y'], s_new)})
r = pd.DataFrame(rows)
q = np.arange(len(r))

fig, (a1, a2, a3) = plt.subplots(1, 3, figsize=(10, 3.4))
a1.plot(q, r['rate'] * 100, 'o-', color='gray', lw=2.5, ms=4, label='actual')
a1.plot(q, r['pred static'] * 100, 'o-', color='#f0883e', lw=2, ms=4, label='static model')
a1.plot(q, r['pred retrained'] * 100, 'o-', color='#3fb950', lw=2, ms=4, label='retrained')
a1.set_ylabel('% subscribing')
a1.set_ylim(0, 70)
a1.set_title('actual vs predicted rate', fontsize=9.5, loc='left')
a1.legend(fontsize=7.5, frameon=False, loc='upper left')

a2.bar(q, r['score PSI'], color='#58a6ff', width=0.6)
a2.axhline(0.25, color='#ff7b72', lw=1, ls='--')
a2.text(-0.4, 0.3, 'alert 0.25', color='#ff7b72', fontsize=8)
a2.set_title('PSI of the static model\'s score', fontsize=9.5, loc='left')
a2.set_ylim(0, 1.6)

a3.plot(q, r['AUC static'], 'o-', color='#f0883e', lw=2, ms=4, label='static model')
a3.plot(q, r['AUC retrained'], 'o-', color='#3fb950', lw=2, ms=4, label='retrained')
a3.axhline(val_auc, color='gray', lw=1, ls=':')
a3.text(len(r) - 1, val_auc - 0.018, f'validation {val_auc:.3f}', color='gray', fontsize=7.5, ha='right')
a3.set_ylim(0.5, 0.85)
a3.set_title('ROC AUC per quarter', fontsize=9.5, loc='left')
a3.legend(fontsize=7.5, frameon=False, loc='upper left')
for ax in (a1, a2, a3):
    ax.set_xticks(q, r['batch'], rotation=45, fontsize=7.5)
    ax.tick_params(axis='y', labelsize=8)
    ax.grid(alpha=0.2)
fig.suptitle('Bank-marketing model deployed in May 2009, followed quarter by quarter', fontsize=10)
plt.tight_layout()

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