#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

def gunning_fog(words_per_sentence, complex_word_ratio):
    # Calculate the Gunning Fog Index using the formula
    return 0.4 * (words_per_sentence + 100 * complex_word_ratio)

# Define discrete ranges for words per sentence and complex word ratio
words_per_sentence_range = np.arange(1, 31, 1)  # Discrete count of words per sentence (x-axis)
complex_word_ratio_range = np.arange(0, 0.51, 0.01)  # Complex word ratio (y-axis, from 0 to 1)

# Create a meshgrid for the heatmap
words_per_sentence_grid, complex_word_ratio_grid = np.meshgrid(words_per_sentence_range, complex_word_ratio_range)

# Calculate Gunning Fog Index for each combination
gunning_fog_heatmap = gunning_fog(words_per_sentence_grid, complex_word_ratio_grid)

# Clip Gunning Fog Index values to a realistic range (e.g., 0 to 20)
gunning_fog_heatmap_clipped = np.clip(gunning_fog_heatmap, 0, 20)

# Plot the heatmap with words per sentence on x-axis and complex word ratio on y-axis
plt.figure(figsize=(10, 8))
plt.contourf(
    words_per_sentence_range,  # X-axis: Words per sentence
    complex_word_ratio_range,  # Y-axis: Complex word ratio
    gunning_fog_heatmap_clipped,  # Z-values: Gunning Fog Index
    levels=40, cmap="plasma"
)

cbar = plt.colorbar()
cbar.set_label("Gunning Fog Index (Clipped 0 to 20)", rotation=270, labelpad=20)
cbar.ax.text(1, -0.025, "Kindergarden", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)
cbar.ax.text(1, 1.025, "Professor", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)

plt.title("Gunning Fog Heatmap")
plt.xlabel("Words per Sentence")
plt.ylabel("Complex Word Ratio (Polysyllabic Words / Total Words)")
plt.grid(alpha=0.3, linestyle="--")
plt.savefig("gunning_fog_plot.png", transparent=False)