"""A three-step Softmax animation for 'Where Do Probabilities Come From?'

The fixed logits and numerical values match the article's toy table:
cats = 0, dogs = 1, pizza = 2. This is a model output,
not an empirical distribution or a trained model result.

To render locally (Python 3.13+):
    pip install manimgx==0.1.1
    manimgx render softmax-from-logits.py
"""

import math

import manimgx as m


class SoftmaxFromLogits(m.Scene):
    def construct(self) -> None:
        names = ("cats", "dogs", "pizza")
        logits = (0.0, 1.0, 2.0)
        weights = tuple(math.exp(z) for z in logits)
        normalizer = sum(weights)
        probabilities = tuple(w / normalizer for w in weights)

        colors = (m.BLUE, m.TEAL, m.YELLOW)
        positions = (-3.4, 0.0, 3.4)
        baseline_y = -2.1
        max_height = 2.4

        def stage(values: tuple[float, ...], maximum: float):
            bars = []
            amounts = []
            for x, value, color in zip(positions, values, colors, strict=True):
                height = max(0.04, max_height * value / maximum)
                bar = m.Rectangle(
                    width=1.05,
                    height=height,
                    fill_color=color,
                    fill_opacity=0.85,
                    stroke_color=color,
                    stroke_width=1,
                ).move_to([x, baseline_y + height / 2, 0])
                amount = m.Text(
                    f"{value:.4f}", font_size=29, color=color
                ).move_to([x, baseline_y + height + 0.30, 0])
                bars.append(bar)
                amounts.append(amount)
            return bars, amounts

        heading = m.Text(
            "Softmax: from logits to probabilities", font_size=42
        ).move_to(m.UP * 3.25)
        equation = m.MathTex(
            r"p_i=\frac{e^{z_i}}{\sum_j e^{z_j}}", font_size=47
        ).move_to(m.UP * 2.3)
        baseline = m.Line(
            [-4.75, baseline_y, 0],
            [4.75, baseline_y, 0],
            stroke_width=2,
            color=m.GREY,
        )
        labels = [
            m.Text(name, font_size=31).move_to([x, -2.68, 0])
            for name, x in zip(names, positions, strict=True)
        ]

        title = m.Text(
            "1 / Logits: unconstrained scores", font_size=31
        ).move_to(m.UP * 1.38)
        note = m.Text(
            "These numbers are not probabilities.", font_size=25
        ).move_to(m.DOWN * 3.36)
        bars, amounts = stage(logits, 2.0)

        self.play(
            m.Write(heading), m.Write(equation), m.Create(baseline),
            run_time=1.7,
        )
        self.play(
            m.Write(title), m.FadeIn(note),
            *(m.FadeIn(label) for label in labels),
            *(m.FadeIn(bar) for bar in bars),
            *(m.FadeIn(amount) for amount in amounts),
            run_time=1.7,
        )
        self.wait(1.0)

        phases = (
            (
                "2 / Exponentiate: positive weights",
                weights,
                max(weights),
                "Weights have a common sum, Z = 11.1073.",
            ),
            (
                "3 / Normalize: a probability distribution",
                probabilities,
                max(probabilities),
                "The shares sum to 1; the weight ratios do not change.",
            ),
        )

        for heading_text, values, maximum, note_text in phases:
            next_title = m.Text(
                heading_text, font_size=31
            ).move_to(m.UP * 1.38)
            next_note = m.Text(
                note_text, font_size=24
            ).move_to(m.DOWN * 3.36)
            next_bars, next_amounts = stage(values, maximum)

            self.play(
                m.ReplacementTransform(title, next_title),
                m.ReplacementTransform(note, next_note),
                *(m.Transform(before, after)
                  for before, after in zip(bars, next_bars, strict=True)),
                *(m.Transform(before, after)
                  for before, after in zip(amounts, next_amounts, strict=True)),
                run_time=2.0,
            )
            title = next_title
            note = next_note
            self.wait(1.6)