"""Compare greedy and optimal breaks in an original fixed-width toy model."""
from pathlib import Path
import math

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

WIDTHS = (3, 2, 2, 5)
LIMIT = 7
COLORS = ("#347ea8", "#329185", "#cb8742", "#8866a7")


def length(group):
    return sum(WIDTHS[i] for i in group) + len(group) - 1


def greedy():
    lines, current = [], []
    for i in range(len(WIDTHS)):
        if current and length(current + [i]) > LIMIT:
            lines.append(current)
            current = []
        current.append(i)
    return lines + [current]


def optimum():
    n = len(WIDTHS)
    costs, previous = [0] + [math.inf] * n, [None] * (n + 1)
    for end in range(1, n + 1):
        for start in range(end):
            extent = length(range(start, end))
            if extent <= LIMIT:
                candidate = costs[start] + (0 if end == n else (LIMIT - extent) ** 2)
                if candidate < costs[end]:
                    costs[end], previous[end] = candidate, start
    groups, end = [], n
    while end:
        start = previous[end]
        groups.append(list(range(start, end)))
        end = start
    return list(reversed(groups))


def main():
    fig, axes = plt.subplots(1, 2, figsize=(9, 4.6), dpi=100, facecolor="white")
    fig.subplots_adjust(left=0.045, right=0.98, bottom=0.2, top=0.76, wspace=0.25)
    fig.suptitle("Choosing breaks for the whole paragraph", y=0.97,
                 fontsize=16, weight="bold", color="#172330")
    fig.text(0.5, 0.885, "Word widths: (3, 2, 2, 5); spaces: 1; maximum line width: 7",
             ha="center", fontsize=11, color="#45535d")
    for ax, groups, name in zip(axes, [greedy(), optimum()],
                                ["Greedy choice", "Global optimum"]):
        total = sum((LIMIT - length(g)) ** 2 for g in groups[:-1])
        ax.set_title(f"{name}: total cost {total}", fontsize=12, pad=17)
        for row, group in enumerate(groups):
            y = 3 - row
            ax.add_patch(Rectangle((0, y - 0.27), LIMIT, 0.54,
                                   facecolor="#f2f4f6", edgecolor="#9aabb5",
                                   linewidth=1, linestyle="--"))
            x = 0
            for index in group:
                width = WIDTHS[index]
                ax.add_patch(Rectangle((x, y - 0.27), width, 0.54,
                                       facecolor=COLORS[index], edgecolor="white"))
                ax.text(x + width / 2, y, f"word {index + 1}",
                        ha="center", va="center", color="white", fontsize=10)
                x += width + 1
            cost = (LIMIT - length(group)) ** 2
            label = f"cost {cost}" if row < len(groups) - 1 else "free"
            ax.text(7.16, y, label, va="center", fontsize=10, color="#45535d")
        ax.set_xlim(-0.1, 8.5)
        ax.set_ylim(0.4, 3.65)
        ax.set_xticks(range(8))
        ax.set_yticks([])
        ax.tick_params(axis="x", labelsize=9)
        for side in ["left", "right", "top"]:
            ax.spines[side].set_visible(False)
        ax.spines["bottom"].set_color("#9aabb5")
    fig.text(0.5, 0.115, "Cost = squared unused width; the final line is unpenalized.",
             ha="center", fontsize=11, color="#172330")
    fig.text(0.5, 0.045, "Original schematic with hypothetical widths and a simplified cost function.",
             ha="center", fontsize=10, color="#45535d")
    fig.savefig(Path.cwd() / (Path(__file__).stem + ".png"), dpi=100,
                facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()