"""Compare an original cubic-cost plot with Knuth's integer badness routine.

Algorithm: badness(t,s) in https://tug.ctan.org/systems/knuth/dist/tex/tex.web.
All plotted inputs describe finite stretching. Write the PNG to cwd.
"""
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


STRETCH = 65536


def tex_badness(excess, stretch):
    """Integer algorithm, including its large-input branches and sentinel."""
    assert excess >= 0
    if excess == 0:
        return 0
    if stretch <= 0:
        return 10000
    if excess <= 7230584:
        scaled_ratio = excess * 297 // stretch
    elif stretch >= 1663497:
        scaled_ratio = excess // (stretch // 297)
    else:
        scaled_ratio = excess
    if scaled_ratio > 1290:
        return 10000
    return (scaled_ratio ** 3 + 131072) // 262144


def samples(ratios):
    return [tex_badness(round(float(r) * STRETCH), STRETCH) for r in ratios]


def main():
    examples = [(0.5, 32768, 12), (1, 65536, 100), (2, 131072, 800)]
    for ratio, excess, expected in examples:
        assert excess / STRETCH == ratio
        assert tex_badness(excess, STRETCH) == expected
    fig, axes = plt.subplots(1, 2, figsize=(9, 4.8), dpi=100, facecolor="white")
    fig.subplots_adjust(left=0.08, right=0.98, bottom=0.22, top=0.72, wspace=0.28)
    fig.suptitle("Spacing strain has a rapidly increasing cost", y=0.96,
                 fontsize=16, weight="bold", color="#172330")
    fig.text(0.5, 0.86, "Finite stretching; integer routine uses coefficient s = 65536 internal units",
             ha="center", fontsize=10, color="#45535d")
    for ax, lo, hi, title in zip(axes, [0, 0.35], [2.05, 0.65],
                                 ["Cubic growth", "Zoom: integer steps"]):
        ratios = np.linspace(lo, hi, 3001)
        ax.plot(ratios, 100 * ratios ** 3, color="#347ea8", linestyle="--",
                linewidth=1.7, label=r"Smooth guide $100r^3$")
        ax.step(ratios, samples(ratios), where="mid", color="#bd6b29",
                linewidth=1.6, label="TeX integer routine")
        ax.set_xlim(lo, hi)
        ax.set_ylim((0, 860) if lo == 0 else (3, 29))
        ax.set_title(title, fontsize=12, pad=12)
        ax.set_xlabel("Adjustment ratio r", fontsize=10)
        ax.set_ylabel("Badness b", fontsize=10)
        ax.grid(alpha=0.18)
        ax.tick_params(labelsize=9)
        for side in ["top", "right"]:
            ax.spines[side].set_visible(False)
    for (ratio, excess, value), xytext in zip(examples, [(0.8, 45), (0.58, 220), (1.22, 735)]):
        axes[0].plot(ratio, value, "o", color="#172330", markersize=4)
        axes[0].annotate(f"({ratio:g}, {value})", (ratio, value), xytext=xytext,
                         fontsize=9, color="#172330",
                         arrowprops={"arrowstyle": "-", "color": "#788994"})
    axes[1].plot(0.5, 12, "o", color="#172330", markersize=4)
    axes[1].annotate("r = 0.5 gives b = 12", (0.5, 12), xytext=(0.365, 24),
                     fontsize=9, color="#172330",
                     arrowprops={"arrowstyle": "-", "color": "#788994"})
    handles, labels = axes[0].get_legend_handles_labels()
    fig.legend(handles, labels, loc="lower center", bbox_to_anchor=(0.5, 0.055),
               ncol=2, frameon=False, fontsize=10)
    fig.text(0.5, 0.018, "Original numerical illustration; saturation and excessive shrinkage are outside the plotted range.",
             ha="center", fontsize=9, color="#45535d")
    fig.savefig(Path.cwd() / (Path(__file__).stem + ".png"),
                dpi=100, facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()