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

# 1. nAtom Data
bin_centers = list(range(1, 61)) # Creates a list from 1 to 60
frequencies = [
    2.796, 0.679, 0.439, 1.038, 0.639, 0.639, 0.839, 0.879, 1.478, 2.236, 
    2.396, 2.236, 2.396, 2.676, 2.037, 2.915, 3.195, 2.995, 4.034, 4.073, 
    4.273, 4.433, 3.754, 3.474, 3.395, 3.474, 3.674, 3.435, 3.035, 3.115, 
    2.676, 2.356, 2.276, 1.597, 1.997, 1.558, 1.238, 0.719, 0.759, 1.158, 
    0.799, 0.919, 0.839, 0.359, 0.519, 0.359, 0.200, 0.240, 0.319, 0.200, 
    0.359, 0.240, 0.240, 0.280, 0.160, 0.080, 0.200, 0.280, 0.240, 0.160
]

# 2. Gaussian Fit Parameters
amplitude = 3.796
mean = 22.51
sd = 10.32

# Generate smooth X data for the curve
x_smooth = np.linspace(0, 65, 300)

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

# 3. Define colors (nAtom Traffic light system)
colors = []
for x in bin_centers:
    # Optimal Range: 15 to 35 heavy atoms
    if 15 <= x <= 35:
        colors.append('green')
    # Caution Range: Small (5-15) or Large (35-50)
    elif (5 <= x < 15) or (35 < x <= 50):
        colors.append('gold')
    # Risk Range: Very small (<5) or Very large (>50)
    else:
        colors.append('firebrick')

# 4. Create the plot
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. Labels and Titles
plt.xlabel('nAtom (Number of Heavy Atoms)', fontsize=12)
plt.ylabel('% Frequency', fontsize=12)
plt.title('nAtom Distribution', fontsize=14)

# Adjust X axis (show every 5 numbers to avoid clutter)
plt.xticks(np.arange(0, 65, 5), fontsize=10)
plt.xlim(0, 62)

# 6. Custom Legend
legend_elements = [
    Line2D([0], [0], color='darkorange', lw=2.5, label=f'Fit (Mean={mean}, SD={sd})'),
    Patch(facecolor='green', edgecolor='black', alpha=0.7, label='Optimal (15-35 Atoms)'),
    Patch(facecolor='gold', edgecolor='black', alpha=0.7, label='Caution (5-15, 35-50 Atoms)'),
    Patch(facecolor='firebrick', edgecolor='black', alpha=0.7, label='Risk / Atypical (<5, >50 Atoms)')
]

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

plt.show()