"""Epicure Explorer: a chef-facing interactive demo of the three sibling embeddings.

Three tabs:
  - Pairings: top-K cosine neighbours + closest emergent mode for a chosen ingredient.
  - Supervised SLERP: rotate a seed toward a supervised pole (cuisine, food group,
    NOVA, sensory, USDA macros) by a chosen angle.
  - Emergent SLERP: rotate a seed toward an emergent factor-mode pole.

Loads all three siblings on startup from their HF model repos.
"""

from __future__ import annotations

import os
import sys
import gradio as gr

# epicure.py is loaded from the cooc repo's snapshot at runtime; alternatively
# copy it into this Space's root for offline development.
try:
    from epicure import Epicure  # noqa: F401
except ImportError:
    from huggingface_hub import hf_hub_download
    epicure_py = hf_hub_download("Kaikaku/epicure-cooc", "epicure.py")
    sys.path.insert(0, os.path.dirname(epicure_py))
    from epicure import Epicure  # noqa: F401

MODELS = {
    "cooc": Epicure.from_pretrained("Kaikaku/epicure-cooc"),
    "core": Epicure.from_pretrained("Kaikaku/epicure-core"),
    "chem": Epicure.from_pretrained("Kaikaku/epicure-chem"),
}

ALL_INGREDIENTS = sorted(MODELS["cooc"].vocab.keys())


def _supervised_choices(sibling: str) -> list[str]:
    return sorted(MODELS[sibling].supervised_poles.keys())


def _factor_modes(sibling: str) -> list[tuple[str, str]]:
    return [
        (f"{m.mode_id} - {m.label}", m.mode_id)
        for m in MODELS[sibling].modes
        if m.kind == "factor"
    ]


def pairings(sibling: str, ingredient: str, k: int):
    if not ingredient or ingredient not in MODELS[sibling].vocab:
        return [], []
    m = MODELS[sibling]
    nb = m.neighbors(ingredient, k=k)
    cm = m.closest_mode(ingredient, kind=None, k=k)
    return (
        [[name, f"{sim:.4f}"] for name, sim in nb],
        [[mid, label, f"{sim:.4f}"] for mid, label, sim in cm],
    )


def supervised_slerp(sibling: str, seed: str, direction: str, theta: float, k: int):
    if not seed or seed not in MODELS[sibling].vocab:
        return []
    if direction not in MODELS[sibling].supervised_poles:
        return []
    r = MODELS[sibling].slerp(seed, direction, theta_deg=theta, k=k)
    return [[name, f"{sim:.4f}"] for name, sim in r]


def emergent_slerp(sibling: str, seed: str, factor_mode_id: str, theta: float, k: int):
    if not seed or seed not in MODELS[sibling].vocab:
        return []
    m = MODELS[sibling]
    pole = None
    for mode in m.modes:
        if mode.mode_id == factor_mode_id:
            pole = mode.pole
            break
    if pole is None:
        return []
    r = m.slerp(seed, pole, theta_deg=theta, k=k)
    return [[name, f"{sim:.4f}"] for name, sim in r]


