encoder/src/demo/copilot_app.py

"""Gradio app for the Dispute Co-Pilot demo.

Single-tab UI. Cast strip up top; click any of the six archetype
cards to select. Hit Analyze. The score + attribution glow paint
together; the reasoning text streams in word-by-word underneath.

CLI:
    python -m encoder.src.demo.copilot_app \\
        --checkpoint encoder/experiments/dispute_legitimacy_v1/step_003000_slim.pt \\
        --model-config encoder/configs/model_dispute_legitimacy.yaml \\
        --schema data/schema.yaml \\
        --histories data/synthetic/token_ids.npy \\
        --cast encoder/data/demo_cast.json \\
        --port 7860

This module owns the layout + event wiring only. All inference goes
through `copilot_inference.py`. All HTML goes through `copilot_render.py`.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import gradio as gr
import torch

from encoder.src.demo.copilot_inference import CastMember, CopilotModel
from encoder.src.demo.copilot_render import (
    render_cast_strip,
    render_complaint,
    render_header,
    render_reasoning,
    render_score,
    render_timeline,
)


# Container width: tuned for projector resolution (1280x800) and
# laptop demos. Anything wider and the cast strip wraps awkwardly.
_CONTAINER_WIDTH_PX = 1100


def _build_tab(model: CopilotModel) -> None:
    """Build the Dispute surface's components into the current Gradio context.

    Caller owns the surrounding Blocks/Tab. This function only adds the
    surface-specific content (cast strip, complaint quote, Analyze
    button, timeline, score panel, reasoning panel) and wires event
    handlers. The shared top-level header is owned by the unified app.
    """
    cast = model.cast

    # --- state ---
    selected_idx = gr.State(value=0)

    # --- cast strip ---
    cast_html = gr.HTML(render_cast_strip(cast, 0))

    with gr.Row():
        cast_buttons: list[gr.Button] = []
        for i, m in enumerate(cast):
            short = " ".join(m.display_name.split(" ")[:2])
            btn = gr.Button(
                value=short,
                variant="primary" if i == 0 else "secondary",
                scale=1,
            )
            cast_buttons.append(btn)

    # --- complaint quote ---
    complaint_html = gr.HTML(render_complaint(cast[0]))

    # --- analyze button ---
    with gr.Row():
        gr.HTML("<div style='flex: 1'></div>")
        analyze_btn = gr.Button(
            value="Analyze",
            variant="primary",
            size="lg",
            scale=0,
            min_width=180,
        )
        gr.HTML("<div style='flex: 1'></div>")

    # --- timeline ---
    timeline_html = gr.HTML(
        render_timeline(disputed_idx=cast[0].disputed_idx),
    )

    # --- score + reasoning side-by-side ---
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            score_html = gr.HTML(render_score(None))
        with gr.Column(scale=2):
            reasoning_html = gr.HTML(render_reasoning(None))

    # --- event wiring ---

    def _select(idx: int) -> tuple:
        member = cast[idx]
        button_updates = tuple(
            gr.update(variant="primary" if i == idx else "secondary")
            for i in range(len(cast))
        )
        return (
            idx,
            render_cast_strip(cast, idx),
            render_complaint(member),
            render_timeline(disputed_idx=member.disputed_idx),
            render_score(None),
            render_reasoning(None),
        ) + button_updates

    for i, btn in enumerate(cast_buttons):
        btn.click(
            fn=lambda i=i: _select(i),
            inputs=None,
            outputs=[selected_idx, cast_html, complaint_html,
                     timeline_html, score_html, reasoning_html,
                     *cast_buttons],
        )

    def _analyze(idx: int):
        member: CastMember = cast[idx]
        result = model.predict(member, top_k=5)
        timeline = render_timeline(
            disputed_idx=member.disputed_idx,
            top_k_positions=result.top_k_positions,
            attribution_probs=result.attribution_probs,
        )
        score = render_score(result)
        yield timeline, score, render_reasoning("")
        for partial in model.stream_reasoning(member, result, chunk_chars=6):
            yield timeline, score, render_reasoning(partial)

    analyze_btn.click(
        fn=_analyze,
        inputs=[selected_idx],
        outputs=[timeline_html, score_html, reasoning_html],
    )


def _build_ui(model: CopilotModel) -> gr.Blocks:
    """Standalone Blocks UI: Dispute-specific header + tab content."""
    with gr.Blocks(title="Dispute Co-Pilot — Liquid AI") as demo:
        gr.HTML(render_header())
        _build_tab(model)
    return demo


def main() -> None:
    parser = argparse.ArgumentParser(description="Dispute Co-Pilot Gradio demo")
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=Path("encoder/experiments/dispute_legitimacy_v7/demo_checkpoint.pt"),
    )
    parser.add_argument(
        "--model-config",
        type=Path,
        default=Path("encoder/configs/model_dispute_legitimacy.yaml"),
    )
    parser.add_argument("--schema", type=Path, default=Path("data/schema.yaml"))
    parser.add_argument(
        "--histories",
        type=Path,
        default=Path("data/synthetic/token_ids.npy"),
    )
    parser.add_argument(
        "--cast",
        type=Path,
        default=Path("encoder/data/demo_cast.json"),
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        choices=["cpu", "cuda", "mps"],
    )
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--share", action="store_true")
    args = parser.parse_args()

    device = torch.device(args.device)
    print(f"Loading CopilotModel on {device} ...")
    model = CopilotModel.from_paths(
        checkpoint_path=args.checkpoint,
        model_config_path=args.model_config,
        schema_path=args.schema,
        histories_path=args.histories,
        cast_path=args.cast,
        device=device,
    )
    print(f"  cast size: {len(model.cast)}")
    print(f"  histories: {model.histories.shape} dtype={model.histories.dtype}")

    demo = _build_ui(model)
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=args.share,
        theme=gr.themes.Default(
            font=["Inter", "system-ui", "sans-serif"],
            font_mono=["JetBrains Mono", "ui-monospace", "monospace"],
        ),
        css=f"""
        .gradio-container {{
            max-width: {_CONTAINER_WIDTH_PX}px !important;
            background: #fafafa !important;
        }}
        """,
    )


if __name__ == "__main__":
    main()