encoder/src/demo/copilot_render_collections.py

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

Mirror of `copilot_render.py` (dispute) adapted to the K-treatment
output. Same visual vocabulary as dispute:
  - Inter for body, JetBrains Mono for numbers.
  - Subtle palette: near-white background, near-black ink, small set
    of brand accents per treatment.
  - Inline styles only; self-contained snippets.

The unique panel for Collections is the treatment-grid score panel:
four horizontal bars showing P(likely_respond) per treatment, with
the dominant treatment highlighted with a recommendation pill.
"""

from __future__ import annotations

import html
from typing import Iterable

import numpy as np

from encoder.src.data.synthetic_collections import (
    BAND_NAMES,
    NUM_TREATMENTS,
    TREATMENT_NAMES,
)
from encoder.src.demo.copilot_inference_collections import (
    TREATMENT_COLORS,
    CollectionsCastMember,
    CollectionsResult,
)


# ---- shared constants ----

_PAGE_BG = "#fafafa"
_INK = "#171717"
_INK_DIM = "#525252"
_BORDER = "rgba(0,0,0,0.08)"
_CARD_BG = "#ffffff"
_BAND_NAME_SHORT = {
    "unlikely_respond": "unlikely",
    "ambiguous": "ambiguous",
    "likely_respond": "LIKELY",
}


# ---- 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;
        ">
            Collections Co-Pilot
        </h1>
    </div>
    """


# ---- cast strip ----

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

    Each card shows the cast member's display name and a small pill
    indicating the expected dominant treatment.
    """
    cards: list[str] = []
    for i, m in enumerate(cast):
        is_sel = i == selected_idx
        accent = TREATMENT_COLORS[m.dominant_treatment]
        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(m.dominant_treatment_name)}
            </div>
        </div>
        """)
    return f"""
    <div style="display: flex; gap: 10px; margin-bottom: 18px;">
        {''.join(cards)}
    </div>
    """


# ---- context text ----

def render_context(member: CollectionsCastMember) -> str:
    """Pull-quote of the analyst-facing delinquency context."""
    return f"""
    <div style="
        background: {_CARD_BG};
        border: 1px solid {_BORDER};
        border-left: 4px solid {TREATMENT_COLORS[member.dominant_treatment]};
        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;
        ">
            Delinquency context &middot; pattern:
            <span style="color: {_INK}; font-weight: 600;">{html.escape(member.pattern)}</span>
        </div>
        <div style="
            font-size: 18px;
            color: {_INK};
            line-height: 1.5;
            font-style: italic;
            font-weight: 400;
        ">
            &ldquo;{html.escape(member.context_text)}&rdquo;
        </div>
        <div style="
            font-size: 13px;
            color: {_INK_DIM};
            margin-top: 12px;
            line-height: 1.5;
        ">
            {html.escape(member.description)}
        </div>
    </div>
    """


# ---- timeline ----

def render_timeline(
    context_idx: int,
    top_k_positions: Iterable[int] | None = None,
    attribution_probs: Iterable[float] | None = None,
    num_positions: int = 64,
) -> str:
    """64 dots; the context position is starred, top-k glow.

    Same dot vocabulary as the dispute timeline — re-using the visual
    grammar so audiences cross-trained on Dispute pick up Collections
    in one frame.
    """
    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 == context_idx:
            dots.append(_dot_context(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 &middot; oldest → newest</span>
            <span>
                <span style="color: #2563eb;">★</span> now &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_context(i: int) -> str:
    return f"""
    <div title="position {i} &middot; current" style="
        width: 14px; height: 14px;
        color: #2563eb;
        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} &middot; 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>
    """


# ---- treatment grid (the unique Collections panel) ----

def render_treatment_grid(result: CollectionsResult | None) -> str:
    """K horizontal bars showing P(likely_respond) per treatment.

    The dominant treatment is highlighted with a pill. When `result`
    is None (pre-analyze), renders a placeholder so the layout doesn't
    shift when the verdict arrives.
    """
    if result is None:
        return _treatment_grid_placeholder()

    rows: list[str] = []
    for t in range(NUM_TREATMENTS):
        score = float(result.likely_scores[t])
        band = result.predicted_bands[t]
        accent = TREATMENT_COLORS[t]
        is_dominant = t == result.dominant_treatment
        bar_w = max(2, int(score * 100))
        rows.append(f"""
        <div style="
            display: flex;
            align-items: center;
            gap: 12px;
            padding: 8px 0;
            border-bottom: 1px solid {_BORDER};
        ">
            <div style="
                width: 130px;
                font-size: 13px;
                font-weight: {600 if is_dominant else 500};
                color: {_INK if is_dominant else _INK_DIM};
                font-family: 'JetBrains Mono', ui-monospace, monospace;
            ">
                {'★ ' if is_dominant else '&nbsp;&nbsp; '}{html.escape(TREATMENT_NAMES[t])}
            </div>
            <div style="
                flex: 1;
                height: 14px;
                background: rgba(0,0,0,0.05);
                border-radius: 7px;
                position: relative;
                overflow: hidden;
            ">
                <div style="
                    position: absolute;
                    left: 0; top: 0; bottom: 0;
                    width: {bar_w}%;
                    background: {accent};
                    border-radius: 7px;
                    transition: width 0.3s ease;
                "></div>
            </div>
            <div style="
                width: 56px;
                text-align: right;
                font-size: 14px;
                font-weight: 600;
                color: {accent};
                font-family: 'JetBrains Mono', ui-monospace, monospace;
            ">
                {score:.2f}
            </div>
            <div style="
                width: 76px;
                text-align: right;
                font-size: 11px;
                color: {_INK_DIM};
                text-transform: uppercase;
                letter-spacing: 0.04em;
            ">
                {html.escape(_BAND_NAME_SHORT.get(BAND_NAMES[band], BAND_NAMES[band]))}
            </div>
        </div>
        """)

    dom_idx = result.dominant_treatment
    dom_color = TREATMENT_COLORS[dom_idx]
    dom_score = float(result.likely_scores[dom_idx])
    dom_name = TREATMENT_NAMES[dom_idx]

    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="
            display: flex;
            justify-content: space-between;
            align-items: baseline;
            margin-bottom: 16px;
        ">
            <div style="
                font-size: 11px;
                color: {_INK_DIM};
                font-family: 'JetBrains Mono', ui-monospace, monospace;
                text-transform: uppercase;
                letter-spacing: 0.05em;
            ">
                Treatment scoreboard
            </div>
            <div style="
                display: inline-block;
                padding: 4px 12px;
                background: {_hex_with_alpha(dom_color, 0.12)};
                color: {dom_color};
                font-size: 12px;
                font-weight: 700;
                border-radius: 999px;
                text-transform: uppercase;
                letter-spacing: 0.04em;
            ">
                recommend: {html.escape(dom_name)} &middot; {dom_score:.2f}
            </div>
        </div>
        {''.join(rows)}
    </div>
    """


def _treatment_grid_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;
        ">
            Treatment scoreboard
        </div>
        <div style="font-size: 14px; color: {_INK_DIM};">
            Click <b>Analyze</b> to see per-treatment response probabilities.
        </div>
    </div>
    """


# ---- reasoning ----

def render_reasoning(text: str | None) -> str:
    """Reasoning panel; same vocabulary as dispute."""
    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;
            font-family: ui-monospace, 'JetBrains Mono', monospace;
        ">{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>
    """


# ---- helpers ----

def _hex_with_alpha(hex_color: str, alpha: float) -> str:
    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})"