import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from imblearn.over_sampling import SMOTE, RandomOverSampler
from imblearn.pipeline import make_pipeline as make_imb_pipeline
from imblearn.under_sampling import RandomUnderSampler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import precision_recall_curve, precision_score, recall_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

# The five remedies of section 9 on the simulated fraud data (same code as in the lesson):
# their precision-recall curves on the test set almost coincide; at 0.5 they sit far apart,
# at the F1 threshold tuned on validation they meet again.
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)

models = {'no weights': make_pipeline(StandardScaler(), LogisticRegression()),
          'balanced weights': make_pipeline(StandardScaler(), LogisticRegression(class_weight='balanced')),
          'undersampling': make_imb_pipeline(StandardScaler(), RandomUnderSampler(random_state=0),
                                             LogisticRegression()),
          'oversampling': make_imb_pipeline(StandardScaler(), RandomOverSampler(random_state=0),
                                            LogisticRegression()),
          'SMOTE': make_imb_pipeline(StandardScaler(), SMOTE(random_state=0), LogisticRegression())}
colors = ['#58a6ff', '#f0883e', '#3fb950', '#bc8cff', '#ff7b72']

fig, ax = plt.subplots(figsize=(8.6, 5.4))
for (name, model), c in zip(models.items(), colors):
    model.fit(X_train, y_train)
    prec, rec, thr = precision_recall_curve(y_val, model.predict_proba(X_val)[:, 1])
    t_best = thr[np.argmax(2 * prec[:-1] * rec[:-1] / (prec[:-1] + rec[:-1] + 1e-12))]
    p = model.predict_proba(X_test)[:, 1]
    prec_t, rec_t, _ = precision_recall_curve(y_test, p)
    ax.plot(rec_t, prec_t, color=c, lw=1.1, alpha=0.75, drawstyle='steps-post')
    for t, marker, size in [(0.5, 'o', 9), (t_best, '*', 15)]:
        f = p >= t
        ax.plot(recall_score(y_test, f), precision_score(y_test, f), marker, color=c, ms=size,
                mfc='none' if marker == 'o' else c, mew=1.8, ls='none')
    ax.plot([], [], color=c, lw=2, label=f'{name}: tuned t = {t_best:.2f}')

ax.plot([], [], 'o', color='gray', mfc='none', mew=1.6, ms=8, label='threshold 0.5')
ax.plot([], [], '*', color='gray', ms=12, label='F1 threshold tuned on validation')
ax.annotate('at 0.5, no weights:\nhigh precision, low recall', xy=(0.231, 0.783), xytext=(0.30, 0.86),
            fontsize=8.5, color='gray', arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.annotate('at 0.5, the four rebalanced models:\nlow precision, high recall', xy=(0.82, 0.10),
            xytext=(0.66, 0.42), fontsize=8.5, color='gray',
            arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.annotate('tuned thresholds:\nthe five meet again', xy=(0.45, 0.47), xytext=(0.50, 0.64),
            fontsize=8.5, color='gray', arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.set_xlabel('recall (test set)', color='gray')
ax.set_ylabel('precision (test set)', color='gray')
ax.set_title('Five remedies, one precision–recall curve', fontsize=10, color='gray')
ax.set_xlim(0, 1)
ax.set_ylim(0, 1.02)
ax.grid(color='gray', alpha=0.2, lw=0.5)
ax.tick_params(colors='gray')
for spine in ax.spines.values():
    spine.set_color('gray')
ax.legend(fontsize=8, loc='upper right', frameon=False, labelcolor='gray')

plt.tight_layout()

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

plt.close()