"""
TSM — Text Similarity Maker
Streamlit pipeline for building VOSviewer-compatible paper networks.
"""

import csv
import io
import json
import os
import tempfile
import uuid
from pathlib import Path

import numpy as np
import streamlit as st

# ── page config ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="TSM Text Similarity Maker",
    page_icon="🔬",
    layout="wide",
)

st.title("🔬 TSM — Text Similarity Maker")
st.header("Create a text similarity science map")
st.write("Generate embeddings of your documents' titles and abstracts — a numerical representation of their semantic content. Then use those embeddings to build a science map. Two map types are available: a text similarity network map, and an embedding space reduction map.")
st.markdown("Built by [Juan Pablo Bascur](https://jpbascur.com) — Problems? Contact [juanpablobascurcifuentes@gmail.com](mailto:juanpablobascurcifuentes@gmail.com)")
st.markdown("Source code: [github.com/jpbascur/text-similarity-maker](https://github.com/jpbascur/text-similarity-maker)")

# ── helpers ──────────────────────────────────────────────────────────────────

def array_to_csv_bytes(arr: np.ndarray, ids: list[str] | None = None) -> bytes:
    buf = io.StringIO()
    if ids is not None:
        for row_id, row in zip(ids, arr):
            buf.write(row_id + "," + ",".join(f"{v:.8f}" for v in row) + "\n")
    else:
        np.savetxt(buf, arr, delimiter=",", fmt="%.8f")
    return buf.getvalue().encode()

def csv_bytes_to_array(data: bytes) -> np.ndarray:
    """Parse an embeddings CSV. First column is always the paper ID and is dropped.
    No header row is expected or supported."""
    lines = [l for l in data.decode("utf-8").splitlines() if l.strip()]
    if not lines:
        raise ValueError("The embeddings file is empty.")
    first_cells = lines[0].split(",")
    try:
        float(first_cells[1])
    except (ValueError, IndexError):
        raise ValueError(
            "The embeddings file appears to have a header row. "
            "This file must not have a header — the first row should be data."
        )
    rows = []
    for line in lines:
        cells = line.split(",")
        rows.append([float(c) for c in cells[1:]])
    arr = np.array(rows, dtype=np.float32)
    if arr.shape[1] != 768:
        raise ValueError(
            f"Expected 768 embedding dimensions per row, got {arr.shape[1]}. "
            "This file does not look like a SPECTER2 embeddings file."
        )
    return arr

def show_array_info(arr: np.ndarray, label: str = "Array"):
    st.caption(f"{label}: {arr.shape[0]} papers × {arr.shape[1]} dimensions")

def parse_pubmed_export(text: str) -> list[dict]:
    """Parse a PubMed .txt or .nbib export into [{id, title, abstract}]."""
    import re
    records, current, current_tag = [], {}, None
    for line in text.splitlines():
        if re.match(r'^ER\s*-', line):
            if current:
                records.append(current)
            current, current_tag = {}, None
            continue
        m = re.match(r'^([A-Z]+)\s*-\s+(.*)', line)
        if m:
            current_tag = m.group(1)
            val = m.group(2).strip()
            current[current_tag] = (current[current_tag] + " " + val) if current_tag in current else val
        elif line.startswith("      ") and current_tag:
            current[current_tag] += " " + line.strip()
        elif not line.strip():
            if current:
                records.append(current)
            current, current_tag = {}, None
    if current:
        records.append(current)
    return [
        {"id": r["PMID"].strip(), "title": r["TI"].strip(), "abstract": r.get("AB", "").strip()}
        for r in records if "PMID" in r and "TI" in r
    ]

