import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

# The simulated fraud screen of the lesson: 20,000 transactions, about 2% fraudulent
rng = np.random.default_rng(0)
n = 20_000
y = (rng.random(n) < 0.02).astype(int)
X = rng.normal(size=(n, 5)) + y[:, None] * np.array([1.2, 1.0, 0.8, 0.8, 0.6])
X_tmp, X_test, y_tmp, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
X_train, X_val, y_train, y_val = train_test_split(X_tmp, y_tmp, test_size=0.25, stratify=y_tmp,
                                                  random_state=0)

fig, axes = plt.subplots(2, 1, figsize=(8, 5.2), sharex=True)
bins = np.linspace(0, 1, 41)
for ax, cw, name in zip(axes, [None, 'balanced'], ['no weights', "class_weight='balanced'"]):
    model = make_pipeline(StandardScaler(), LogisticRegression(class_weight=cw)).fit(X_train, y_train)
    p = model.predict_proba(X_test)[:, 1]
    neg, pos = p[y_test == 0], p[y_test == 1]
    ax.hist(neg, bins=bins, color='#58a6ff', alpha=0.75, label=f'legitimate ({len(neg):,})')
    ax.hist(pos, bins=bins, color='#ff7b72', alpha=0.85, label=f'fraud ({len(pos)})')
    ax.axvline(0.5, color='gray', ls='--', lw=1)
    ax.set_yscale('log')
    ax.set_ylim(0.7, 6000)
    ax.set_ylabel('transactions (log)')
    ax.text(0.52, 1500, f'flagged at 0.5: {(pos >= 0.5).sum()} of {len(pos)} frauds '
            f'(recall {(pos >= 0.5).mean():.2f})\n'
            f'+ {(neg >= 0.5).sum()} legitimate  ·  mean score {p.mean():.3f}  ·  '
            f'ROC AUC {roc_auc_score(y_test, p):.3f}', fontsize=8, va='top')
    ax.set_title(f'Logistic regression, {name}', fontsize=9.5, loc='left')
    ax.legend(fontsize=7.5, loc='upper left', bbox_to_anchor=(0.12, 1.0))
    ax.grid(color='gray', alpha=0.2, lw=0.5)
axes[1].set_xlabel('predicted probability of fraud (test set, 4,000 transactions, 78 frauds)')

plt.tight_layout()

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

plt.close()