import marimo

__generated_with = "0.25.1"
app = marimo.App(width="medium")


@app.cell
def _():
    import math
    import marimo as mo

    return math, mo


@app.cell
def _(mo):
    delta = mo.ui.slider(
        -8.0,
        8.0,
        step=0.25,
        value=0.0,
        label="Logit difference Δ = z₁ − z₂",
        show_value=True,
    )
    delta
    return (delta,)


@app.cell
def _(delta, math):
    p1 = 1.0 / (1.0 + math.exp(-delta.value))
    p2 = 1.0 - p1
    return p1, p2


@app.cell
def _(delta, math, mo, p1, p2):
    width = 760
    height = 430
    left = 72
    right = 24
    top = 28
    bottom = 62
    plot_width = width - left - right
    plot_height = height - top - bottom

    def scale_x(value: float) -> float:
        return left + ((value + 8.0) / 16.0) * plot_width

    def scale_y(value: float) -> float:
        return top + (1.0 - value) * plot_height

    values = [-8.0 + i * 0.1 for i in range(161)]
    p1_points = " ".join(
        f"{scale_x(value):.2f},{scale_y(1.0 / (1.0 + math.exp(-value))):.2f}"
        for value in values
    )
    p2_points = " ".join(
        f"{scale_x(value):.2f},{scale_y(1.0 - 1.0 / (1.0 + math.exp(-value))):.2f}"
        for value in values
    )

    x_ticks = [-8, -4, 0, 4, 8]
    y_ticks = [0.0, 0.25, 0.5, 0.75, 1.0]
    x_tick_svg = "".join(
        f'<line x1="{scale_x(value):.2f}" y1="{top + plot_height}" '
        f'x2="{scale_x(value):.2f}" y2="{top + plot_height + 6}" stroke="#6b7280" />'
        f'<text x="{scale_x(value):.2f}" y="{top + plot_height + 26}" '
        f'text-anchor="middle" font-size="13" fill="#4b5563">{value}</text>'
        for value in x_ticks
    )
    y_tick_svg = "".join(
        f'<line x1="{left - 6}" y1="{scale_y(value):.2f}" '
        f'x2="{left}" y2="{scale_y(value):.2f}" stroke="#6b7280" />'
        f'<text x="{left - 12}" y="{scale_y(value) + 4:.2f}" '
        f'text-anchor="end" font-size="13" fill="#4b5563">{value:.2g}</text>'
        for value in y_ticks
    )

    current_x = scale_x(delta.value)
    current_y1 = scale_y(p1)
    current_y2 = scale_y(p2)

    svg = f"""
    <div style="font-family: system-ui, -apple-system, sans-serif; width: 100%;">
      <svg viewBox="0 0 {width} {height}" role="img"
           aria-label="Two-class Softmax probabilities as a function of the logit difference"
           style="width: 100%; height: auto; background: white; border-radius: 10px;">
        <line x1="{left}" y1="{top}" x2="{left}" y2="{top + plot_height}" stroke="#374151" stroke-width="1.5"/>
        <line x1="{left}" y1="{top + plot_height}" x2="{left + plot_width}" y2="{top + plot_height}" stroke="#374151" stroke-width="1.5"/>

        <line x1="{scale_x(0):.2f}" y1="{top}" x2="{scale_x(0):.2f}" y2="{top + plot_height}"
              stroke="#cbd5e1" stroke-dasharray="5 5"/>
        <line x1="{left}" y1="{scale_y(0.5):.2f}" x2="{left + plot_width}" y2="{scale_y(0.5):.2f}"
              stroke="#cbd5e1" stroke-dasharray="5 5"/>

        {x_tick_svg}
        {y_tick_svg}

        <polyline points="{p1_points}" fill="none" stroke="#2563eb" stroke-width="4"/>
        <polyline points="{p2_points}" fill="none" stroke="#f97316" stroke-width="4"/>

        <line x1="{current_x:.2f}" y1="{top}" x2="{current_x:.2f}" y2="{top + plot_height}"
              stroke="#94a3b8" stroke-dasharray="3 4"/>
        <circle cx="{current_x:.2f}" cy="{current_y1:.2f}" r="6" fill="#2563eb"/>
        <circle cx="{current_x:.2f}" cy="{current_y2:.2f}" r="6" fill="#f97316"/>

        <text x="{left + plot_width / 2:.2f}" y="{height - 14}" text-anchor="middle"
              font-size="15" fill="#374151">logit difference Δ = z₁ − z₂</text>
        <text x="18" y="{top + plot_height / 2:.2f}" text-anchor="middle"
              transform="rotate(-90 18 {top + plot_height / 2:.2f})"
              font-size="15" fill="#374151">probability</text>

      </svg>

      <div style="display:flex; gap:1.5rem; justify-content:center; align-items:center; flex-wrap:wrap; margin-top:0.5rem; color:#374151;">
        <strong>Δ = {delta.value:.2f}</strong>
        <span style="display:inline-flex; align-items:center; gap:0.4rem;">
          <span style="width:1.5rem; border-top:4px solid #2563eb;"></span>
          <span>p₁ = {p1:.3f}</span>
        </span>
        <span style="display:inline-flex; align-items:center; gap:0.4rem;">
          <span style="width:1.5rem; border-top:4px solid #f97316;"></span>
          <span>p₂ = {p2:.3f}</span>
        </span>
      </div>
    </div>
    """
    mo.Html(svg)
    return


if __name__ == "__main__":
    app.run()