_html.py

from __future__ import annotations

# index 0-5 = teachers (aligns with MODEL_COLORS[1:] in _fig.py)
_TEACHER_COLORS = ["#7c3aed", "#06b6d4", "#f59e0b", "#34d399", "#f472b6", "#76b900"]


def gate_html(gate_weights: list[float], teacher_names: list[str],
              ranked: bool = True) -> str:
    """Horizontal bar chart. ranked=False keeps teacher order fixed — used
    while streaming so bars animate in place instead of swapping rows."""
    pairs = list(enumerate(zip(gate_weights, teacher_names)))
    if ranked:
        pairs.sort(key=lambda x: x[1][0], reverse=True)
    rows = []
    for orig_idx, (w, name) in pairs:
        color = _TEACHER_COLORS[orig_idx % len(_TEACHER_COLORS)]
        bar_pct = int(w * 100)
        rows.append(
            f'<div style="display:flex;align-items:center;gap:8px;margin:5px 0;">'
            f'<span style="width:80px;font-family:monospace;font-size:11px;color:{color};'
            f'white-space:nowrap;overflow:hidden;text-overflow:ellipsis;">{name}</span>'
            f'<div style="flex:1;height:5px;background:#21262d;border-radius:3px;overflow:hidden;">'
            f'<div style="width:{bar_pct}%;height:100%;background:{color};border-radius:3px;'
            f'transition:width 0.18s ease;"></div>'
            f'</div>'
            f'<span style="width:36px;font-family:monospace;font-size:11px;'
            f'color:#8b949e;text-align:right;">{w:.2f}</span>'
            f'</div>'
        )
    return (
        '<div style="background:#161b22;border:1px solid #30363d;border-radius:6px;padding:14px;margin-top:8px;">'
        '<div style="font-size:10px;color:#8b949e;font-family:monospace;margin-bottom:10px;letter-spacing:0.04em;">'
        'TEACHER INFLUENCE</div>'
        + "".join(rows)
        + "</div>"
    )


def task_html(gate_weights: list[float], teacher_names: list[str]) -> str:
    """Badge for argmax teacher + runner-up."""
    best_idx = max(range(len(gate_weights)), key=lambda i: gate_weights[i])
    second_idx = sorted(range(len(gate_weights)), key=lambda i: gate_weights[i])[-2]
    color = _TEACHER_COLORS[best_idx % len(_TEACHER_COLORS)]
    best = teacher_names[best_idx].upper()
    second = teacher_names[second_idx].upper()
    w2 = gate_weights[second_idx]
    n = len(teacher_names)
    return (
        '<div style="background:#161b22;border:1px solid #30363d;border-radius:6px;padding:14px;margin-top:8px;">'
        '<div style="font-size:10px;color:#8b949e;font-family:monospace;margin-bottom:10px;letter-spacing:0.04em;">'
        'DOMINANT TEACHER</div>'
        '<div style="display:flex;align-items:center;gap:12px;">'
        f'<span style="display:inline-flex;align-items:center;gap:6px;background:{color}22;'
        f'border:1px solid {color}55;border-radius:999px;padding:4px 14px;'
        f'font-family:monospace;font-weight:600;color:{color};font-size:13px;">'
        f'<span style="width:6px;height:6px;background:{color};border-radius:50%;'
        f'box-shadow:0 0 6px {color};"></span>{best}</span>'
        f'<span style="font-size:11px;color:#8b949e;font-family:monospace;">'
        f'runner-up <b style="color:#e6edf3;">{second}</b> · {w2:.2f}</span>'
        f'</div>'
        f'<div style="font-size:10px;color:#6e7681;margin-top:8px;font-family:monospace;">'
        f'gate picked this teacher for your prompt ({n} total)</div>'
        '</div>'
    )


def header_html(n_teachers: int = 6, backend: str = "torch") -> str:
    """App header with ∀ logo, wordmark, and status pills."""
    def pill(color: str, dot_color: str, text: str, pulse: bool = False) -> str:
        dot_class = ' class="live-dot"' if pulse else ''
        return (
            f'<span style="display:inline-flex;align-items:center;gap:8px;height:28px;'
            f'padding:0 12px;border-radius:999px;font-family:monospace;font-size:10px;'
            f'font-weight:600;letter-spacing:0.06em;color:{color};'
            f'background:rgba(13,17,23,0.8);border:1px solid {color}33;'
            f'backdrop-filter:blur(8px);white-space:nowrap;">'
            f'<span{dot_class} style="width:6px;height:6px;border-radius:50%;'
            f'background:{dot_color};box-shadow:0 0 7px {dot_color};flex-shrink:0;"></span>'
            f'{text}</span>'
        )
    return (
        '<div style="display:flex;align-items:center;justify-content:space-between;'
        'padding:20px 2px 14px;gap:16px;flex-wrap:wrap;">'

        # Left: logo + title
        '<div style="display:flex;align-items:center;gap:14px;">'
        '<div style="width:46px;height:46px;display:grid;place-items:center;flex-shrink:0;'
        'border-radius:10px;'
        'background:radial-gradient(circle at 40% 35%, rgba(124,58,237,0.45) 0%, rgba(8,11,16,0.9) 70%);'
        'border:1px solid rgba(124,58,237,0.35);'
        'box-shadow:0 0 20px rgba(124,58,237,0.2),inset 0 1px 0 rgba(255,255,255,0.07);">'
        '<span style="font-family:monospace;font-weight:800;font-size:26px;'
        'background:linear-gradient(135deg,#a78bfa,#7c3aed);'
        '-webkit-background-clip:text;-webkit-text-fill-color:transparent;'
        'background-clip:text;">∀</span>'
        '</div>'
        '<div>'
        '<div style="font-weight:700;font-size:18px;letter-spacing:-0.01em;'
        'background:linear-gradient(90deg,#e6edf3 0%,#a78bfa 100%);'
        '-webkit-background-clip:text;-webkit-text-fill-color:transparent;'
        'background-clip:text;">One for All</div>'
        '<div style="font-size:11px;color:#484f58;margin-top:2px;font-family:monospace;'
        'letter-spacing:0.05em;text-transform:uppercase;">multi-teacher soul distillation</div>'
        '</div></div>'

        # Right: pills
        '<div style="display:flex;gap:8px;flex-wrap:wrap;">'
        + pill("#06b6d4", "#06b6d4", "STUDENT · Qwen2.5-0.5B", pulse=True)
        + pill("#7c3aed", "#7c3aed", f"{n_teachers} TEACHERS")
        + pill("#f59e0b", "#f59e0b", "PATH B · GEOMETRY")
        + (pill("#76b900", "#76b900", "🦙 llama.cpp · GGUF") if backend == "llamacpp" else "")
        + '</div></div>'

        # Divider with glow
        '<div style="height:1px;margin-bottom:22px;'
        'background:linear-gradient(90deg,transparent 0%,rgba(124,58,237,0.6) 40%,'
        'rgba(6,182,212,0.4) 70%,transparent 100%);'
        'box-shadow:0 0 12px rgba(124,58,237,0.15);"></div>'
    )