encoder/src/demo/copilot_render.py

"""HTML rendering for the Co-Pilot demo.

Each function takes typed data (a CastMember, a CopilotResult, etc.)
and returns a self-contained HTML snippet. Gradio injects the snippets
into `gr.HTML` components.

Visual discipline:
  - Single typeface (Inter for body, JetBrains Mono for numbers).
  - Three colors only: page background near-white, near-black ink,
    one accent color per class (green/yellow/red). No gradient palettes,
    no shadows, no neon. Liquid demo aesthetic.
  - Inline styles only (no external CSS file). The renderer is one
    self-contained module so a future maintainer can edit visuals
    without touching the inference path.

The renderer NEVER reaches into model internals. It consumes the
serializable result objects produced by `copilot_inference.py`.
"""

from __future__ import annotations

import html
from typing import Iterable

from encoder.src.demo.copilot_inference import (
    LABEL_COLORS,
    LABEL_NAMES,
    CastMember,
    CopilotResult,
)


# ---- shared constants ----

_PAGE_BG = "#fafafa"
_INK = "#171717"
_INK_DIM = "#525252"
_BORDER = "rgba(0,0,0,0.08)"
_CARD_BG = "#ffffff"

_TONE_COLORS: dict[str, str] = {
    "angry": "#ef4444",
    "confused": "#a855f7",
    "casual": "#3b82f6",
    "formal": "#0ea5e9",
    "terse": "#737373",
}


# ---- cast strip ----

def render_cast_strip(cast: list[CastMember], selected_idx: int) -> str:
    """Six clickable archetype cards across the top.

    The card buttons are styled as anchors but the *click* is wired
    by Gradio via separate `gr.Button` components — this function
    only renders the visual chrome that mirrors the selected state.
    Selection is driven by the `selected_idx` argument.
    """
    cards: list[str] = []
    for i, m in enumerate(cast):
        is_sel = i == selected_idx
        accent = LABEL_COLORS[m.expected_label]
        border = accent if is_sel else _BORDER
        bg = "#ffffff" if not is_sel else _hex_with_alpha(accent, 0.06)
        cards.append(f"""
        <div style="
            flex: 1 1 0;
            min-width: 0;
            padding: 12px 14px;
            background: {bg};
            border: 1.5px solid {border};
            border-radius: 10px;
            transition: all 0.15s ease;
        ">
            <div style="
                font-size: 13px;
                font-weight: 600;
                color: {_INK};
                margin-bottom: 4px;
                overflow: hidden;
                text-overflow: ellipsis;
                white-space: nowrap;
            ">{html.escape(m.display_name)}</div>
            <div style="
                font-size: 11px;
                color: {_INK_DIM};
                font-family: 'JetBrains Mono', ui-monospace, monospace;
                text-transform: uppercase;
                letter-spacing: 0.04em;
            ">
                <span style="color: {accent};">●</span>
                expected: {html.escape(_short_label(m.expected_label_name))}
            </div>
        </div>
        """)
    return f"""
    <div style="
        display: flex;
        gap: 10px;
        margin-bottom: 18px;
    ">
        {''.join(cards)}
    </div>
    """


# ---- complaint quote ----

def render_complaint(member: CastMember) -> str:
    """Large pull-quote of what the customer wrote."""
    tone_color = _TONE_COLORS.get(member.tone, _INK_DIM)
    return f"""
    <div style="
        background: {_CARD_BG};
        border: 1px solid {_BORDER};
        border-left: 4px solid {tone_color};
        border-radius: 12px;
        padding: 22px 26px;
        margin-bottom: 18px;
    ">
        <div style="
            font-size: 11px;
            color: {_INK_DIM};
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            margin-bottom: 10px;
        ">
            Complaint received · tone:
            <span style="color: {tone_color}; font-weight: 600;">{html.escape(member.tone)}</span>
        </div>
        <div style="
            font-size: 18px;
            color: {_INK};
            line-height: 1.5;
            font-style: italic;
            font-weight: 400;
        ">
            &ldquo;{html.escape(member.complaint_text)}&rdquo;
        </div>
    </div>
    """


