#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

def smog(polysyllables, sentences):
    # Calculate SMOG using the formula
    return 1.043 * np.sqrt(polysyllables * (30 / sentences)) + 3.1291

# Define discrete ranges for polysyllabic words and sentences
polysyllables_range = np.arange(1, 51, 1)  # Discrete count of polysyllabic words (x-axis)
sentences_range = np.arange(1, 31, 1)  # Discrete count of sentences (y-axis)

# Create a meshgrid for the heatmap, where the dimensions match
sentences_grid, polysyllables_grid = np.meshgrid(sentences_range, polysyllables_range)

# Calculate SMOG for each combination
smog_heatmap = smog(polysyllables_grid, sentences_grid)

# Clip SMOG values to a realistic range (e.g., 0 to 20)
smog_heatmap_clipped = np.clip(smog_heatmap, 0, 20)

# Plot the heatmap with flipped axes
plt.figure(figsize=(10, 8))
plt.contourf(
    sentences_range,  # X-axis: Number of sentences
    polysyllables_range,  # Y-axis: Polysyllabic words
    smog_heatmap_clipped,  # Z-values: SMOG index
    levels=40, cmap="plasma"
)

cbar = plt.colorbar()
cbar.set_label("SMOG Score (Clipped 0 to 20)", rotation=270, labelpad=20)
cbar.ax.text(1, -0.025, "Elementary (4th Grade)", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)
cbar.ax.text(1, 1.025, "Post-Graduate", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)

plt.title("SMOG Heatmap")
plt.xlabel("Number of Sentences")
plt.ylabel("Number of Polysyllabic Words")
plt.grid(alpha=0.3, linestyle="--")
plt.savefig("smog_plot.png", transparent=False)