"""Interactive Gradio demo for the encoder-on-LFM2.5 transaction model.

Three-tab structure (Demo / Why Liquid / Integration). Demo is the
landing tab so a customer immediately sees what the model does; Why
Liquid is the architectural pitch; Integration is the build-it-yourself
playbook. Same Gradio theme + CSS as the rest of Liquid's customer-
facing demos for visual consistency.

The demo is intentionally self-contained: no side-by-side comparison
against an alternative architecture, no per-tab references to other
work. The argument is the architecture pattern itself — encoder +
frozen LFM2.5 backbone + LoRA + multi-head — and the reader gets a
clean read of it.

Usage:
    python -m encoder.src.demo.app \\
        --checkpoint encoder/experiments/.../step_004999.pt \\
        --model-config encoder/configs/model_nocompress.yaml \\
        --schema data/schema.yaml \\
        --data-dir data/synthetic \\
        --port 7860
"""

from __future__ import annotations

import argparse
import time
from pathlib import Path

import gradio as gr
import torch

from src.data.schema import load_schema
from src.demo.app import DemoData  # reuse parent's curated-test-set loader
from src.demo.decode import TransactionDecoder
from src.demo.merchant_catalog import DemoMerchantCatalog
from src.demo.profile_inference import format_profile_html, infer_profile

from encoder.src.demo.cold_start import render_cold_start
from encoder.src.demo.inference import EncoderDemoModel
from encoder.src.demo.render import (
    format_amount_predictions,
    format_fraud_score,
    format_mcc_predictions,
    format_merchant_predictions,
    format_timeline,
    render_encoder_integration,
    render_why_encoder,
)


# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------

_HEADER_HTML = """
<div style="text-align: center; margin-bottom: 16px;">
    <h1 style="margin: 0; font-size: 26px; font-weight: 700; color: #171717; letter-spacing: -0.02em;">
        Encoder on LFM2.5 — Transaction Foundation Model
    </h1>
    <p style="color: #737373; margin: 6px 0 0 0; font-size: 13px;
              font-family: JetBrains Mono, ui-monospace, monospace;">
        Liquid AI &middot; LFM2.5-350M base &middot; Encoder + LoRA + multi-head
    </p>
</div>
"""


# Container width applied to every tab's content so the three surfaces
# read at the same width. Kept in sync with render.py's _CONTAINER_WIDTH.
_CONTAINER_WIDTH = "1180px"


# ---------------------------------------------------------------------------
# Demo tab intro — standalone, no comparison framing
# ---------------------------------------------------------------------------

_DEMO_INTRO_HTML = f"""
<div style="max-width: {_CONTAINER_WIDTH}; margin: 16px auto 8px auto; padding: 14px 18px;
            background: #ffffff; border: 1px solid rgba(0,0,0,0.1);
            border-radius: 12px;">
    <div style="font-size: 14px; font-weight: 600; color: #171717; margin-bottom: 6px;">
        Live multi-head inference on a frozen LFM2.5-350M backbone
    </div>
    <div style="font-size: 12px; color: #525252; line-height: 1.55;">
        Pick a curated customer archetype, click <b>Run Inference</b>, watch four task
        heads predict in parallel from one shared backbone — fraud probability,
        next-merchant, amount bucket, and merchant category (MCC). The encoder turns
        the 64-transaction history into 960 pseudo-tokens; the frozen LFM2.5-350M
        backbone with LoRA processes them; per-task heads pool the hidden states.
        For the architectural rationale see <i>Why Liquid</i>; for the build-it-yourself
        playbook see <i>Integration</i>.
    </div>
</div>
"""


# ---------------------------------------------------------------------------
# CSS — scoped to the encoder demo
# ---------------------------------------------------------------------------

