import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import pandas as pd

# style
plt.style.use(["science", "grid"])

# read in SHAP and feat values
feat_val = pd.read_csv("Frontend/xai_experiment/immobilien_features.csv", index_col=0)
shap_val = pd.read_csv("Frontend/xai_experiment/immobilien_shap.csv", index_col=0)
# # ... combine SHAPS for plot
# shap_val["Frankfurt"] = shap_val["Frankfurt"] + shap_val["Muenchen"]
shap_val = shap_val.applymap(lambda x: round(x/50)*50)

# get figures and subplots
fig = plt.figure(figsize=(7.5, 5), constrained_layout=False, dpi=200)
spec = gridspec.GridSpec(ncols=3, nrows=1, figure=fig, width_ratios=[1, 1, 2])
ax = [fig.add_subplot(spec[0, 0]), fig.add_subplot(spec[0, 1]), fig.add_subplot(spec[0, 2])]
titles = ["Stadt", "Balkon/Terasse", "Anteil Grünenwähler im Stadtteil\n(innerstädt. Vergleich):"]
x_ranges = [[-0.5, 1.5], [-0.5, 1.5], [0, 4]]
x_ticks = [[0, 1], [0, 1], [1, 2, 3]]
y_ticks = [-5000, -4000, -3000, -2000, -1000, 0, 1000, 2000, 3000, 4000, 5000]
x_ticklabels = [["Hamburg", "München"], ["Nein", "Ja"], ["unterdurch-\nschnittlich", "durchschnittlich", "überdurch-\nschnittlich"]]
y_ticklabels = ["-5000", "-4000", "-3000", "-2000", "-1000", "0", "+1000", "+2000", "+3000", "+4000", "+5000"]
y_lims = [-5000, 5000]

for i, col in enumerate(["Stadt", "balcony", "Anteil Gruenenwaehler"]):
    foc_ax = ax[i]
    foc_ax.set_xlim(x_ranges[i])
    foc_ax.set_xticks(x_ticks[i])
    foc_ax.set_xticklabels(x_ticklabels[i], rotation=45, size=13)
    foc_ax.set_ylim(y_lims)
    if i == 0:
        foc_ax.set_ylabel("Einfluss auf Preis [EUR/qm]\n", size=12)
    foc_ax.set_yticks(y_ticks)
    foc_ax.set_yticklabels(y_ticklabels)
    foc_ax.set_title(titles[i], size=14)
    foc_ax.axhspan(0, y_lims[1], facecolor='green', alpha=0.05)
    foc_ax.axhspan(y_lims[0], 0, facecolor='red', alpha=0.05)
    if col != "Stadt":
        foc_ax.scatter(x=feat_val[col].values, y=shap_val[col].values, s=10)
    else:
        foc_ax.scatter(x=feat_val["Muenchen"].values, y=shap_val[col].values, s=10)

plt.tight_layout()
# plt.show()
plt.savefig("Frontend/_static/global_explanations_M_HH.jpg")