import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

# Labels
x_labels = ['B5', 'B4', 'B3', 'B2', 'B1', 'N0', 'H1', 'H2', 'H3', 'H4', 'H5']
y_labels = ['B5', 'B4', 'B3', 'B2', 'B1', 'N0', 'H1', 'H2', 'H3', 'H4', 'H5']

# Data
data = np.array([
    [6.1, 11.7, 15.6, 20.6, 24.0, 26.0, 30.9, 40.7, 51.4, 54.9, 76.2],
    [10.4, 13.2, 19.6, 22.7, 31.9, 36.8, 44.5, 50.8, 68.0, 74.6, 88.2],
    [10.2, 16.7, 24.4, 25.1, 32.2, 42.6, 53.7, 60.0, 74.3, 78.8, 88.2],
    [11.5, 18.0, 24.8, 29.1, 35.6, 44.9, 54.1, 66.4, 75.5, 81.3, 90.0],
    [10.2, 18.8, 26.0, 31.6, 39.3, 48.7, 60.7, 71.4, 78.5, 83.4, 90.5],
    [12.5, 22.6, 26.4, 34.3, 42.2, 56.8, 63.3, 71.4, 80.7, 83.3, 89.5],
    [14.4, 24.5, 28.7, 40.0, 49.2, 60.2, 68.0, 72.3, 82.2, 83.4, 92.5],
    [12.7, 26.3, 33.5, 42.6, 53.6, 61.9, 68.8, 75.1, 80.9, 83.8, 92.0],
    [13.2, 26.9, 40.9, 46.8, 56.9, 65.7, 72.5, 75.8, 84.4, 86.9, 93.2],
    [15.8, 31.3, 43.6, 49.9, 64.2, 68.6, 75.9, 77.0, 85.4, 88.9, 95.0],
    [18.8, 39.1, 54.7, 64.0, 70.1, 79.6, 77.7, 81.8, 89.9, 89.4, 96.4]

])

# Meshgrid for 3D plotting
x = np.arange(len(x_labels))
y = np.arange(len(y_labels))
x, y = np.meshgrid(x, y)

# Plot
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(x, y, data, cmap='viridis', edgecolor='k')

# Axis formatting
ax.set_xticks(np.arange(len(x_labels)))
ax.set_yticks(np.arange(len(y_labels)))
ax.set_xticklabels(x_labels)
ax.set_yticklabels(y_labels)
ax.set_xlabel("sma5_deriv1_1h_bin")
ax.set_ylabel("mid_smooth_3_deriv1_bin")
ax.set_zlabel("Probabilité de hausse (%)")
ax.set_title("Probabilité de hausse pour futur_percent_3h en 3D")

fig.colorbar(surf, shrink=0.5, aspect=10)

plt.tight_layout()
plt.show()