#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

def flesch_reading_ease(words_per_sentence, syllables_per_word):
    # Calculate the Flesch Reading Ease using the formula
    return 206.835 - 1.015 * words_per_sentence - 84.6 * syllables_per_word

# Define discrete ranges for words per sentence and syllables per word
words_per_sentence_range = np.arange(1, 31, 1)  # Discrete count of words per sentence (x-axis)
syllables_per_word_range = np.arange(1, 3.1, 0.1)  # Average syllables per word (y-axis)

# Create a meshgrid for the heatmap
words_per_sentence_grid, syllables_per_word_grid = np.meshgrid(words_per_sentence_range, syllables_per_word_range)

# Calculate Flesch Reading Ease for each combination
flesch_heatmap = flesch_reading_ease(words_per_sentence_grid, syllables_per_word_grid)

# Clip Flesch Reading Ease values to a realistic range (-100 to 120)
flesch_heatmap_clipped = np.clip(flesch_heatmap, 0, 100)

# Plot the heatmap with words per sentence on x-axis and syllables per word on y-axis
plt.figure(figsize=(10, 8))
plt.contourf(
    words_per_sentence_range,  # X-axis: Words per sentence
    syllables_per_word_range,  # Y-axis: Syllables per word
    flesch_heatmap_clipped,  # Z-values: Flesch Reading Ease
    levels=40, cmap="plasma_r"
)

cbar = plt.colorbar()
cbar.set_label("Flesch Reading Ease (Clipped 0 to 100)", rotation=270, labelpad=20)
cbar.ax.invert_yaxis()
cbar.ax.text(1, -0.025, "Very Easy", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)
cbar.ax.text(1, 1.025, "Very Difficult", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)

plt.title("Flesch Reading Ease Heatmap")
plt.xlabel("Words per Sentence")
plt.ylabel("Syllables per Word")
plt.grid(alpha=0.3, linestyle="--")

plt.savefig("flesch_reading_ease_plot.png", transparent=False)