import pandas as pd
import numpy as np
from matplotlib import pyplot as plt

from sklearn.datasets import make_hastie_10_2

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.calibration import CalibratedClassifierCV

X, y = make_hastie_10_2(50000)
y[y == -1] = 0

X_0 = X[y == 0, :]
X_1 = X[y == 1, :]

y_0 = y[y == 0]
y_1 = y[y == 1]

# Discard half of the sample with y==0
X = np.vstack([X_0[:int(len(y_1) / 2), :], X_1])
y = np.hstack([y_0[:int(len(y_1) / 2)], y_1])

# Compute weigths to 'unbalance' the dataset'
weight = (y==0) + 1

X_train, X_test, y_train, y_test, weight_train, weight_test = train_test_split(X, y, weight)

calib = CalibratedClassifierCV(RandomForestClassifier(n_estimators=5, max_depth=3), ensemble=False, n_jobs=-1)
calib.fit(X_train, y_train, sample_weight=weight_train)

pred = calib.predict_proba(X_test)[:, 1]

# Check calibration in a way which takes into account that both classes have equal importance despite
# class 0 being less frequent

df_target_pred = pd.DataFrame([y_test, pred]).transpose()
df_target_pred.columns = ["target", "pred"]

hist_0 = np.histogram(df_target_pred.loc[df_target_pred.target == 0, 'pred'], bins=np.linspace(0, 1, 6), density=True)
hist_1 = np.histogram(df_target_pred.loc[df_target_pred.target == 1,  'pred'], bins=np.linspace(0, 1, 6), density=True)

plt.bar(hist_0[1][:-1], hist_0[0], align='edge', label='0', alpha=0.5, width=0.2)
plt.bar(hist_1[1][:-1], hist_1[0], align='edge', label='1', alpha=0.5, width=0.2)
plt.ylabel('Prediction histograms')
plt.xlabel('Predicted score')
plt.legend()
ax2 = plt.gca().twinx()
ax2.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated")
ax2.plot((hist_0[1][:-1] + hist_0[1][1:]) / 2, hist_1[0] / (hist_0[0] + hist_1[0] + 1e-10))
ax2.grid(False)
ax2.set_ylabel('Fraction of positives')
plt.show()