import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import numpy as np

# 1. Original Data (Nephrotoxicity)
bin_centers = [0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45,
               0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95]

frequencies = [0.12, 3.76, 8.62, 8.78, 7.36, 7.20, 5.74, 8.05, 8.01,
               8.86, 7.61, 6.31, 6.63, 5.83, 3.48, 1.78, 0.89, 0.85, 0.12]

# 2. Gaussian parameters
amp = 8.593
mean = 0.4039
std_dev = 0.2577

# 3. Generate Gaussian curve
x_fit = np.linspace(0, 1, 500)
y_fit = amp * np.exp(-((x_fit - mean)**2) / (2 * std_dev**2))

# 4. Colors by risk category
colors = []
for val in bin_centers:
    if val < 0.3:
        colors.append('green')
    elif val < 0.7:
        colors.append('gold')
    else:
        colors.append('firebrick')

# 5. Create figure
plt.figure(figsize=(7, 6))

# Bars (Observed data)
plt.bar(
    bin_centers,
    frequencies,
    width=0.04,
    color=colors,
    edgecolor='black',
    alpha=0.6,
    label='Observed Data'
)

# Gaussian fit (FIX: raw f-string para evitar warnings)
plt.plot(
    x_fit,
    y_fit,
    color='orange',
    linewidth=3,
    label=rf'Gaussian Fit\n($\mu$={mean:.4f}, $\sigma$={std_dev:.4f})'
)

# 6. Labels and title
plt.xlabel('Probability of Nephrotoxicity', fontsize=12)
plt.ylabel('% Frequency', fontsize=12)
plt.title('Nephrotoxicity Distribution', fontsize=14)

# Axes settings
plt.xticks(np.arange(0.0, 1.05, 0.1))
plt.xlim(0.0, 1.0)
plt.ylim(0, 12)

# 7. Custom legend
legend_elements = [
    Patch(facecolor='green', edgecolor='black', label='Low Risk'),
    Patch(facecolor='gold', edgecolor='black', label='Average Load'),
    Patch(facecolor='firebrick', edgecolor='black', label='High Risk'),
    plt.Line2D([0], [0], color='orange', lw=3, label='Fit (Mean=8.59, SD=0.40)'),
]

plt.legend(handles=legend_elements, loc='upper right', framealpha=0.95, fontsize=10)

# Grid and layout
plt.grid(axis='y', linestyle='--', alpha=0.5)
plt.tight_layout()

plt.show()