def parse_ris_export(text: str) -> list[dict]:
    """Parse a RIS (.ris) export into [{id, title, abstract}]."""
    import re
    records, current, current_tag = [], {}, None
    for line in text.splitlines():
        if re.match(r'^ER\s*-', line):
            if current:
                records.append(current)
            current, current_tag = {}, None
            continue
        m = re.match(r'^([A-Z][A-Z0-9])\s+-\s+(.*)', line)
        if m:
            current_tag = m.group(1)
            val = m.group(2).strip()
            current[current_tag] = (current[current_tag] + " " + val) if current_tag in current else val
        elif line.startswith("  ") and current_tag:
            current[current_tag] += " " + line.strip()
        elif not line.strip():
            if current:
                records.append(current)
            current, current_tag = {}, None
    if current:
        records.append(current)
    result = []
    for i, r in enumerate(records):
        id_ = r.get("ID") or r.get("AN") or r.get("UT") or r.get("DO") or str(i + 1)
        title = r.get("TI") or r.get("T1", "")
        abstract = r.get("AB") or r.get("N2", "")
        if title:
            result.append({"id": id_.strip(), "title": title.strip(), "abstract": abstract.strip()})
    return result

def parse_bibtex_export(text: str) -> list[dict]:
    """Parse a BibTeX (.bib) export into [{id, title, abstract}]."""
    import re
    result = []
    for entry in re.split(r'(?=@\w+\{)', text):
        key_m = re.match(r'@\w+\{([^,\n]+),', entry)
        if not key_m:
            continue
        key = key_m.group(1).strip()
        def _field(name):
            m = re.search(rf'\b{name}\s*=\s*\{{([^{{}}]*(?:\{{[^{{}}]*\}}[^{{}}]*)*)\}}', entry, re.IGNORECASE)
            if not m:
                m = re.search(rf'\b{name}\s*=\s*"([^"]*)"', entry, re.IGNORECASE)
            if m:
                return re.sub(r'\{([^{}]*)\}', r'\1', m.group(1)).strip()
            return ""
        title = _field("title")
        abstract = _field("abstract")
        if title:
            result.append({"id": key, "title": title, "abstract": abstract})
    return result

def save_upload(file_obj, state_key: str):
    if file_obj is not None:
        st.session_state[state_key] = (file_obj.name, file_obj.read())

# Session-unique ID for static file naming (avoids collisions between users)
if "session_id" not in st.session_state:
    st.session_state["session_id"] = uuid.uuid4().hex

if "running" not in st.session_state:
    st.session_state["running"] = False


_is_running = st.session_state["running"]
if _is_running:
    st.warning("A job is already running in this session. Please wait for it to finish.")

try:
    def _read_cgroup_mem():
        """Read container memory limits from cgroup (accurate inside Docker)."""
        # Try cgroups v2 first
        try:
            limit = int(Path("/sys/fs/cgroup/memory.max").read_text().strip())
            usage = int(Path("/sys/fs/cgroup/memory.current").read_text().strip())
            return usage, limit
        except Exception:
            pass
        # Fall back to cgroups v1
        limit = int(Path("/sys/fs/cgroup/memory/memory.limit_in_bytes").read_text().strip())
        usage = int(Path("/sys/fs/cgroup/memory/memory.usage_in_bytes").read_text().strip())
        return usage, limit

    _mem_used, _mem_limit = _read_cgroup_mem()
    _mem_used_gb  = _mem_used  / 1024**3
    _mem_total_gb = _mem_limit / 1024**3
    _mem_free_gb  = (_mem_limit - _mem_used) / 1024**3
    _mem_pct      = _mem_used / _mem_limit * 100
    _mem_caption  = (
        f"Container memory: {_mem_used_gb:.1f} GB used / {_mem_total_gb:.1f} GB total — "
        f"{_mem_free_gb:.1f} GB free ({100 - _mem_pct:.0f}% available). "
        "This tool has limited memory shared across all users. "
        "If memory is low, please wait for it to be released by another user."
    )
except Exception:
    _mem_caption = None

_STATIC_DIR = Path(__file__).parent / "static"
_STATIC_DIR.mkdir(exist_ok=True)

