#! /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 realistic ranges
polysyllables_per_word_range = np.linspace(0, 0.5, 50)  # Fraction of polysyllabic words (x-axis)
words_per_sentence_range = np.linspace(5, 30, 50)  # Average words per sentence (y-axis)

# Create a meshgrid for the heatmap
polysyllables_grid, words_grid = np.meshgrid(polysyllables_per_word_range, words_per_sentence_range)

# Calculate the number of polysyllables and sentences
words = words_grid  # Total words
polysyllables = polysyllables_grid * words  # Total polysyllabic words
sentences = 1  # Assume one sentence for simplicity

# Calculate SMOG for each combination
smog_heatmap = smog(polysyllables, sentences)

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

# Plot the heatmap
plt.figure(figsize=(10, 8))
plt.contourf(
    polysyllables_per_word_range, words_per_sentence_range, smog_heatmap_clipped,
    levels=20, cmap="cividis"
)
cbar = plt.colorbar()
cbar.set_label("SMOG Score (Clipped 0 to 20)", rotation=270, labelpad=20)
plt.title("SMOG Heatmap: Polysyllables per Word vs. Words per Sentence")
plt.xlabel("Polysyllables per Word")
plt.ylabel("Words per Sentence")
plt.grid(alpha=0.3, linestyle="--")
plt.savefig("smog_plot_clipped.png", transparent=False)