with gr.Blocks(title="Epicure Explorer") as demo:
    gr.Markdown(
        """# Epicure Explorer

Interactive chef-facing operators over the three Epicure sibling embeddings (Cooc, Core, Chem),
from the paper [Epicure: Navigating the Emergent Geometry of Food Ingredient Embeddings](https://arxiv.org/abs/2605.22391).

Each sibling sits at a different point on the chemistry-vs-recipe-context spectrum:
- **Cooc** walks recipe co-occurrence only. Neighbours are recipe companions.
- **Core** blends typed FlavorDB compound walks with injected ingredient-ingredient walks. Concentrated geometry, tightest modes.
- **Chem** walks typed FlavorDB compound metapaths only. Strongest supervised-direction recovery; neighbours are flavour-profile peers.
"""
    )

    sibling = gr.Radio(
        choices=["cooc", "core", "chem"],
        value="chem",
        label="Sibling embedding",
    )

    with gr.Tab("Pairings"):
        ingredient = gr.Dropdown(
            choices=ALL_INGREDIENTS, value="chicken", label="Ingredient", allow_custom_value=False
        )
        k_pair = gr.Slider(1, 10, value=5, step=1, label="K")
        pair_btn = gr.Button("Find pairings", variant="primary")
        with gr.Row():
            nb_table = gr.Dataframe(
                headers=["Neighbour", "Cosine"], label="Top-K nearest neighbours", interactive=False
            )
            mode_table = gr.Dataframe(
                headers=["Mode id", "Label", "Cosine"], label="Closest modes", interactive=False
            )
        pair_btn.click(
            pairings, inputs=[sibling, ingredient, k_pair], outputs=[nb_table, mode_table]
        )

    with gr.Tab("Supervised SLERP"):
        sup_seed = gr.Dropdown(
            choices=ALL_INGREDIENTS, value="rice", label="Seed ingredient"
        )
        sup_dir = gr.Dropdown(
            choices=_supervised_choices("chem"),
            value="cuisine:South_Asian",
            label="Supervised direction",
        )
        sup_theta = gr.Slider(0, 90, value=30, step=5, label="Rotation angle (deg)")
        sup_k = gr.Slider(1, 10, value=5, step=1, label="K")
        sup_btn = gr.Button("Rotate", variant="primary")
        sup_table = gr.Dataframe(
            headers=["Ingredient", "Cosine"], label="Top-K rotated-query neighbours"
        )
        sup_btn.click(
            supervised_slerp,
            inputs=[sibling, sup_seed, sup_dir, sup_theta, sup_k],
            outputs=sup_table,
        )
        sibling.change(
            lambda s: gr.Dropdown(choices=_supervised_choices(s), value=None),
            inputs=sibling,
            outputs=sup_dir,
        )

    with gr.Tab("Emergent SLERP"):
        em_seed = gr.Dropdown(
            choices=ALL_INGREDIENTS, value="chocolate", label="Seed ingredient"
        )
        factor_options = _factor_modes("chem")
        em_mode = gr.Dropdown(
            choices=[label for label, _ in factor_options],
            value=factor_options[0][0] if factor_options else None,
            label="Emergent factor mode (label - mode_id)",
        )
        em_theta = gr.Slider(0, 90, value=30, step=5, label="Rotation angle (deg)")
        em_k = gr.Slider(1, 10, value=5, step=1, label="K")
        em_btn = gr.Button("Rotate", variant="primary")
        em_table = gr.Dataframe(
            headers=["Ingredient", "Cosine"], label="Top-K rotated-query neighbours"
        )

        def _resolve_factor(sib, label, seed, theta, k):
            options = _factor_modes(sib)
            mode_id = None
            for lab, mid in options:
                if lab == label:
                    mode_id = mid
                    break
            if mode_id is None and options:
                mode_id = options[0][1]
            if mode_id is None:
                return []
            return emergent_slerp(sib, seed, mode_id, theta, k)

        em_btn.click(
            _resolve_factor,
            inputs=[sibling, em_mode, em_seed, em_theta, em_k],
            outputs=em_table,
        )
        sibling.change(
            lambda s: gr.Dropdown(choices=[label for label, _ in _factor_modes(s)], value=None),
            inputs=sibling,
            outputs=em_mode,
        )

    gr.Markdown(
        """---
**Cite:** Radzikowski and Chen 2026, *Epicure: Navigating the Emergent Geometry of Food Ingredient Embeddings*, arXiv:2605.22391.

Models: [epicure-cooc](https://huggingface.co/Kaikaku/epicure-cooc), [epicure-core](https://huggingface.co/Kaikaku/epicure-core), [epicure-chem](https://huggingface.co/Kaikaku/epicure-chem).
Dataset: [epicure-corpus-resources](https://huggingface.co/datasets/Kaikaku/epicure-corpus-resources).
"""
    )

if __name__ == "__main__":
    demo.launch()