_CSS = f"""
/* Force light mode regardless of system preference */
:root, .dark {{ color-scheme: light !important; }}

.gradio-container {{
    background: #f5f5f5 !important;
    max-width: 1280px !important;
    margin: auto !important;
    padding: 1.5rem !important;
}}

/* Tabs: pill-style matching the rest of Liquid's customer-facing demos */
.tabs {{ background: transparent !important; }}
.tab-nav {{
    background: #f5f5f5 !important;
    border: none !important;
    border-bottom: 1px solid rgba(0,0,0,0.1) !important;
    gap: 4px !important;
    padding: 4px 0 !important;
}}
.tab-nav button {{
    font-weight: 500 !important;
    font-size: 14px !important;
    color: #737373 !important;
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    padding: 8px 16px !important;
    border-radius: 0 !important;
    transition: all 0.15s ease !important;
}}
.tab-nav button:hover {{
    color: #171717 !important;
    background: rgba(0,0,0,0.03) !important;
}}
.tab-nav button.selected {{
    color: #171717 !important;
    font-weight: 600 !important;
    border-bottom: 2px solid #171717 !important;
    background: transparent !important;
}}

/* Run button: pill style, dark fill. Scoped by elem_id so it does not
   bleed into Gradio's internal radio/checkbox <button> wrappers. */
#run-inference-btn button,
button#run-inference-btn {{
    background: #171717 !important;
    color: #ffffff !important;
    border: none !important;
    border-radius: 9999px !important;
    padding: 10px 24px !important;
    font-weight: 600 !important;
    letter-spacing: -0.01em !important;
}}
#run-inference-btn button:hover,
button#run-inference-btn:hover {{
    background: #404040 !important;
}}

/* Wrap the Demo tab's inner content at the same width as the HTML tabs
   so all three tabs read at consistent width. */
#demo-tab-container {{
    max-width: {_CONTAINER_WIDTH} !important;
    margin: 0 auto !important;
    padding: 0 16px !important;
}}
"""


# Gradio theme matching parent demo's design system (monochrome neutral
# with Inter + JetBrains Mono fonts). Kept in this file rather than a
# shared module to keep encoder/ a self-contained directory.
def _build_theme() -> gr.themes.Soft:
    return gr.themes.Soft(
        primary_hue="neutral",
        secondary_hue="neutral",
        neutral_hue="neutral",
        font=gr.themes.GoogleFont("Inter"),
        font_mono=gr.themes.GoogleFont("JetBrains Mono"),
    ).set(
        body_background_fill="#f5f5f5",
        body_text_color="#171717",
        body_text_color_subdued="#737373",
        block_background_fill="#ffffff",
        block_border_color="rgba(0,0,0,0.1)",
        block_label_background_fill="#f5f5f5",
        block_label_text_color="#525252",
        block_title_text_color="#171717",
        block_shadow="0 1px 3px rgba(0,0,0,0.04)",
        input_background_fill="#ffffff",
        input_border_color="rgba(0,0,0,0.1)",
        input_border_color_focus="#171717",
        input_placeholder_color="#a3a3a3",
        panel_background_fill="#fafafa",
        panel_border_color="rgba(0,0,0,0.06)",
        border_color_primary="rgba(0,0,0,0.1)",
        button_primary_background_fill="#171717",
        button_primary_background_fill_hover="#404040",
        button_primary_text_color="#ffffff",
        button_secondary_background_fill="#ffffff",
        button_secondary_text_color="#525252",
        button_secondary_border_color="rgba(0,0,0,0.1)",
        slider_color="#171717",
        table_border_color="rgba(0,0,0,0.06)",
        table_even_background_fill="#fafafa",
        table_odd_background_fill="#ffffff",
        shadow_spread="0px",
        color_accent_soft="rgba(0,0,0,0.04)",
    )


# ---------------------------------------------------------------------------
# App builder
# ---------------------------------------------------------------------------


