import time
from io import StringIO

import matplotlib.pyplot as plt
import numpy as np
from sklearn.compose import make_column_transformer
from sklearn.datasets import fetch_openml
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# The lesson's credit pipeline, timed for calls carrying 1 to 10,000 rows
# (timings depend on the machine that runs this script)
credit = fetch_openml(data_id=31, as_frame=True)
num = ['duration', 'credit_amount', 'age']
cat = ['checking_status', 'savings_status', 'purpose', 'employment', 'housing']
X, y = credit.data[num + cat], (credit.target == 'bad').astype(int)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=0)
pipe = make_pipeline(make_column_transformer((StandardScaler(), num),
                                             (OneHotEncoder(handle_unknown='ignore'), cat)),
                     LogisticRegression(max_iter=1000)).fit(X_train, y_train)
big = X_test.sample(10_000, replace=True, random_state=0)


def call_ms(rows, repeats=7):
    best = np.inf
    for _ in range(repeats):
        start = time.perf_counter()
        pipe.predict_proba(rows)
        best = min(best, time.perf_counter() - start)
    return 1000 * best


sizes = np.array([1, 3, 10, 30, 100, 300, 1000, 3000, 10000])
per_call = np.array([np.median([call_ms(big.iloc[i:i + n]) for i in range(0, 50, 10)]) for n in sizes])
per_row = per_call / sizes

fig, (a1, a2) = plt.subplots(1, 2, figsize=(9, 3.4))
a1.plot(sizes, per_call, 'o-', color='#58a6ff', ms=5, lw=2)
a1.set_title('time per call', fontsize=10, loc='left')
a1.set_ylabel('milliseconds per call')
a1.annotate(f'{per_call[0]:.1f} ms for 1 row', xy=(1, per_call[0]), xytext=(2, per_call[0] * 2.4),
            fontsize=8, color='gray', arrowprops=dict(arrowstyle='->', color='gray'))
a1.annotate(f'{per_call[-1]:.1f} ms for 10,000 rows', xy=(10000, per_call[-1]), xytext=(40, per_call[-1] * 1.4),
            fontsize=8, color='gray', arrowprops=dict(arrowstyle='->', color='gray'))
a1.set_ylim(per_call.min() / 2, per_call.max() * 5)
a2.plot(sizes, per_row, 'o-', color='#f0883e', ms=5, lw=2)
a2.set_title('time per row', fontsize=10, loc='left')
a2.set_ylabel('milliseconds per row')
a2.annotate(f'{per_row[0] / per_row[-1]:,.0f}x cheaper per row', xy=(10000, per_row[-1]),
            xytext=(12, per_row[-1] * 3), fontsize=8, color='gray',
            arrowprops=dict(arrowstyle='->', color='gray'))
for ax in (a1, a2):
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_xlabel('rows per call')
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.25, which='major')
fig.suptitle('predict_proba of the credit pipeline: the cost is per call, not per row  (machine-dependent)',
             fontsize=10)
plt.tight_layout()

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