# ---- timeline ----

def render_timeline(
    disputed_idx: int,
    top_k_positions: Iterable[int] | None = None,
    attribution_probs: Iterable[float] | None = None,
    num_positions: int = 64,
) -> str:
    """64 dots; the disputed one is a red star, top-k are accent glows.

    If `top_k_positions` is None (pre-inference), all dots are neutral
    except the disputed star. If `attribution_probs` is provided, dot
    opacity scales with per-position attribution probability so the
    cluster around the highlighted positions reads as a glow gradient.
    """
    # Explicit None check: numpy arrays raise ValueError under `or []`.
    top_k_set: set[int] = set()
    if top_k_positions is not None:
        top_k_set = {int(i) for i in top_k_positions}
    probs = list(attribution_probs) if attribution_probs is not None else None

    dots: list[str] = []
    for i in range(num_positions):
        if i == disputed_idx:
            dots.append(_dot_disputed(i))
        elif i in top_k_set:
            alpha = float(probs[i]) if probs is not None else 1.0
            dots.append(_dot_attributed(i, alpha))
        else:
            faint = float(probs[i]) * 0.4 if probs is not None else 0.0
            dots.append(_dot_neutral(i, faint))

    return f"""
    <div style="
        background: {_CARD_BG};
        border: 1px solid {_BORDER};
        border-radius: 12px;
        padding: 20px 22px;
        margin-bottom: 18px;
    ">
        <div style="
            font-size: 11px;
            color: {_INK_DIM};
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            margin-bottom: 14px;
            display: flex;
            justify-content: space-between;
        ">
            <span>64-transaction history · oldest → newest</span>
            <span>
                <span style="color: #ef4444;">★</span> disputed &nbsp;
                <span style="color: #f59e0b;">●</span> contributed
            </span>
        </div>
        <div style="
            display: flex;
            justify-content: space-between;
            align-items: center;
            gap: 2px;
        ">
            {''.join(dots)}
        </div>
    </div>
    """


def _dot_disputed(i: int) -> str:
    return f"""
    <div title="position {i} · disputed" style="
        width: 14px; height: 14px;
        color: #ef4444;
        font-size: 18px;
        line-height: 14px;
        text-align: center;
        font-weight: 700;
    ">★</div>
    """


def _dot_attributed(i: int, alpha: float) -> str:
    glow = max(0.4, min(1.0, alpha))
    return f"""
    <div title="position {i} · contributed ({alpha:.2f})" style="
        width: 12px; height: 12px;
        border-radius: 50%;
        background: rgba(245, 158, 11, {glow});
        box-shadow: 0 0 8px rgba(245, 158, 11, {glow * 0.7});
    "></div>
    """


def _dot_neutral(i: int, faint: float = 0.0) -> str:
    bg = f"rgba(245, 158, 11, {faint:.2f})" if faint > 0.0 else "rgba(0,0,0,0.18)"
    return f"""
    <div title="position {i}" style="
        width: 8px; height: 8px;
        border-radius: 50%;
        background: {bg};
    "></div>
    """


# ---- score panel ----

def render_score(result: CopilotResult | None) -> str:
    """Big number, color band, class label.

    When `result` is None (idle / pre-click), shows a placeholder
    that occupies the same vertical space so the layout doesn't jump
    when the score arrives.
    """
    if result is None:
        return _score_placeholder()

    pct = int(round(result.score * 100))
    color = result.color
    label = result.predicted_label

    return f"""
    <div style="
        background: {_CARD_BG};
        border: 1px solid {_BORDER};
        border-radius: 12px;
        padding: 22px 26px;
        height: 100%;
        box-sizing: border-box;
    ">
        <div style="
            font-size: 11px;
            color: {_INK_DIM};
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            margin-bottom: 16px;
        ">
            Model verdict
        </div>
        <div style="
            font-size: 64px;
            font-weight: 700;
            color: {color};
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            line-height: 1;
            letter-spacing: -0.04em;
            margin-bottom: 12px;
        ">
            0.{pct:02d}
        </div>
        <div style="
            display: inline-block;
            padding: 5px 12px;
            background: {_hex_with_alpha(color, 0.12)};
            color: {color};
            font-size: 13px;
            font-weight: 600;
            border-radius: 999px;
            text-transform: uppercase;
            letter-spacing: 0.04em;
        ">
            {html.escape(label)}
        </div>
    </div>
    """