def _build_demo_tab_contents(
    model: EncoderDemoModel,
    data: DemoData,
    decoder: TransactionDecoder,
    merchant_catalog: DemoMerchantCatalog,
    app: gr.Blocks,
) -> None:
    """Build the original Multi-Head Demo tab into the current Gradio context.

    Caller owns the outer Blocks/Tab. The `app` argument is the host
    Blocks instance, needed so we can register the auto-run `app.load(...)`
    hook against the right Blocks.
    """

    def on_customer_select(curated_name: str) -> tuple[str, str, str, str, str, str, str, str]:
        if not curated_name or curated_name not in data.curated_names:
            return ("Select a customer to see predictions.", "", "", "", "", "", "", "")

        idx = data.get_curated_index(curated_name)
        token_ids = data.token_ids[idx]
        is_fraud = bool(data.labels[idx])

        summary = decoder.summarize_customer(token_ids, is_fraud)

        t0 = time.perf_counter()
        preds = model.run_inference(token_ids)
        latency_ms = (time.perf_counter() - t0) * 1000

        timeline_html = format_timeline(decoder, token_ids)
        fraud_html = format_fraud_score(float(preds["fraud"][0]))
        merchant_html = format_merchant_predictions(
            preds["next_merchant"], merchant_catalog, k=5,
        )
        amount_html = format_amount_predictions(preds["amount_range"], k=5)
        mcc_html = format_mcc_predictions(preds["mcc"], k=5)

        profile_match = infer_profile(token_ids)
        profile_html = format_profile_html(profile_match)

        latency_html = (
            f"<div style='font-family: JetBrains Mono, ui-monospace, monospace; "
            f"font-size: 11px; color: #737373; padding: 6px 10px; "
            f"background: #fafafa; border-radius: 6px; display: inline-block;'>"
            f"Inference: <b style='color: #171717;'>{latency_ms:.0f} ms</b> "
            f"({'CPU' if model.device.type == 'cpu' else 'GPU'}) "
            f"&middot; ground truth: "
            f"<b style='color: {'#EF4444' if is_fraud else '#10B981'};'>"
            f"{'FRAUD' if is_fraud else 'LEGITIMATE'}</b>"
            f"</div>"
        )

        return (
            summary,
            timeline_html,
            profile_html,
            fraud_html,
            merchant_html,
            amount_html,
            mcc_html,
            latency_html,
        )

    gr.HTML(_DEMO_INTRO_HTML)

    with gr.Column(elem_id="demo-tab-container"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML(
                    "<h3 style='margin: 0 0 8px 0; color: #171717;'>"
                    "Select Customer</h3>"
                )
                curated_dropdown = gr.Dropdown(
                    choices=data.curated_names,
                    value=data.curated_names[0] if data.curated_names else None,
                    label="Curated archetypes",
                    info="5 legitimate profiles + 3 fraud archetypes",
                )
                run_btn = gr.Button(
                    "Run Inference",
                    variant="primary",
                    size="lg",
                    elem_id="run-inference-btn",
                )
                latency_html = gr.HTML("")
            with gr.Column(scale=2):
                summary_html = gr.HTML("")

        gr.HTML(
            "<h3 style='margin: 20px 0 8px 0; color: #171717;'>"
            "Transaction Timeline</h3>"
        )
        timeline_html = gr.HTML("")

        with gr.Row():
            with gr.Column():
                gr.HTML(
                    "<h3 style='margin: 16px 0 8px 0; color: #171717;'>"
                    "Behavioral Profile</h3>"
                )
                profile_html = gr.HTML("")
            with gr.Column():
                gr.HTML(
                    "<h3 style='margin: 16px 0 8px 0; color: #171717;'>"
                    "Fraud Score</h3>"
                )
                fraud_html = gr.HTML("")

        gr.HTML(
            "<h3 style='margin: 20px 0 8px 0; color: #171717;'>"
            "Next-Transaction Predictions</h3>"
        )
        with gr.Row():
            with gr.Column():
                gr.HTML(
                    "<h4 style='margin: 4px 0; color: #525252; font-size: 13px;'>"
                    "Next merchant</h4>"
                )
                merchant_html = gr.HTML("")
            with gr.Column():
                gr.HTML(
                    "<h4 style='margin: 4px 0; color: #525252; font-size: 13px;'>"
                    "Amount bucket</h4>"
                )
                amount_html = gr.HTML("")
            with gr.Column():
                gr.HTML(
                    "<h4 style='margin: 4px 0; color: #525252; font-size: 13px;'>"
                    "Merchant category (MCC)</h4>"
                )
                mcc_html = gr.HTML("")

    outputs = [
        summary_html,
        timeline_html,
        profile_html,
        fraud_html,
        merchant_html,
        amount_html,
        mcc_html,
        latency_html,
    ]

    run_btn.click(
        fn=on_customer_select,
        inputs=[curated_dropdown],
        outputs=outputs,
    )

    # Auto-run on first load so the user lands on populated outputs.
    app.load(
        fn=on_customer_select,
        inputs=[curated_dropdown],
        outputs=outputs,
    )


def _build_why_liquid_tab_contents() -> None:
    """Render the Why Liquid tab content (HTML pitch)."""
    gr.HTML(render_why_encoder())


def _build_integration_tab_contents() -> None:
    """Render the Integration tab content (HTML build-it-yourself guide)."""
    gr.HTML(render_encoder_integration())


def _build_cold_start_tab_contents() -> None:
    """Render the Cold Start tab content (HTML self-supervised-pretraining result)."""
    gr.HTML(render_cold_start())


def create_app(
    model: EncoderDemoModel,
    data: DemoData,
    decoder: TransactionDecoder,
    merchant_catalog: DemoMerchantCatalog,
) -> gr.Blocks:
    """Standalone 3-tab Gradio app for the multi-head encoder demo."""
    with gr.Blocks(
        title="Encoder on LFM2.5 — Transaction Foundation Model",
        css=_CSS,
        theme=_build_theme(),
    ) as app:
        gr.HTML(_HEADER_HTML)
        with gr.Tabs():
            with gr.Tab("Demo"):
                _build_demo_tab_contents(model, data, decoder, merchant_catalog, app)
            with gr.Tab("Cold Start"):
                _build_cold_start_tab_contents()
            with gr.Tab("Why Liquid"):
                _build_why_liquid_tab_contents()
            with gr.Tab("Integration"):
                _build_integration_tab_contents()
    return app


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Encoder + LFM2.5-350M interactive demo",
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=None,
        help="Trained checkpoint path. Omit to run with random-init (debug only).",
    )
    parser.add_argument(
        "--model-config",
        type=Path,
        default=Path("encoder/configs/model_nocompress.yaml"),
    )
    parser.add_argument(
        "--schema",
        type=Path,
        default=Path("data/schema.yaml"),
    )
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path("data/synthetic"),
        help="Test-set data directory (token_ids.npy, sequence_labels.npy, etc.)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        choices=["cpu", "cuda", "mps"],
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="float32",
        choices=["float32", "bfloat16"],
    )
    parser.add_argument(
        "--port",
        type=int,
        default=7860,
    )
    args = parser.parse_args()

    dtype = torch.float32 if args.dtype == "float32" else torch.bfloat16

    print("Loading schema, data, model...")
    schema = load_schema(args.schema)
    data = DemoData(args.data_dir, schema)
    merchant_catalog = DemoMerchantCatalog(schema)
    decoder = TransactionDecoder(schema, merchant_catalog)

    model = EncoderDemoModel(
        model_config_path=args.model_config,
        schema_path=args.schema,
        checkpoint_path=args.checkpoint,
        dtype=dtype,
        device=args.device,
    )
    print(f"Model loaded: {model.checkpoint_status}")
    pc = model.num_params()
    print(f"Params — total: {pc['total']:,} / trainable: {pc['trainable']:,}")

    app = create_app(model, data, decoder, merchant_catalog)

    # Walk a few ports if 7860 is in use.
    for port in range(args.port, args.port + 10):
        try:
            app.launch(server_port=port, server_name="0.0.0.0")
            break
        except OSError as e:
            if "Cannot find empty port" in str(e) or "Address already in use" in str(e):
                print(f"  port {port} in use, trying {port + 1}...")
                continue
            raise


if __name__ == "__main__":
    main()