"""Epicure Explorer: chef-facing operators over the three sibling embeddings.

Four tabs:
  - Basket pairings: pick 1+ ingredients, get neighbours and closest modes of the basket centroid.
  - Supervised SLERP: rotate a (possibly multi-ingredient) seed toward 1+ supervised poles.
  - Emergent SLERP: rotate a (possibly multi-ingredient) seed toward 1+ emergent factor modes.
  - Arithmetic: Mikolov-style 'positives - negatives' returning nearest neighbours.

All three siblings (Cooc, Core, Chem) load on startup from public HF model repos.
"""

from __future__ import annotations

import os
import sys
import numpy as np
import gradio as gr

# epicure.py is shipped alongside this app.py in the Space; fall back to HF if absent.
try:
    from epicure import Epicure
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

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 _unit(v: np.ndarray, eps: float = 1e-9) -> np.ndarray:
    n = np.linalg.norm(v)
    return v / max(n, eps)


def _basket_centroid(m: Epicure, names: list[str]) -> np.ndarray:
    """L2-normalised mean of the unit vectors of the named ingredients."""
    valid = [n for n in (names or []) if n in m.vocab]
    if not valid:
        return None
    idxs = [m.vocab[n] for n in valid]
    centroid = m.E[idxs].mean(axis=0)
    return _unit(centroid)


def _stack_directions(m: Epicure, keys: list[str], use_factor_pole: bool = False) -> np.ndarray:
    """L2-normalised sum of the named supervised pole vectors (or factor mode poles)."""
    poles = []
    for k in keys or []:
        if use_factor_pole:
            for mode in m.modes:
                if mode.mode_id == k:
                    poles.append(_unit(mode.pole))
                    break
        else:
            if k in m.supervised_poles:
                poles.append(_unit(m.supervised_poles[k]))
    if not poles:
        return None
    return _unit(np.stack(poles, axis=0).sum(axis=0))


def _topk_from_query(m: Epicure, q: np.ndarray, k: int, exclude: list[str]) -> list[tuple[str, float]]:
    sims = m.E @ q
    for name in exclude or []:
        if name in m.vocab:
            sims[m.vocab[name]] = -np.inf
    order = np.argsort(-sims)
    return [(m.itos[int(i)], float(sims[i])) for i in order[:k]]


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


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


# ===== Tab 1: Basket pairings =====
def basket_pairings(sibling: str, basket: list[str], k: int):
    m = MODELS[sibling]
    centroid = _basket_centroid(m, basket)
    if centroid is None:
        return [], []
    nb = _topk_from_query(m, centroid, k=k, exclude=basket or [])
    # Closest modes to the basket centroid
    scored = [
        (mode.mode_id, mode.label, mode.kind, float(_unit(mode.pole) @ centroid))
        for mode in m.modes
    ]
    scored.sort(key=lambda x: -x[3])
    return (
        [[name, f"{sim:.4f}"] for name, sim in nb],
        [[mid, label, kind, f"{sim:.4f}"] for mid, label, kind, sim in scored[:k]],
    )


# ===== Tab 2: Supervised SLERP (multi-direction, multi-seed) =====
def supervised_slerp_multi(sibling: str, basket: list[str], directions: list[str], theta: float, k: int):
    m = MODELS[sibling]
    v = _basket_centroid(m, basket)
    d = _stack_directions(m, directions, use_factor_pole=False)
    if v is None or d is None:
        return []
    # SLERP from v toward d
    d_perp = d - (d @ v) * v
    n_perp = np.linalg.norm(d_perp)
    if n_perp < 1e-9:
        return _topk_from_query(m, v, k=k, exclude=basket or [])
    d_perp = d_perp / n_perp
    theta_rad = np.deg2rad(float(theta))
    q = _unit(np.cos(theta_rad) * v + np.sin(theta_rad) * d_perp)
    hits = _topk_from_query(m, q, k=k, exclude=basket or [])
    return [[name, f"{sim:.4f}"] for name, sim in hits]


# ===== Tab 3: Emergent SLERP (multi-direction, multi-seed) =====
def emergent_slerp_multi(sibling: str, basket: list[str], mode_labels: list[str], theta: float, k: int):
    m = MODELS[sibling]
    # Resolve label strings back to mode_ids
    label_to_id = {f"{mode.label} ({mode.mode_id})": mode.mode_id for mode in m.modes if mode.kind == "factor"}
    mode_ids = [label_to_id[lab] for lab in (mode_labels or []) if lab in label_to_id]
    v = _basket_centroid(m, basket)
    d = _stack_directions(m, mode_ids, use_factor_pole=True)
    if v is None or d is None:
        return []
    d_perp = d - (d @ v) * v
    n_perp = np.linalg.norm(d_perp)
    if n_perp < 1e-9:
        return [[n, f"{s:.4f}"] for n, s in _topk_from_query(m, v, k=k, exclude=basket or [])]
    d_perp = d_perp / n_perp
    theta_rad = np.deg2rad(float(theta))
    q = _unit(np.cos(theta_rad) * v + np.sin(theta_rad) * d_perp)
    hits = _topk_from_query(m, q, k=k, exclude=basket or [])
    return [[name, f"{sim:.4f}"] for name, sim in hits]


# ===== Tab 4: Mikolov arithmetic =====
def arithmetic(sibling: str, positives: list[str], negatives: list[str], k: int):
    m = MODELS[sibling]
    pos = _basket_centroid(m, positives)
    if pos is None:
        return []
    neg = _basket_centroid(m, negatives) if negatives else None
    if neg is None:
        q = pos
    else:
        # pos - neg, then renormalise. This is the king - man + woman pattern reshaped:
        # the user supplies the 'positives' and 'negatives' sets directly.
        q = _unit(pos - neg)
    exclude = (positives or []) + (negatives or [])
    hits = _topk_from_query(m, q, k=k, exclude=exclude)
    return [[name, f"{sim:.4f}"] for name, sim in hits]


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

