#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

def colemanliau(characters, words, sentences):
    L = (characters / words) * 100  # Average characters per 100 words
    S = (sentences / words) * 100  # Average sentences per 100 words
    return 0.0588 * L - 0.296 * S - 15.8

# Define realistic ranges for characters per word and words per sentence
characters_per_words_range = np.linspace(1, 15, 50)  # Characters per word (x-axis)
sentences_per_words_range = np.linspace(1, 30, 50)  # Words per sentence (y-axis)

# Create a meshgrid for the heatmap
characters_grid, sentences_grid = np.meshgrid(characters_per_words_range, sentences_per_words_range)

words = 100

# Calculate CLI for each combination in the grid
cli_heatmap = colemanliau(
    characters_grid * words,  # Total characters
    words,  # Total words
    sentences_grid  # Total sentences
)

# Clip CLI values to a realistic range (e.g., -10 to 20)
cli_heatmap_clipped = np.clip(cli_heatmap, 1, 20)

# Plot the heatmap
plt.figure(figsize=(10, 8))
plt.contourf(
    characters_per_words_range, #x-axis
    sentences_per_words_range,  #y-axis
    cli_heatmap_clipped,       #z-axis
    levels=40, cmap="coolwarm"
)
cbar = plt.colorbar()
cbar.set_label("CLI Score (Clipped 1 to 20)", rotation=270, labelpad=20)
cbar.ax.text(1, 1.025, "Professor", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)
cbar.ax.text(1, -0.025, "1st Grade (US)", ha='left', va='center', fontsize=10, transform=cbar.ax.transAxes)

plt.title("Coleman-Liau Heatmap")
plt.xlabel("Average Characters per 100 Word")
plt.ylabel("Average Sentences per 100 Words")
plt.grid(alpha=0.3, linestyle="--")
plt.savefig("coleman_liau_plot.png", transparent=False)