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

# 1. nRot Data (Rotable Links)
bin_centers = list(range(0, 21))
frequencies = [
    10.78, 7.21, 12.25, 10.78, 11.97, 10.74, 8.44, 6.74, 4.40, 4.00, 
    3.53, 2.30, 1.74, 1.47, 0.67, 0.83, 0.55, 0.32, 0.48, 0.44, 0.36
]

# 2. Gaussian Fit Parameters
amplitude = 11.22
mean = 2.8
sd = 4.47

# Generate smooth X data for the curve
x_smooth = np.linspace(-2, 22, 300)

# Calculate Y using the Gaussian equation
y_smooth = amplitude * np.exp(-0.5 * ((x_smooth - mean) / sd)**2)

# 3. Define colors (nRot Traffic Light - Veber's Rule)
colors = []
for x in bin_centers:
    # Optimal Range: 0 to 9 (Safe Veber Zone)
    if 0 <= x <= 9:
        colors.append('green')
    # Caution Range: 10 to 15 (Flexibility Limit)
    elif 10 <= x <= 15:
        colors.append('gold')
    # Risk Rank: > 15 (Too flexible/entropic penalty)
    else:
        colors.append('firebrick')

# 4. Create the chart
plt.figure(figsize=(7, 6))

# A. Draw Bars
plt.bar(bin_centers, frequencies, width=0.8, color=colors, edgecolor='black', alpha=0.7, label='Data Frequency')

# B. Draw Trend Line
plt.plot(x_smooth, y_smooth, color='darkorange', linewidth=2, label='Gaussian Fit')

# 5. Tags and Titles
plt.xlabel('Number of Rotatable Bonds (nRot)', fontsize=12)
plt.ylabel('% Frequency', fontsize=12)
plt.title('Rotatable Bonds Distribution', fontsize=14)

# Adjust X axis
plt.xticks(bin_centers)
plt.xlim(-1, 21)

# 6. Custom Legend
legend_elements = [
    Line2D([0], [0], color='darkorange', lw=2, label=f'Fit (Mean={mean}, SD={sd})'),
    Patch(facecolor='green', edgecolor='black', alpha=0.7, label='Optimal (Veber Rule: ≤ 10)'),
    Patch(facecolor='gold', edgecolor='black', alpha=0.7, label='Caution (10 - 15)'),
    Patch(facecolor='firebrick', edgecolor='black', alpha=0.7, label='High Entropy Risk (> 15)')
]

plt.legend(handles=legend_elements, loc='upper right')
plt.grid(axis='y', linestyle='--', alpha=0.5)
plt.tight_layout()

plt.show()