"""Jazz Harmony Explorer — nearest-neighbor search over published tune embeddings.

Runs entirely from the released bundle (no chord data): cosine similarity
over 6,900 x 128 vectors, plus the precomputed 2-D UMAP layout for context.
"""

import csv

import gradio as gr
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from huggingface_hub import hf_hub_download

DATASET = "eigenben/jazz-harmony-embeddings"

_bundle = np.load(hf_hub_download(DATASET, "embeddings.npz", repo_type="dataset"))
EMB = _bundle["embeddings"]
with open(hf_hub_download(DATASET, "metadata.csv", repo_type="dataset"), newline="") as _handle:
    ROWS = list(csv.DictReader(_handle))
PROJ = np.load(hf_hub_download(DATASET, "projection_2d.npz", repo_type="dataset"))["projection"]

SURFACE = "#fcfcfb"
BASE = "#c3c2b7"
BLUE = "#2a78d6"
RED = "#e34948"
MUTED = "#898781"


def find_anchor(title: str) -> int | None:
    needle = title.strip().lower()
    if not needle:
        return None
    exact = next((i for i, r in enumerate(ROWS) if r["title"].strip().lower() == needle), None)
    if exact is not None:
        return exact
    return next((i for i, r in enumerate(ROWS) if needle in r["title"].lower()), None)


def neighbors(anchor: int, k: int) -> list[tuple[int, float]]:
    scores = EMB @ EMB[anchor]
    out = []
    seen = {ROWS[anchor]["title"].strip().lower()}
    for i in np.argsort(-scores):
        title = ROWS[i]["title"].strip().lower()
        if i == anchor or title in seen:
            continue
        seen.add(title)
        out.append((int(i), float(scores[i])))
        if len(out) == k:
            break
    return out


def map_figure(anchor: int, hits: list[tuple[int, float]]):
    fig, ax = plt.subplots(figsize=(7, 5.6), dpi=120)
    fig.patch.set_facecolor(SURFACE)
    ax.set_facecolor(SURFACE)
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.scatter(PROJ[:, 0], PROJ[:, 1], s=3, c=BASE, alpha=0.5, linewidths=0)
    hit_rows = [i for i, _ in hits]
    ax.scatter(PROJ[hit_rows, 0], PROJ[hit_rows, 1], s=42, c=BLUE, linewidths=0, zorder=3)
    ax.scatter([PROJ[anchor, 0]], [PROJ[anchor, 1]], s=110, c=RED, linewidths=0, zorder=4)
    ax.set_title(
        f'"{ROWS[anchor]["title"]}" (red) and its neighbors (blue)',
        fontsize=11, loc="left", color="#0b0b0b",
    )
    ax.text(
        0, -0.04,
        "2-D UMAP of all 6,900 tunes. Axes are arbitrary; only closeness is meaningful.",
        transform=ax.transAxes, color=MUTED, fontsize=8,
    )
    fig.tight_layout()
    return fig


def search(title: str, k: int):
    anchor = find_anchor(title)
    if anchor is None:
        return (
            gr.Markdown(f"No tune matches **{title}** — try a shorter substring."),
            None,
            None,
        )
    hits = neighbors(anchor, int(k))
    table = [
        [rank, f"{score:.3f}", ROWS[i]["title"], ROWS[i]["composer"], ROWS[i]["source"]]
        for rank, (i, score) in enumerate(hits, start=1)
    ]
    heading = gr.Markdown(
        f"### {ROWS[anchor]['title']}"
        + (f" — {ROWS[anchor]['composer']}" if ROWS[anchor]["composer"] else "")
        + f"\n`{ROWS[anchor]['id']}` — closest tunes by learned harmonic similarity:"
    )
    return heading, table, map_figure(anchor, hits)


with gr.Blocks(title="Jazz Harmony Explorer") as demo:
    gr.Markdown(
        "# Jazz Harmony Explorer\n"
        "A small transformer, trained from scratch on ~8,000 chord charts, placed every "
        "jazz standard in a 128-d vector space where related harmony sits close together. "
        "Search a tune to see its nearest harmonic neighbors — contrafacts (same chords, "
        "different melody) should surface without the model ever being told about them. "
        "Try **Oleo**, **Giant Steps**, **Cherokee**, or **Ornithology**.\n\n"
        "[Model](https://huggingface.co/eigenben/jazz-harmony-embeddings) · "
        "[Embeddings](https://huggingface.co/datasets/eigenben/jazz-harmony-embeddings) · "
        "[Code & write-up](https://github.com/eigenben/jazz-harmony-embeddings)"
    )
    with gr.Row():
        title_box = gr.Textbox(value="Oleo", label="tune title", scale=3)
        k_slider = gr.Slider(5, 30, value=10, step=1, label="neighbors", scale=1)
        go = gr.Button("Search", variant="primary", scale=1)
    heading = gr.Markdown()
    with gr.Row():
        table = gr.Dataframe(
            headers=["rank", "similarity", "title", "composer", "source"],
            interactive=False,
        )
        plot = gr.Plot(label="corpus map")
    for trigger in (go.click, title_box.submit, k_slider.release):
        trigger(search, inputs=[title_box, k_slider], outputs=[heading, table, plot])
    demo.load(search, inputs=[title_box, k_slider], outputs=[heading, table, plot])

demo.launch()