def _write_static_json(filename: str, data: dict) -> str:
    """Write a VOSviewer JSON to the static directory and return its public URL."""
    path = _STATIC_DIR / filename
    path.write_text(json.dumps(data), encoding="utf-8")
    space_id = os.environ.get("SPACE_ID", "")
    if space_id:
        slug = space_id.replace("/", "-").lower()
        return f"https://{slug}.hf.space/app/static/{filename}"
    return f"http://localhost:8501/app/static/{filename}"


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# ROOT — EMBEDDINGS
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
with st.expander("Text to Embeddings", expanded=True):
    st.subheader("Text to Embeddings")
    st.caption("Encodes each paper into a numerical vector that captures its semantic meaning. Uses [SPECTER2](https://github.com/allenai/SPECTER2) with the proximity adapter, a transformer model optimized to generate embeddings of paper titles and abstracts such that semantically similar papers end up with similar embeddings.")

    st.caption("Requirements: CSV columns: id, title, abstract. Avoid including papers without a title or abstract — they will produce poor-quality embeddings.")
    save_upload(
        st.file_uploader("Upload papers file", type=["csv"], key="s1_upload"),
        "s1_file",
    )
    ignore_incomplete = st.checkbox("Ignore documents without title or abstract", value=True, key="s1_ignore_incomplete")

    if "s1_file" in st.session_state and "step1_papers" not in st.session_state:
        from pipeline.embed import load_papers
        fname, raw = st.session_state["s1_file"]
        suffix = Path(fname).suffix
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(raw)
            tmp_path = Path(tmp.name)
        try:
            st.session_state["step1_papers"] = load_papers(tmp_path)
        except Exception as e:
            st.error(f"Could not load papers: {e}")
        finally:
            tmp_path.unlink(missing_ok=True)

    if "step1_papers" in st.session_state:
        st.caption(f"{len(st.session_state['step1_papers'])} papers loaded.")

    col_run, col_fallback = st.columns([3, 2])
    s1_run = col_run.button("Run Embeddings", key="run_embed",
                            disabled=_is_running or "step1_papers" not in st.session_state)
    use_fallback = col_fallback.checkbox(
        "Use fallback resources", key="use_fallback", value=False,
        help=(
            "By default, embeddings are computed on your GPU using WebGPU — fast and uses no server resources. "
            "Check this to use the web server instead, which runs on CPU and is much slower. "
            "Use this if your browser does not support WebGPU "
            "(Chrome, Edge and Opera support it by default; Firefox and Safari need additional configuration)."
        ),
    )
    if _mem_caption:
        st.caption(_mem_caption)

    has_papers_to_run = "step1_papers" in st.session_state

    if s1_run and has_papers_to_run:
        papers = st.session_state["step1_papers"]
        if ignore_incomplete:
            papers = [p for p in papers if p.get("title", "").strip() and p.get("abstract", "").strip()]

        if use_fallback:
            from pipeline.embed import embed_papers
            st.session_state["running"] = True
            prog = st.progress(0, text="Loading SPECTER2 model…")
            def _cb(cur, tot):
                prog.progress(cur / tot, text=f"Encoding {cur}/{tot}…")
            try:
                embeddings = embed_papers(papers, progress_callback=_cb)
                prog.progress(1.0, text="Done.")
                st.session_state["step1_embeddings"] = embeddings
            finally:
                st.session_state["running"] = False
            st.session_state["use_step1_net"]          = True
            st.session_state["use_step1_umap"]         = True
            st.session_state["use_step1_meta"]         = True
            st.session_state["use_step1_vos_map_meta"] = True
            st.rerun()
        else:
            st.session_state["webgpu_papers"]  = papers
            st.session_state["webgpu_run"]     = True

    # WebGPU component (only shown when not using fallback)
    if not use_fallback and has_papers_to_run:
        from component import webgpu_embed
        papers_for_gpu = st.session_state.get("webgpu_papers", [])
        run_flag       = st.session_state.get("webgpu_run", False)
        try:
            result = webgpu_embed(papers=papers_for_gpu, run=run_flag, key="webgpu_embedder")
            if result is not None:
                st.session_state["step1_embeddings"] = result
                st.session_state["webgpu_run"]       = False
                st.session_state["use_step1_net"]          = True
                st.session_state["use_step1_umap"]         = True
                st.session_state["use_step1_meta"]         = True
                st.session_state["use_step1_vos_map_meta"] = True
                st.rerun()
        except RuntimeError as e:
            if str(e) != "webgpu_not_supported":
                st.error(f"Embedding error: {e}")

    has_embed_dl = "step1_embeddings" in st.session_state
    if has_embed_dl:
        embeddings = st.session_state["step1_embeddings"]
        papers     = st.session_state["step1_papers"]
        show_array_info(embeddings, "Embeddings ready")
        embed_dl_data = array_to_csv_bytes(embeddings, ids=[p["id"] for p in papers])
    else:
        embed_dl_data = b""
    st.download_button(
        "Download embeddings (.csv)", embed_dl_data,
        "embeddings.csv", mime="text/csv", key="dl_embed_csv",
        disabled=not has_embed_dl,
    )

    with st.expander("Transform file to supported format", expanded=False):
        st.caption("Converts reference exports to the CSV format the app needs. Supported formats:")
        st.caption("**PubMed** (.txt, .nbib) — Send to → File → Format: PubMed  \n**RIS** (.ris) — exported by Scopus, Web of Science, Zotero, Mendeley, EndNote  \n**BibTeX** (.bib) — exported by Google Scholar, Zotero, most reference managers")
        ref_file = st.file_uploader("Upload export file", type=["txt", "nbib", "ris", "bib"], key="ref_upload")
        if ref_file:
            try:
                text = ref_file.read().decode("utf-8", errors="replace")
                ext = Path(ref_file.name).suffix.lower()
                if ext == ".bib":
                    papers_ref = parse_bibtex_export(text)
                elif ext == ".ris":
                    papers_ref = parse_ris_export(text)
                else:
                    papers_ref = parse_pubmed_export(text)
                n_total = len(papers_ref)
                n_abstract = sum(1 for p in papers_ref if p["abstract"])
                st.caption(f"{n_total} papers found — {n_abstract} with abstracts, {n_total - n_abstract} without.")
                if n_total > 0:
                    buf = io.StringIO()
                    w = csv.DictWriter(buf, fieldnames=["id", "title", "abstract"])
                    w.writeheader()
                    w.writerows(papers_ref)
                    csv_bytes = buf.getvalue().encode()
                    col_a, col_b = st.columns(2)
                    if col_a.button("Load into app", key="load_ref"):
                        st.session_state["s1_file"] = ("papers.csv", csv_bytes)
                        st.session_state.pop("step1_papers", None)
                        st.session_state.pop("step1_embeddings", None)
                        st.rerun()
                    col_b.download_button(
                        "Download as CSV", csv_bytes,
                        "papers.csv", mime="text/csv", key="dl_ref",
                    )
            except Exception as e:
                st.error(f"Could not parse file: {e}")

    with st.expander("Don't have data? Try the demo data to start", expanded=False):
        st.caption("500 sample papers to try the tool without your own data.")
        demo_path = Path(__file__).parent / "sample_papers.csv"
        demo_bytes = demo_path.read_bytes()
        col_a, col_b = st.columns(2)
        if col_a.button("Load demo data", key="load_demo"):
            st.session_state["s1_file"] = ("sample_papers.csv", demo_bytes)
            st.session_state.pop("step1_papers", None)
            st.session_state.pop("step1_embeddings", None)
            st.rerun()
        col_b.download_button(
            "Download demo data", demo_bytes,
            "sample_papers.csv", mime="text/csv", key="dl_demo",
        )

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# BRANCH A — NETWORK MAP (VOSviewer)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
with st.expander("Text Similarity Network Map", expanded=False):
    st.header("Text Similarity Network Map")
    st.caption("Traditional science map. Builds a cosine similarity network and exports it to VOSviewer — emulating the citation-based approach used in bibliometrics.")
    st.caption("Each document is connected to its most similar documents based on cosine similarity of their embeddings. The resulting network can be visualized and explored in VOSviewer.")

    # ── Create Network ──
    with st.container(border=True):
        st.subheader("Create Network")

        embed_src = None
        has_embeddings = "step1_embeddings" in st.session_state
        st.caption("Requirements: CSV with no header. First column: paper ID. Remaining 768 columns: embedding values. Generated by the Text to Embeddings step.")
        use_above_net = st.checkbox("Use embeddings from Text to Embeddings", value=has_embeddings, key="use_step1_net", disabled=not has_embeddings)
        if use_above_net and has_embeddings:
            embed_src = st.session_state["step1_embeddings"]
        save_upload(
            st.file_uploader("Or upload an embeddings file (.csv)", type=["csv"], key="s2_upload"),
            "s2_file",
        )
        if not use_above_net and "s2_file" in st.session_state:
            try:
                embed_src = csv_bytes_to_array(st.session_state["s2_file"][1])
            except Exception as e:
                st.error(f"Could not load file: {e}")

        n_papers = len(embed_src) if embed_src is not None else None
        c1, c2 = st.columns(2)
        c1.caption("A document connects to its most similar documents. Higher values produce a denser network.")
        top_k = c1.number_input(
            "Maximum number of connections per document", min_value=1, max_value=n_papers - 1 if n_papers else None,
            value=20, step=1, key="top_k",
        )
        c2.caption("Minimum similarity required to keep a connection. Higher values produce a sparser network.")
        min_sim = c2.slider("Min similarity", 0.0, 1.0, 0.0, 0.01, key="min_sim")

        s2_run = st.button("Build Network", key="run_network", disabled=_is_running or embed_src is None)

        if s2_run and embed_src is not None:
            from pipeline.network import build_edge_list
            st.session_state["running"] = True
            prog = st.progress(0, text="Building edge list…")
            def _cb(cur, tot):
                prog.progress(cur / tot, text=f"Processing {cur}/{tot} papers…")
            try:
                with st.spinner("Computing cosine similarities…"):
                    edges = build_edge_list(embed_src, top_k=int(top_k), min_similarity=float(min_sim), progress_callback=_cb)
                prog.progress(1.0, text="Done.")
                st.session_state["step2_edges"] = edges
            finally:
                st.session_state["running"] = False
            st.session_state["use_step2_edges"] = True
            st.rerun()

        has_edges_dl = "step2_edges" in st.session_state
        if has_edges_dl:
            edges = st.session_state["step2_edges"]
            st.caption(f"{len(edges)} edges ready.")
            edge_buf = io.StringIO()
            ew = csv.writer(edge_buf)
            ew.writerow(["source", "target", "weight"])
            ew.writerows(edges)
            edge_dl_data = edge_buf.getvalue().encode()
        else:
            edge_dl_data = b""
        st.download_button(
            "Download edge list (.csv)", edge_dl_data,
            "network.csv", mime="text/csv", key="dl_network",
            disabled=not has_edges_dl,
        )

    # ── VOSviewer Export ──
    with st.container(border=True):
        st.subheader("Visualize Network with VOSviewer")

        edges_src = None
        has_edges = "step2_edges" in st.session_state
        st.caption("Requirements: CSV columns: source, target, weight. Generated by the Create Network step.")
        use_above_edges = st.checkbox("Use edge list from Create Network", value=has_edges, key="use_step2_edges", disabled=not has_edges)
        if use_above_edges and has_edges:
            edges_src = st.session_state["step2_edges"]
        save_upload(
            st.file_uploader("Or upload an edge list CSV", type=["csv"], key="s3_edges_upload"),
            "s3_edges_file",
        )
        if not use_above_edges and "s3_edges_file" in st.session_state:
            _, raw = st.session_state["s3_edges_file"]
            edges_src = [
                (int(r["source"]), int(r["target"]), float(r["weight"]))
                for r in csv.DictReader(io.StringIO(raw.decode()))
            ]

        papers_src = None
        has_papers = "step1_papers" in st.session_state
        st.caption("Requirements: CSV columns: id, title. Your original papers CSV works here — it already has these columns.")
        use_above_meta = st.checkbox("Use papers from Text to Embeddings", value=has_papers, key="use_step1_meta", disabled=not has_papers)
        if use_above_meta and has_papers:
            papers_src = st.session_state["step1_papers"]
        save_upload(
            st.file_uploader("Or upload a papers CSV (id, title)", type=["csv"], key="s3_meta_upload"),
            "s3_meta_file",
        )
        if not use_above_meta and "s3_meta_file" in st.session_state:
            _, raw = st.session_state["s3_meta_file"]
            papers_src = list(csv.DictReader(io.StringIO(raw.decode())))

        s3_run = st.button("Generate VOSviewer Map", key="run_vos",
                           disabled=(edges_src is None or papers_src is None))

        if s3_run and edges_src is not None and papers_src is not None:
            with st.spinner("Generating VOSviewer map…"):
                items = []
                for idx, paper in enumerate(papers_src):
                    items.append({
                        "id": str(idx + 1),
                        "label": paper.get("title", paper.get("id", str(idx + 1))),
                        "description": str(paper.get("id", "")),
                    })
                items_no_cluster = [{**item, "cluster": 1} for item in items]
                links = [
                    {"source_id": str(i + 1), "target_id": str(j + 1), "strength": round(w, 6)}
                    for i, j, w in edges_src
                ]
                sid = st.session_state["session_id"]
                vos_data_auto  = {"network": {"items": items,            "links": links}}
                vos_data_fixed = {"network": {"items": items_no_cluster, "links": links}}
                st.session_state["vos_json_auto"]    = json.dumps(vos_data_auto,  indent=2)
                st.session_state["vos_json_fixed"]   = json.dumps(vos_data_fixed, indent=2)
                st.session_state["vos_json_url_auto"]  = _write_static_json(f"{sid}_network_auto.json",  vos_data_auto)
                st.session_state["vos_json_url_fixed"] = _write_static_json(f"{sid}_network_fixed.json", vos_data_fixed)
            st.session_state["vos_do_cluster"] = True
            st.rerun()

        has_vos_json = "vos_json_auto" in st.session_state
        do_cluster = st.checkbox(
            "Open and cluster (may be slow in dense networks)",
            key="vos_do_cluster", value=True, disabled=not has_vos_json,
        )
        vos_dl_data = st.session_state["vos_json_auto" if do_cluster else "vos_json_fixed"].encode() if has_vos_json else b""
        vos_url = (
            f"https://app.vosviewer.com/?json={st.session_state['vos_json_url_auto' if do_cluster else 'vos_json_url_fixed']}&max_n_links=0"
            if has_vos_json else "https://app.vosviewer.com/"
        )
        st.download_button(
            "Download VOSviewer map (.json)", vos_dl_data,
            "vosviewer_network.json", mime="application/json", key="dl_vos_json",
            disabled=not has_vos_json,
        )
        st.link_button("🗺️ Open in VOSviewer Online", vos_url, type="primary", use_container_width=True, disabled=not has_vos_json)

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# BRANCH B — EMBEDDING SPACE MAP (UMAP)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
with st.expander("Embedding Space Reduction Map", expanded=False):
    st.header("Embedding Space Reduction Map")
    st.caption("Alternative science map. Projects embeddings directly into 2D space with UMAP — more faithful to the semantic structure of the embeddings.")
    st.caption("UMAP finds a 2D layout that preserves the high-dimensional relationships between documents as faithfully as possible. Similar documents end up close together; dissimilar ones far apart.")

    # ── UMAP ──
    with st.container(border=True):
        st.subheader("Generate 2D Coordinates")

        viz_embed_src = None
        viz_embed_ids = []
        has_embeddings_umap = "step1_embeddings" in st.session_state
        st.caption("Requirements: CSV with no header. First column: paper ID. Remaining 768 columns: embedding values. Generated by the Text to Embeddings step.")
        use_above_umap = st.checkbox("Use embeddings from Text to Embeddings", value=has_embeddings_umap, key="use_step1_umap", disabled=not has_embeddings_umap)
        if use_above_umap and has_embeddings_umap:
            viz_embed_src = st.session_state["step1_embeddings"]
            viz_embed_ids = [p["id"] for p in st.session_state["step1_papers"]]
        save_upload(
            st.file_uploader("Or upload an embeddings file (.csv)", type=["csv"], key="viz_embed_upload"),
            "viz_embed_file",
        )
        if not use_above_umap and "viz_embed_file" in st.session_state:
            _, raw = st.session_state["viz_embed_file"]
            try:
                viz_embed_src = csv_bytes_to_array(raw)
                viz_embed_ids = [l.split(",")[0] for l in raw.decode("utf-8").splitlines() if l.strip()]
            except Exception as e:
                st.error(f"Could not load file: {e}")

        c1, c2 = st.columns(2)
        c1.caption("Number of documents to consider at a time when preserving the embedding space structure. Low values keep local structure, high values keep global structure.")
        umap_n_neighbors = c1.number_input(
            "n_neighbors", min_value=2, value=15, step=1, key="umap_n_neighbors",
        )
        c2.caption("Minimum distance between documents in the 2D projection. Low values place similar documents into tight clumps, high values spread them more uniformly.")
        umap_min_dist = c2.slider(
            "min_dist", 0.0, 1.0, 0.1, 0.01, key="umap_min_dist",
        )

        if viz_embed_src is not None:
            n = len(viz_embed_src)
            estimate = "a few seconds" if n < 500 else "~30 seconds" if n < 2000 else "a few minutes"
            st.caption(f"Expected time: {estimate} ({n} papers)")

        sa_run = st.button("Run UMAP", key="run_umap", disabled=_is_running or viz_embed_src is None)

        if sa_run and viz_embed_src is not None:
            from pipeline.reduce import umap_reduce
            st.session_state["running"] = True
            try:
                with st.spinner(f"Running UMAP on {n} papers… ({estimate})"):
                    coords = umap_reduce(viz_embed_src, n_neighbors=int(umap_n_neighbors), min_dist=float(umap_min_dist))
                st.session_state["viz_coords"] = coords
                st.session_state["viz_ids"]    = viz_embed_ids
            finally:
                st.session_state["running"] = False
            st.session_state["use_viz_coords_vos"] = True
            st.rerun()

        has_coords_dl = "viz_coords" in st.session_state
        if has_coords_dl:
            st.caption(f"{len(st.session_state['viz_coords'])} points projected.")
            coords_buf = io.StringIO()
            coords_buf.write("id,x,y\n")
            for pid, (x, y) in zip(st.session_state["viz_ids"], st.session_state["viz_coords"]):
                coords_buf.write(f"{pid},{x:.6f},{y:.6f}\n")
            coords_dl_data = coords_buf.getvalue().encode()
        else:
            coords_dl_data = b""
        st.download_button(
            "Download coords (.csv)", coords_dl_data,
            "coords.csv", mime="text/csv", key="dl_coords",
            disabled=not has_coords_dl,
        )

    # ── VOSviewer Map Export ──
    with st.container(border=True):
        st.subheader("Visualize with VOSviewer")
        st.caption("Generates a VOSviewer map file with UMAP coordinates, so VOSviewer positions nodes according to the projection.")

        vos_coords = None
        vos_coord_ids = []
        has_coords = "viz_coords" in st.session_state
        st.caption("Requirements: CSV columns: id, x, y. Generated by the Generate 2D Coordinates step.")
        use_above_vos_coords = st.checkbox("Use coordinates from Generate 2D Coordinates", value=has_coords, key="use_viz_coords_vos", disabled=not has_coords)
        if use_above_vos_coords and has_coords:
            vos_coords    = st.session_state["viz_coords"]
            vos_coord_ids = st.session_state["viz_ids"]
        save_upload(
            st.file_uploader("Or upload a coordinates CSV (id, x, y)", type=["csv"], key="vos_coords_upload"),
            "vos_coords_file",
        )
        if not use_above_vos_coords and "vos_coords_file" in st.session_state:
            _, raw = st.session_state["vos_coords_file"]
            rows = list(csv.DictReader(io.StringIO(raw.decode())))
            vos_coord_ids = [r["id"] for r in rows]
            vos_coords    = np.array([[float(r["x"]), float(r["y"])] for r in rows], dtype=np.float32)

        vos_map_papers = None
        has_papers_vos = "step1_papers" in st.session_state
        st.caption("Requirements: CSV columns: id, title. Your original papers CSV works here — it already has these columns.")
        use_above_vos_meta = st.checkbox("Use papers from Text to Embeddings", value=has_papers_vos, key="use_step1_vos_map_meta", disabled=not has_papers_vos)
        if use_above_vos_meta and has_papers_vos:
            vos_map_papers = st.session_state["step1_papers"]
        save_upload(
            st.file_uploader("Or upload a papers CSV (id, title)", type=["csv"], key="vos_map_meta_upload"),
            "vos_map_meta_file",
        )
        if not use_above_vos_meta and "vos_map_meta_file" in st.session_state:
            _, raw = st.session_state["vos_map_meta_file"]
            vos_map_papers = list(csv.DictReader(io.StringIO(raw.decode())))

        sc_run = st.button("Generate VOSviewer Map", key="run_vos_map",
                           disabled=(vos_coords is None or vos_map_papers is None))

        if sc_run and vos_coords is not None and vos_map_papers is not None:
            with st.spinner("Generating VOSviewer map…"):
                coord_lookup = {pid: (float(x), float(y)) for pid, (x, y) in zip(vos_coord_ids, vos_coords)}
                items = []
                for idx, paper in enumerate(vos_map_papers):
                    pid   = str(paper.get("id", idx + 1))
                    label = paper.get("title", pid)
                    x, y  = coord_lookup.get(pid, (0.0, 0.0))
                    items.append({
                        "id": str(idx + 1),
                        "label": label,
                        "description": pid,
                        "x": round(x, 6),
                        "y": round(y, 6),
                        "cluster": 1,
                    })
                vos_data = {"network": {"items": items, "links": []}}
                sid = st.session_state["session_id"]
                url = _write_static_json(f"{sid}_umap.json", vos_data)
                st.session_state["viz_vos_json"]     = json.dumps(vos_data, indent=2)
                st.session_state["viz_vos_json_url"] = url
            st.rerun()

        has_viz_vos_json = "viz_vos_json" in st.session_state
        viz_vos_dl_data = st.session_state["viz_vos_json"].encode() if has_viz_vos_json else b""
        viz_vos_url = f"https://app.vosviewer.com/?json={st.session_state['viz_vos_json_url']}" if has_viz_vos_json else "https://app.vosviewer.com/"
        st.download_button(
            "Download VOSviewer map (.json)", viz_vos_dl_data,
            "vosviewer_umap.json", mime="application/json", key="dl_viz_vos_json",
            disabled=not has_viz_vos_json,
        )
        st.link_button("🗺️ Open in VOSviewer Online", viz_vos_url, type="primary", use_container_width=True, disabled=not has_viz_vos_json)