Chef-facing operators over the three Epicure sibling embeddings (Cooc, Core, Chem),
from [arXiv:2605.22391](https://arxiv.org/abs/2605.22391).

- **Cooc** walks recipe co-occurrence only. Neighbours are recipe companions.
- **Core** blends typed FlavorDB compound walks with injected I-I 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")

    # -------- Tab 1: Basket pairings --------
    with gr.Tab("Basket pairings"):
        gr.Markdown(
            "Pick one or more ingredients. The tool averages their unit vectors and returns "
            "what is nearest to that centroid in the embedding. Useful for 'what should I add "
            "to the ingredients I already have?'"
        )
        basket = gr.Dropdown(
            choices=ALL_INGREDIENTS,
            value=["chicken", "lemon", "garlic"],
            label="Ingredient basket (pick 1+)",
            multiselect=True,
            max_choices=10,
        )
        k_pair = gr.Slider(1, 15, value=8, 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 to basket centroid", interactive=False
            )
            mode_table = gr.Dataframe(
                headers=["Mode id", "Label", "Kind", "Cosine"],
                label="Closest modes (factor + supervised)", interactive=False
            )
        pair_btn.click(basket_pairings, inputs=[sibling, basket, k_pair], outputs=[nb_table, mode_table])

    # -------- Tab 2: Supervised SLERP (multi) --------
    with gr.Tab("Supervised SLERP"):
        gr.Markdown(
            "Rotate the (possibly multi-ingredient) seed toward one or more supervised direction poles. "
            "Multiple directions are summed and L2-normalised before rotation, matching the paper's "
            "'chicken + processed + Western_Atlantic' style multi-constraint queries."
        )
        sup_basket = gr.Dropdown(
            choices=ALL_INGREDIENTS, value=["rice"], label="Seed basket (pick 1+)",
            multiselect=True, max_choices=10,
        )
        sup_dirs = gr.Dropdown(
            choices=_supervised_choices("chem"),
            value=["cuisine:South_Asian"],
            label="Supervised directions (pick 1+; summed before rotation)",
            multiselect=True, max_choices=5,
        )
        sup_theta = gr.Slider(0, 90, value=30, step=5, label="Rotation angle (deg)")
        sup_k = gr.Slider(1, 15, value=8, 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_multi,
            inputs=[sibling, sup_basket, sup_dirs, sup_theta, sup_k],
            outputs=sup_table,
        )
        sibling.change(
            lambda s: gr.Dropdown(choices=_supervised_choices(s), value=[]),
            inputs=sibling, outputs=sup_dirs,
        )

    # -------- Tab 3: Emergent SLERP (multi) --------
    with gr.Tab("Emergent SLERP"):
        gr.Markdown(
            "Rotate the seed basket toward one or more emergent factor-mode poles discovered "
            "by multi-seed-stable FastICA + GMM. Stack mode targets to combine culinary axes."
        )
        em_basket = gr.Dropdown(
            choices=ALL_INGREDIENTS, value=["chocolate"], label="Seed basket (pick 1+)",
            multiselect=True, max_choices=10,
        )
        factor_opts = _factor_mode_choices("chem")
        em_modes = gr.Dropdown(
            choices=[label for label, _ in factor_opts],
            value=[factor_opts[0][0]] if factor_opts else [],
            label="Factor modes (pick 1+; summed before rotation)",
            multiselect=True, max_choices=5,
        )
        em_theta = gr.Slider(0, 90, value=30, step=5, label="Rotation angle (deg)")
        em_k = gr.Slider(1, 15, value=8, step=1, label="K")
        em_btn = gr.Button("Rotate", variant="primary")
        em_table = gr.Dataframe(headers=["Ingredient", "Cosine"], label="Top-K rotated-query neighbours")
        em_btn.click(
            emergent_slerp_multi,
            inputs=[sibling, em_basket, em_modes, em_theta, em_k],
            outputs=em_table,
        )
        sibling.change(
            lambda s: gr.Dropdown(choices=[label for label, _ in _factor_mode_choices(s)], value=[]),
            inputs=sibling, outputs=em_modes,
        )

    # -------- Tab 4: Mikolov arithmetic --------
    with gr.Tab("Arithmetic"):
        gr.Markdown(
            "Classic Mikolov-style vector arithmetic: `centroid(positives) - centroid(negatives)`, "
            "then top-K nearest neighbours. Try `miso - salty` (no negative-set), or `chicken - "
            "Western + Asian` style queries (split your own intuition into positives and negatives)."
        )
        pos_box = gr.Dropdown(
            choices=ALL_INGREDIENTS, value=["miso"],
            label="Positives (added)", multiselect=True, max_choices=10,
        )
        neg_box = gr.Dropdown(
            choices=ALL_INGREDIENTS, value=[],
            label="Negatives (subtracted)", multiselect=True, max_choices=10,
        )
        ar_k = gr.Slider(1, 15, value=8, step=1, label="K")
        ar_btn = gr.Button("Compute", variant="primary")
        ar_table = gr.Dataframe(headers=["Ingredient", "Cosine"], label="Top-K nearest to result vector")
        ar_btn.click(arithmetic, inputs=[sibling, pos_box, neg_box, ar_k], outputs=ar_table)

    gr.Markdown(
        """---
**Cite:** Radzikowski and Chen, 2026, *Epicure: Navigating the Emergent Geometry of Food Ingredient Embeddings*, [arXiv:2605.22391](https://arxiv.org/abs/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()