def _score_placeholder() -> str:
    return f"""
    <div style="
        background: {_CARD_BG};
        border: 1px solid {_BORDER};
        border-radius: 12px;
        padding: 22px 26px;
        height: 100%;
        box-sizing: border-box;
        color: {_INK_DIM};
    ">
        <div style="
            font-size: 11px;
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            margin-bottom: 16px;
        ">
            Model verdict
        </div>
        <div style="
            font-size: 64px;
            font-weight: 300;
            color: rgba(0,0,0,0.15);
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            line-height: 1;
            letter-spacing: -0.04em;
            margin-bottom: 12px;
        ">
            0.&mdash;
        </div>
        <div style="
            font-size: 13px;
            color: {_INK_DIM};
        ">
            Awaiting analysis &middot; click <b>Analyze</b>
        </div>
    </div>
    """


# ---- reasoning panel (streams in) ----

def render_reasoning(text: str | None) -> str:
    """Model's LM-decoded reasoning, with a blinking caret while streaming."""
    if text is None:
        body = f"""
        <div style="
            font-size: 14px;
            color: {_INK_DIM};
            font-style: italic;
        ">
            Reasoning will appear here once the model has read the customer's history.
        </div>
        """
    else:
        body = f"""
        <div style="
            font-size: 14px;
            color: {_INK};
            line-height: 1.6;
            white-space: pre-wrap;
        ">{html.escape(text)}<span style="
            display: inline-block;
            width: 7px;
            height: 14px;
            background: {_INK};
            margin-left: 2px;
            vertical-align: text-bottom;
            animation: copilot-blink 1s infinite;
        "></span></div>
        <style>
            @keyframes copilot-blink {{
                0%, 49% {{ opacity: 1; }}
                50%, 100% {{ opacity: 0; }}
            }}
        </style>
        """

    return f"""
    <div style="
        background: {_CARD_BG};
        border: 1px solid {_BORDER};
        border-radius: 12px;
        padding: 22px 26px;
        height: 100%;
        box-sizing: border-box;
    ">
        <div style="
            font-size: 11px;
            color: {_INK_DIM};
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            text-transform: uppercase;
            letter-spacing: 0.05em;
            margin-bottom: 16px;
        ">
            Reasoning
        </div>
        {body}
    </div>
    """


# ---- header ----

def render_header() -> str:
    return f"""
    <div style="text-align: center; margin-bottom: 22px;">
        <div style="
            font-size: 11px;
            color: {_INK_DIM};
            font-family: 'JetBrains Mono', ui-monospace, monospace;
            text-transform: uppercase;
            letter-spacing: 0.08em;
            margin-bottom: 6px;
        ">
            Liquid AI &middot; LFM2.5-350M backbone &middot; encoder + LoRA
        </div>
        <h1 style="
            margin: 0;
            font-size: 28px;
            font-weight: 700;
            color: {_INK};
            letter-spacing: -0.02em;
        ">
            Dispute Co-Pilot
        </h1>
    </div>
    """


# ---- helpers ----

def _short_label(name: str) -> str:
    return {
        "unlikely_friendly_fraud": "likely legit",
        "ambiguous": "ambiguous",
        "likely_friendly_fraud": "likely fraud",
    }.get(name, name)


def _hex_with_alpha(hex_color: str, alpha: float) -> str:
    """Convert #RRGGBB + alpha float into an rgba() string."""
    h = hex_color.lstrip("#")
    if len(h) != 6:
        return hex_color
    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
    return f"rgba({r}, {g}, {b}, {alpha:.3f})"