"""
DiscourseGeometry
H4rmony Project · Jorge Vallego · https://theh4rmonyproject.org

Reference implementation of
  Vallego, J. (2026). A spatial-geometric framework for discourse comparison: Centroid projection,
  Frobenius norm, and eigenanalysis in a constructed semantic measurement space. SSRN.

Compares two collections of texts (e.g. the responses of two language models to the same prompts,
or two corpora) on a primary semantic axis and up to four secondary axes, each defined by two sets
of exemplar texts (poles). Reports the six analytical layers of the framework.
"""
import glob
import hashlib
import json
import os
import tempfile

import numpy as np
import pandas as pd

import geometry as geo

try:
    import gradio as gr
except ImportError:          # the analysis can be tested without gradio
    gr = None

try:
    import spaces             # Hugging Face ZeroGPU
    GPU = spaces.GPU
except ImportError:
    def GPU(fn=None, **kwargs):
        return fn if fn is not None else (lambda f: f)

APP_NAME = "DiscourseGeometry"
APP_URL = "https://huggingface.co/spaces/neovalle/DiscourseGeometry"
PAPER_URL = "https://papers.ssrn.com/sol3/papers.cfm?abstract_id=6387958"
ENCODERS = {
    "all-mpnet-base-v2 (paper default)": ("sentence-transformers/all-mpnet-base-v2", ""),
    "bge-base-en-v1.5": ("BAAI/bge-base-en-v1.5", ""),
    "e5-base-v2": ("intfloat/e5-base-v2", "query: "),
    "paraphrase-multilingual-mpnet-base-v2 (multilingual)": ("sentence-transformers/paraphrase-multilingual-mpnet-base-v2", ""),
}
DEFAULT_ENCODER = next(iter(ENCODERS))
N_SEC = 5
PRESET_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "presets")

# ------------------------------------------------------------------------------------------ encoders
_encoders, _cache = {}, {}


def get_encoder(model_id):
    if model_id not in _encoders:
        from sentence_transformers import SentenceTransformer
        _encoders[model_id] = SentenceTransformer(model_id)
    return _encoders[model_id]


@GPU(duration=120)
def _encode_gpu(model_id, prefix, texts):
    import torch
    m = get_encoder(model_id)
    if torch.cuda.is_available():
        m = m.to("cuda")
    return m.encode([prefix + t for t in texts], batch_size=32, normalize_embeddings=True,
                    convert_to_numpy=True, show_progress_bar=False)


def make_embed(encoder_label):
    """Returns embed(texts) with an in-memory cache, so poles are encoded once per encoder."""
    model_id, prefix = ENCODERS[encoder_label]

    def embed(texts):
        texts = list(texts)
        keys = [hashlib.md5((model_id + "\x00" + t).encode()).hexdigest() for t in texts]
        todo = [t for t, k in zip(texts, keys) if k not in _cache]
        if todo:
            E = np.asarray(_encode_gpu(model_id, prefix, todo), dtype=np.float64)
            for t, e in zip(todo, E):
                _cache[hashlib.md5((model_id + "\x00" + t).encode()).hexdigest()] = e
        if len(_cache) > 50000:
            _cache.clear()
        return np.array([_cache[k] for k in keys]) if texts else np.zeros((0, 768))
    return embed


# ------------------------------------------------------------------------------------------ generation
@GPU(duration=300)
def _generate_gpu(model_id, prompts, max_new_tokens):
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(model_id)
    tok.padding_side = "left"
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        model_id, torch_dtype=torch.float16 if dev == "cuda" else torch.float32).to(dev).eval()
    chat = bool(getattr(tok, "chat_template", None))
    fmt = (lambda p: tok.apply_chat_template([{"role": "user", "content": p}], tokenize=False,
                                             add_generation_prompt=True)) if chat else (lambda p: p)
    out = []
    for i in range(0, len(prompts), 16):
        batch = [fmt(p) for p in prompts[i:i + 16]]
        enc = tok(batch, return_tensors="pt", padding=True, add_special_tokens=not chat).to(dev)
        with torch.no_grad():
            g = model.generate(**enc, max_new_tokens=int(max_new_tokens), do_sample=False, pad_token_id=tok.pad_token_id)
        out += [tok.decode(r[enc["input_ids"].shape[1]:], skip_special_tokens=True).strip() for r in g]
    return out


def generate(prompts_text, model_a, model_b, max_new_tokens):
    prompts = geo.split_units(prompts_text, "line")
    if not prompts:
        raise gr.Error("Enter at least one prompt (one per line).")
    ra = _generate_gpu(model_a.strip(), prompts, max_new_tokens)
    rb = _generate_gpu(model_b.strip(), prompts, max_new_tokens)
    join = lambda rs: "\n---\n".join(r.replace("\n---\n", "\n") for r in rs)
    return join(ra), join(rb), "Paired", model_a.split("/")[-1], model_b.split("/")[-1]


# ------------------------------------------------------------------------------------------ presets
def list_presets():
    found = []
    for f in sorted(glob.glob(os.path.join(PRESET_DIR, "*.json"))):
        try:
            p = json.load(open(f, encoding="utf8"))
            found.append((not p.get("default", False), p["title"], f))
        except Exception as e:      # a broken preset must not stop the app
            print("preset skipped:", f, e)
    return {t: f for _, t, f in sorted(found)}     # the preset marked "default" comes first


PRESETS = list_presets()


def load_preset(title):
    """Returns values for every preset-controlled input (see PRESET_FIELDS order in build_ui)."""
    empty_axis = ["", "Pole A", "Pole B", "", ""]
    if not title or title not in PRESETS:
        return ["", "Pole A", "Pole B", "", ""] + empty_axis * N_SEC + ["Independent", "Collection A", "", "Collection B",
                                                                      "", "", "", "", ""]
    p = json.load(open(PRESETS[title], encoding="utf8"))
    vals = []
    axes = p.get("axes", [])
    for i in range(1 + N_SEC):
        if i < len(axes):
            a = axes[i]
            vals += [a["name"], a.get("pos_label", "Pole A"), a.get("neg_label", "Pole B"),
                     "\n---\n".join(a["pos"]), "\n---\n".join(a["neg"])]
        else:
            vals += empty_axis
    c = p.get("collections", {})
    vals += ["Paired" if c.get("paired") else "Independent",
             c.get("a_name", "Collection A"), "\n---\n".join(c.get("a", [])),
             c.get("b_name", "Collection B"), "\n---\n".join(c.get("b", [])),
             "\n---\n".join(c.get("baseline", [])),
             "\n".join(p.get("prompts", [])),
             p.get("models", {}).get("a", ""), p.get("models", {}).get("b", "")]
    return vals


# ------------------------------------------------------------------------------------------ reporting
def f4(x):
    return "n/a" if x is None or x != x else f"{x:+.4f}" if abs(x) < 10 else f"{x:+.2f}"


def fp(p):
    return "< .0001" if p < 1e-4 else f"{p:.4f}"


def table(rows, header):
    out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)]
    out += ["| " + " | ".join(str(c) for c in r) + " |" for r in rows]
    return "\n".join(out)


def format_report(G, R, name_a, name_b, encoder_label):
    P, secs = G["primary"], G["secondary"]
    paired = R["paired"]
    L1 = R["L1"]
    test = "paired t-test; d = Cohen's d_z" if paired else "Welch t-test; d = Cohen's d (pooled SD)"
    md = [f"### Comparison: **{name_b}** relative to **{name_a}**",
          f"{'Paired' if paired else 'Independent'} design · {R['n_a']} vs {R['n_b']} units · encoder: {encoder_label} · "
          f"{R['n_perm']} permutations, {R['n_boot']} bootstrap resamples. Scores are dot products with each axis "
          f"(positive = toward pole A); positions are on a 0–100% scale (0% = pole B centre, 100% = pole A centre).\n"]

    # validation
    md.append("#### Axis validation")
    rows = [[a["name"], f"{a['pos_label']} / {a['neg_label']}", f"{a['n_pos']} / {a['n_neg']}", f"{a['separation']:.4f}",
             f"{a['snr']:.1f}", f"{a['dprime']:.1f}", len(a["stray"])] for a in G["axes"]]
    md.append(table(rows, ["axis", "pole A / pole B", "exemplars", "separation", "separation ÷ mean SD", "d′",
                           "exemplars on wrong side (leave-one-out)"]))
    if G["angles"]:
        md.append("\nAngles between axes before orthogonalisation (90° = independent):\n")
        md.append(table([[f"{x['a']} ↔ {x['b']}", f"{x['angle']:.1f}°", f"{x['cosine']:.3f}"] for x in G["angles"]],
                        ["axes", "angle", "cosine"]))
    if secs:
        md.append(f"\nPrimary–secondary overlap ‖V_sec·v_primary‖ = {G['overlap']:.3f}: "
                  f"{100 * G['overlap'] ** 2:.0f}% of the primary direction lies in the secondary subspace.")
    strays = [s for a in G["axes"] for s in a["stray"]]
    if strays:
        md.append("\nExemplars closer to the opposite pole when left out (consider revising): "
                  + "; ".join(f"*{lab}*: “{t[:80]}…”" for lab, t in strays[:8]))

    # L1
    md.append(f"\n#### Layer 1 · Primary axis: {P['name']}")
    rows = [[name_a, f"{L1['mean_a']:+.4f}", f"{L1['pos_a']:.1f}%", f"{L1['sd_a']:.4f}"],
            [name_b, f"{L1['mean_b']:+.4f}", f"{L1['pos_b']:.1f}%", f"{L1['sd_b']:.4f}"]]
    if "baseline" in R:
        rows.append([f"Neutral baseline ({R['baseline']['n']})", f"{R['baseline']['score']:+.4f}",
                     f"{R['baseline']['position']:.1f}%", ""])
    md.append(table(rows, ["collection", "mean score", "position", "SD of score"]))
    md.append(f"\nShift {f4(L1['delta'])} toward *{P['pos_label'] if L1['delta'] > 0 else P['neg_label']}* · "
              f"t = {L1['t']:.2f}, p = {fp(L1['p'])} ({test}) · d = {L1['d']:.3f} ({geo.size_label(L1['d'])}) · "
              f"permutation p = {fp(L1['p_perm'])}")

    # L2–3
    if R["L2"]:
        md.append(f"\n#### Layer 2 · Secondary axes (Gram-Schmidt order: {', '.join(a['name'] for a in secs)})")
        rows = [[x["name"], f4(x["gs"]["delta"]), f"{x['gs']['t']:.2f}", fp(x["gs"]["p"]), f"{x['gs']['d']:.2f}",
                 "yes" if x["significant"] else "no"] for x in R["L2"]]
        md.append(table(rows, ["axis", "shift", "t", "p", "d", f"significant at Bonferroni α = {R['bonferroni_alpha']:.4f}"]))
        md.append("\n#### Layer 3 · Order check (Gram-Schmidt vs raw projections)")
        rows = [[x["name"], f4(x["gs"]["delta"]), f4(x["raw_delta"]),
                 "n/a" if x["ratio"] == float("inf") else f"{x['ratio']:.2f}×", x["flag"]] for x in R["L2"]]
        md.append(table(rows, ["axis", "GS shift", "raw shift", "GS ÷ raw", "reading"]))

    # L4–5
    L5 = R["L5"]
    if paired:
        if "L4" in R:
            md.append("\n#### Layer 4 · Shift within the secondary dimensions")
            md.append(table([[d["name"], f"{d['col_norm']:.4f}", f"{100 * d['share']:.1f}%", f"{d['raw_norm']:.4f}"]
                             for d in R["L4"]["dims"]], ["axis", "column norm (GS)", "share of secondary shift", "raw norm"]))
            md.append(f"\nFrobenius norm in secondary score space: {R['L4']['fro']:.4f}")
        md.append(f"\n#### Layer 5 · Total shift in the full {R['dim']}-dimensional space")
        md.append(f"‖ΔE‖_F = {L5['fro']:.4f} · per-unit shift mean {L5['per_mean']:.4f}, SD {L5['per_sd']:.4f}, "
                  f"mean ÷ SD = {L5['consistency']:.2f} (higher = more uniform shift across units)\n")
        rows = [["Primary axis", f"{L5['primary_norm']:.4f}", f"{100 * L5['primary_norm'] / L5['fro']:.1f}%",
                 f"{100 * L5['primary_frac_var']:.2f}%", f"{L5['primary_frac_var'] / L5['random_1']:.1f}×"]]
        if "sec_norm" in L5:
            rows.append(["Secondary subspace", f"{L5['sec_norm']:.4f}", f"{100 * L5['sec_norm'] / L5['fro']:.1f}%",
                         f"{100 * L5['sec_frac_var']:.2f}%", f"{L5['sec_frac_var'] / L5['random_sec']:.1f}×"])
        rows.append(["All named axes (GS, primary first)", f"{L5['full_norm']:.4f}", f"{100 * L5['full_norm'] / L5['fro']:.1f}%",
                     f"{100 * L5['full_frac_var']:.2f}%", f"{L5['full_frac_var'] / L5['random_full']:.1f}×"])
        rows.append(["Residual (unnamed directions)", f"{L5['residual_norm']:.4f}",
                     f"{100 * L5['residual_norm'] / L5['fro']:.1f}%", f"{100 * (1 - L5['full_frac_var']):.2f}%", ""])
        md.append(table(rows, ["subspace", "norm", "share of norm", "share of variance", "vs random directions"]))
        md.append("\nShares of norm do not add up; shares of variance do (Pythagorean decomposition).")
    else:
        md.append("\n#### Layers 4–5 · Distance between the collections (independent design)")
        md.append(f"Centroid separation ‖μ_B − μ_A‖ = {L5['centroid_sep']:.4f} · shape difference "
                  f"‖Σ_B − Σ_A‖_F relative to their mean size = {L5['shape_diff']:.2f}\n")
        rows = [["Primary axis", f"{100 * L5['primary_frac_var']:.2f}%", f"{L5['primary_frac_var'] / L5['random_1']:.1f}×"]]
        if "sec_frac_var" in L5:
            rows.append(["Secondary subspace", f"{100 * L5['sec_frac_var']:.2f}%", f"{L5['sec_frac_var'] / L5['random_sec']:.1f}×"])
        rows.append(["All named axes", f"{100 * L5['full_frac_var']:.2f}%", f"{L5['full_frac_var'] / L5['random_full']:.1f}×"])
        md.append(table(rows, ["subspace", "share of squared centroid separation", "vs random directions"]))
        if "L4" in R:
            md.append("\n" + table([[d["name"], f"{d['component']:+.4f}", f"{100 * d['share']:.1f}%"] for d in R["L4"]["dims"]],
                                   ["secondary axis", "component of centroid difference", "share within secondary"]))

    # L6
    L6 = R["L6"]
    md.append("\n#### Layer 6 · Shape and orientation of each collection")
    if "sec" in L6:
        s = L6["sec"]
        names = [a["name"] for a in secs]
        rows = []
        for lab, d in [(name_a, s["a"]), (name_b, s["b"])]:
            rows.append([lab, " / ".join(f"{e:.5f}" for e in d["eigenvalues"]), f"{d['total_var']:.5f}",
                         f"{d['isotropy']:.3f}", f"{d['eff_dim']:.2f}",
                         "[" + ", ".join(f"{x:+.2f}" for x in d["pc1"]) + "]", f"{d['pc1_to_primary']:.1f}°"])
        md.append(table(rows, ["collection", "eigenvalues", "total variance", "isotropy (λmin/λmax)",
                               "effective dimensionality", f"PC1 loadings ({', '.join(names)})", "PC1 angle to primary axis"]))
        md.append(f"\nPrincipal-axis rotation (PC1 → PC1, secondary space): **{s['rotation']:.2f}°**, permutation "
                  f"p = {fp(s['rotation_p'])}, 95% bootstrap interval {s['rotation_ci'][0]:.1f}°–{s['rotation_ci'][1]:.1f}°.")
        md.append(f"Change in PC1 angle to the primary axis: {s['b']['pc1_to_primary'] - s['a']['pc1_to_primary']:+.1f}° "
                  f"(95% bootstrap interval {s['pc1_to_primary_change_ci'][0]:+.1f}° to {s['pc1_to_primary_change_ci'][1]:+.1f}°). "
                  f"The primary axis is represented inside the secondary space by its projection "
                  f"({100 * s['primary_in_sec'] ** 2:.0f}% of it lies there).")
        md.append("Eigenvector signs are arbitrary: loadings are shown with the largest one positive, and only angles "
                  "(0–90°) are interpreted, not the direction in which PC1 points.")
    fu = L6["full"]
    md.append(f"\nFull space: PC1 explains {100 * fu['pc1_share_a']:.1f}% ({name_a}) and {100 * fu['pc1_share_b']:.1f}% "
              f"({name_b}) of variance; PC1 angle to primary axis {fu['pc1_to_primary_a']:.1f}° → {fu['pc1_to_primary_b']:.1f}°; "
              f"rotation {fu['rotation']:.1f}°, permutation p = {fp(fu['rotation_p'])}.")
    md.append("\n<sub>Permutation tests: " + ("responses to the same prompt are swapped between collections for a random "
              "half of the prompts." if paired else "units are randomly reassigned between collections.")
              + " p is the share of permutations giving a result at least as extreme as the observed one.</sub>")
    return "\n".join(md)


# ------------------------------------------------------------------------------------------ plots
COL_A, COL_B, COL_BASE = "#7a7a7a", "#2a7f62", "#b08a1e"


def wrap(s, n=70):
    words, lines, cur = str(s).split(), [], ""
    for w in words:
        if len(cur) + len(w) > n:
            lines.append(cur); cur = w
        else:
            cur = (cur + " " + w).strip()
    lines.append(cur)
    return "<br>".join(lines[:8]) + ("…" if len(lines) > 8 else "")


def plot_primary(G, R, ua, ub, name_a, name_b, base_pos=None):
    import plotly.graph_objects as go
    fig = go.Figure()
    for name, pos, units, col in [(name_a, R["units_a"]["position"], ua, COL_A), (name_b, R["units_b"]["position"], ub, COL_B)]:
        fig.add_trace(go.Violin(x=pos, name=f"{name} ({len(pos)})", orientation="h", side="positive", points="all",
                                pointpos=-0.4, jitter=0.3, marker=dict(size=4, opacity=0.6), line_color=col,
                                meanline_visible=True, text=[wrap(u) for u in units],
                                hovertemplate="%{x:.1f}%<br>%{text}<extra></extra>"))
    P = G["primary"]
    for x, lab in [(0, P["neg_label"]), (100, P["pos_label"])]:
        fig.add_vline(x=x, line_color="#999", annotation_text=lab, annotation_position="top")
    if base_pos is not None:
        fig.add_vline(x=base_pos, line_dash="dash", line_color=COL_BASE, annotation_text="neutral baseline",
                      annotation_position="bottom")
    fig.update_layout(title=f"Primary axis: {P['name']}", xaxis_title=f"position (%)   ·   0 = {P['neg_label']}, 100 = {P['pos_label']}",
                      height=380, showlegend=False, margin=dict(l=10, r=10, t=60, b=40))
    return fig


def plot_secondary(G, R, name_a, name_b):
    import plotly.graph_objects as go
    if not R["L2"]:
        return None
    names = [x["name"] for x in R["L2"]]
    fig = go.Figure()
    fig.add_trace(go.Bar(x=names, y=[x["gs"]["mean_a"] for x in R["L2"]], name=name_a, marker_color=COL_A))
    fig.add_trace(go.Bar(x=names, y=[x["gs"]["mean_b"] for x in R["L2"]], name=name_b, marker_color=COL_B))
    for i, x in enumerate(R["L2"]):
        fig.add_annotation(x=names[i], y=max(x["gs"]["mean_a"], x["gs"]["mean_b"]), yshift=14, showarrow=False,
                           text=f"Δ {x['gs']['delta']:+.4f}{' *' if x['significant'] else ''}")
    fig.update_layout(barmode="group", title="Secondary axes: mean projected score (Gram-Schmidt); * = significant after Bonferroni",
                      yaxis_title="score (positive = toward pole A)", height=380, margin=dict(l=10, r=10, t=60, b=40))
    return fig


def default_view(G):
    """Default 3D view: the first three secondary axes, orthogonalised (the paper's secondary space)."""
    names = [a["name"] for a in G["axes"]]
    if len(G["secondary"]) >= 3:
        return names[1], names[2], names[3], True
    pick = (names + [None, None])[:3]
    return pick[0], pick[1], pick[2], True


VIEW_MAP = "Semantic map: x = one axis, y and z = main remaining dimensions"
VIEW_AXES = "Chosen axes: any three axes"


def plot_map(state, xn, show_poles=True):
    """
    The semantic map of SpatialSemanticAnalysis: x = position on the chosen axis (true scale, 0% = pole B centre,
    100% = pole A centre); y, z = the two main directions of variation of all points once the axis direction is
    removed. Pole exemplars and pole centres are shown with the two collections.
    """
    import plotly.graph_objects as go
    G = state["G"]
    axes = {a["name"]: a for a in G["axes"]}
    ax = axes.get(xn) or G["primary"]
    v = ax["v"]
    blocks = []
    if show_poles:
        blocks += [(f"Pole A · {ax['pos_label']}", ax["Ep"], ax["pos_texts"], "#2a78d6", "diamond", 6),
                   (f"Pole B · {ax['neg_label']}", ax["En"], ax["neg_texts"], "#eb6834", "diamond", 6)]
    blocks += [(state["name_a"], state["EA"], state["ua"], COL_A, "circle", 3.5),
               (state["name_b"], state["EB"], state["ub"], COL_B, "circle", 3.5)]
    allE = np.vstack([b[1] for b in blocks])
    resid = allE - np.outer(allE @ v, v)
    resid = resid - resid.mean(0)
    _, sv, Vt = np.linalg.svd(resid, full_matrices=False)
    share = sv[:2] ** 2 / (sv ** 2).sum() if (sv ** 2).sum() > 0 else np.zeros(2)
    fig = go.Figure()
    start = 0
    for name, E, texts, col, sym, size in blocks:
        n = len(E)
        x = geo.position(E, ax)
        yz = resid[start:start + n] @ Vt[:2].T
        start += n
        fig.add_trace(go.Scatter3d(x=x, y=yz[:, 0], z=yz[:, 1], mode="markers", name=name,
                                   text=[wrap(t) for t in texts],
                                   hovertemplate="%{x:.0f}%<br>%{text}<extra>" + name + "</extra>",
                                   marker=dict(size=size, color=col, symbol=sym, opacity=0.85 if sym == "diamond" else 0.6)))
        m = np.array([x.mean(), yz[:, 0].mean(), yz[:, 1].mean()])
        fig.add_trace(go.Scatter3d(x=[m[0]], y=[m[1]], z=[m[2]], mode="markers", name=f"centre · {name}",
                                   hoverinfo="name", showlegend=False,
                                   marker=dict(size=9, color=col, symbol="diamond-open" if sym == "diamond" else "diamond",
                                               line=dict(color="black", width=1))))
    fig.update_layout(title=f"Semantic map · {ax['name']}: 0% = {ax['neg_label']}, 100% = {ax['pos_label']}",
                      scene=dict(xaxis_title=f"{ax['name']} position (%)",
                                 yaxis_title=f"other dimension 1 ({100 * share[0]:.0f}%)",
                                 zaxis_title=f"other dimension 2 ({100 * share[1]:.0f}%)"),
                      height=640, margin=dict(l=0, r=0, t=60, b=0), legend=dict(itemsizing="constant"))
    return fig


def draw_view(state, view, xn, yn, zn, ortho):
    if not state:
        return None
    if view == VIEW_AXES:
        return plot_3d(state, xn, yn, zn, ortho)
    return plot_map(state, xn)


def plot_3d(state, xn, yn, zn, ortho):
    """Collections, their centroids and PC1 in the space spanned by any three chosen axes."""
    import plotly.graph_objects as go
    if not state:
        return None
    G = state["G"]
    axes = {a["name"]: a for a in G["axes"]}
    chosen = [n for n in (xn, yn, zn) if n]
    if len(set(chosen)) != len(chosen) or any(n not in axes for n in chosen):
        raise gr.Error("Choose different axes for x, y and z.")
    if len(chosen) < 2:
        return None
    V = np.array([axes[n]["v"] for n in chosen])
    if ortho:
        V = geo.gram_schmidt(list(V))
    three = len(chosen) == 3
    fig = go.Figure()
    pcs = []
    for name, E, units, col in [(state["name_a"], state["EA"], state["ua"], COL_A),
                                (state["name_b"], state["EB"], state["ub"], COL_B)]:
        C = E @ V.T
        kw = dict(mode="markers", name=name, text=[wrap(u) for u in units],
                  hovertemplate="%{text}<extra>" + name + "</extra>",
                  marker=dict(size=3.5 if three else 6, color=col, opacity=0.6))
        m = C.mean(0)
        w, vec = geo.cov_eig(C)
        pc1 = vec[:, 0]
        pcs.append(pc1)
        half = 2 * np.sqrt(max(w[0], 0))
        ends = np.array([m - half * pc1, m + half * pc1])
        line = dict(mode="lines", name=f"PC1 · {name}", line=dict(color=col, width=6 if three else 3), hoverinfo="skip")
        cen = dict(mode="markers", name=f"centroid · {name}", hoverinfo="name",
                   marker=dict(size=9 if three else 12, color=col, symbol="diamond", line=dict(color="black", width=1)))
        if three:
            fig.add_trace(go.Scatter3d(x=C[:, 0], y=C[:, 1], z=C[:, 2], **kw))
            fig.add_trace(go.Scatter3d(x=ends[:, 0], y=ends[:, 1], z=ends[:, 2], **line))
            fig.add_trace(go.Scatter3d(x=[m[0]], y=[m[1]], z=[m[2]], **cen))
        else:
            fig.add_trace(go.Scatter(x=C[:, 0], y=C[:, 1], **kw))
            fig.add_trace(go.Scatter(x=ends[:, 0], y=ends[:, 1], **line))
            fig.add_trace(go.Scatter(x=[m[0]], y=[m[1]], **cen))
    rot = geo.angle(pcs[0], pcs[1])
    if ortho:
        how = "orthogonalised in the order " + " → ".join(chosen)
    else:
        angs = ", ".join(f"{geo.angle(axes[a]['v'], axes[b]['v']):.0f}°" for k, a in enumerate(chosen) for b in chosen[k + 1:])
        how = f"raw projections; angles between the axes {angs} (90° = independent)"
    title = f"3D view ({how}) · PC1 rotation {rot:.1f}°"
    lab = [n + (" (GS)" if ortho and k else "") for k, n in enumerate(chosen)]
    if three:
        fig.update_layout(title=title, height=620, margin=dict(l=0, r=0, t=60, b=0),
                          scene=dict(xaxis_title=lab[0], yaxis_title=lab[1], zaxis_title=lab[2]))
    else:
        fig.update_layout(title=title, height=480, xaxis_title=lab[0], yaxis_title=lab[1])
    return fig


def plot_shift(R, ua):
    import plotly.graph_objects as go
    if not R["paired"]:
        return None
    per = R["L5"]["per_unit"]
    fig = go.Figure(go.Histogram(x=per, nbinsx=20, marker_color=COL_B))
    fig.add_vline(x=R["L5"]["per_mean"], line_dash="dash", annotation_text="mean")
    fig.update_layout(title="Per-unit shift ‖e_B − e_A‖ in the full space", xaxis_title="shift", yaxis_title="units",
                      height=320, margin=dict(l=10, r=10, t=50, b=40))
    return fig


# ------------------------------------------------------------------------------------------ pipeline
def collect_specs(sep, axis_values):
    specs = []
    for i in range(1 + N_SEC):
        name, la, lb, ta, tb = axis_values[5 * i:5 * i + 5]
        pos, neg = geo.split_units(ta, sep), geo.split_units(tb, sep)
        if i == 0 and (not pos or not neg):
            raise ValueError("The primary axis needs exemplar texts for both poles.")
        if pos or neg:
            specs.append(dict(name=(name or f"Axis {i + 1}").strip(), pos_label=(la or "Pole A").strip(),
                              neg_label=(lb or "Pole B").strip(), pos=pos, neg=neg))
    return specs


def analyse_all(sep, axis_values, mode, name_a, text_a, name_b, text_b, baseline, unit_mode,
                encoder_label, n_perm, n_boot, seed, embed=None):
    specs = collect_specs(sep, axis_values)
    split = (lambda t: geo.split_sentences(t)) if unit_mode == "Split into sentences" else (lambda t: geo.split_units(t, sep))
    ua, ub = split(text_a), split(text_b)
    paired = mode == "Paired"
    if paired and unit_mode == "Split into sentences":
        raise ValueError("Paired mode compares whole units matched by position; use 'As separated'.")
    embed = embed or make_embed(encoder_label)
    G = geo.build_axes(specs, embed)
    EA, EB = embed(ua), embed(ub)
    bl = split(baseline) if baseline and baseline.strip() else []
    R = geo.analyse(G, EA, EB, paired, n_perm=int(n_perm), n_boot=int(n_boot), seed=int(seed),
                    E_base=embed(bl) if bl else None)
    R["_EA"], R["_EB"] = EA, EB
    return G, R, ua, ub


def units_table(G, R, ua, ub, name_a, name_b):
    rows = []
    sec_names = [a["name"] for a in G["secondary"]]
    for name, units, U in [(name_a, ua, R["units_a"]), (name_b, ub, R["units_b"])]:
        for i, u in enumerate(units):
            r = {"collection": name, "unit": i + 1, "position_%": round(U["position"][i], 1),
                 "primary_score": round(U["score"][i], 5)}
            for j, s in enumerate(sec_names):
                r[f"{s} (GS)"] = round(U["sec"][i][j], 5)
            if R["paired"] and name == name_b:
                r["shift_norm"] = round(R["L5"]["per_unit"][i], 4)
            r["text"] = u
            rows.append(r)
    return pd.DataFrame(rows)


def _jsonable(x):
    if isinstance(x, dict):
        return {k: _jsonable(v) for k, v in x.items() if k not in ("units_a", "units_b", "_EA", "_EB")}
    if isinstance(x, (list, tuple)):
        return [_jsonable(v) for v in x]
    if isinstance(x, (np.floating, np.integer)):
        return x.item()
    if isinstance(x, np.ndarray):
        return x.tolist()
    return x


def run(sep, *args):
    axis_values = list(args[:5 * (1 + N_SEC)])
    (mode, name_a, text_a, name_b, text_b, baseline, unit_mode, encoder_label, n_perm, n_boot, seed) = args[5 * (1 + N_SEC):]
    try:
        G, R, ua, ub = analyse_all(sep, axis_values, mode, name_a or "A", text_a, name_b or "B", text_b, baseline,
                                   unit_mode, encoder_label, n_perm, n_boot, seed)
    except ValueError as e:
        raise gr.Error(str(e))
    name_a, name_b = name_a or "A", name_b or "B"
    report = format_report(G, R, name_a, name_b, encoder_label)
    df = units_table(G, R, ua, ub, name_a, name_b)
    tmp = tempfile.mkdtemp()
    csv_path = os.path.join(tmp, "units.csv")
    df.to_csv(csv_path, index=False)
    js_path = os.path.join(tmp, "results.json")
    G_out = {k: v for k, v in G.items() if k not in ("sec_basis", "full_basis")}
    G_out["axes"] = [{k: v for k, v in a.items() if k not in ("cp", "cn", "v", "Ep", "En", "pos_texts", "neg_texts")}
                     for a in G["axes"]]
    G_out.pop("primary", None); G_out.pop("secondary", None)
    json.dump(_jsonable(dict(encoder=encoder_label, axes=G_out, results=R)), open(js_path, "w"), indent=1)
    base_pos = R["baseline"]["position"] if "baseline" in R else None
    state = dict(G=G, EA=np.array(R.pop("_EA")), EB=np.array(R.pop("_EB")), ua=ua, ub=ub, name_a=name_a, name_b=name_b)
    names = [a["name"] for a in G["axes"]]
    xv, yv, zv, ov = default_view(G)
    dd = lambda v: gr.update(choices=names, value=v) if gr is not None else v
    return (report, plot_primary(G, R, ua, ub, name_a, name_b, base_pos), plot_secondary(G, R, name_a, name_b),
            plot_map(state, names[0]), plot_shift(R, ua), df.drop(columns=["text"]).head(500),
            [csv_path, js_path], state, dd(names[0]), dd(yv), dd(zv), ov, VIEW_MAP)


# ------------------------------------------------------------------------------------------ UI
INTRO = f"""
# {APP_NAME}
Compare two collections of texts on theoretically defined semantic axes, using the spatial-geometric framework of
[Vallego (2026)]({PAPER_URL}): **centroid projection** (where does each collection sit, and did it move?),
**Frobenius norm** (how large is the difference, and how much of it do the named axes capture?) and
**eigenanalysis** (did the shape and orientation of the collection change?).

Each axis is defined by two sets of exemplar texts (pole A and pole B). The first axis is the primary axis; further
secondary axes (up to five) decompose it and are orthogonalised in the order given. Use **Paired** for two sets of responses to the
same prompts in the same order (e.g. a base model and its fine-tuned version), and **Independent** for two corpora.
Separate units with a line containing only `---`, with blank lines, or one per line.
"""

CITE_MD = "Part of the [H4rmony Project](https://theh4rmonyproject.org) · **Cite this work** (expand below)"
CITE_TEXT = f"""Vallego, J. (2026). *A spatial-geometric framework for discourse comparison: Centroid projection, Frobenius
norm, and eigenanalysis in a constructed semantic measurement space*. SSRN. {PAPER_URL}

Vallego, J. (2026). *{APP_NAME}* [Software]. Hugging Face Spaces. {APP_URL}

```bibtex
@misc{{vallego2026framework,
  author = {{Vallego, Jorge}},
  title  = {{A Spatial-Geometric Framework for Discourse Comparison: Centroid Projection,
            Frobenius Norm, and Eigenanalysis in a Constructed Semantic Measurement Space}},
  year   = {{2026}},
  howpublished = {{SSRN}},
  url    = {{{PAPER_URL}}}
}}
@software{{vallego2026discoursegeometry,
  author    = {{Vallego, Jorge}},
  title     = {{{APP_NAME}}},
  year      = {{2026}},
  publisher = {{Hugging Face Spaces}},
  url       = {{{APP_URL}}}
}}
```
"""


def axis_inputs(label, open_=True):
    with gr.Row():
        name = gr.Textbox(label=f"{label}: name", scale=2)
        la = gr.Textbox(label="Pole A label (positive)", value="Pole A", scale=1)
        lb = gr.Textbox(label="Pole B label (negative)", value="Pole B", scale=1)
    with gr.Row():
        ta = gr.Textbox(label="Pole A exemplars", lines=8 if open_ else 5)
        tb = gr.Textbox(label="Pole B exemplars", lines=8 if open_ else 5)
    return [name, la, lb, ta, tb]


def build_ui():
    with gr.Blocks(title=APP_NAME) as demo:
        gr.Markdown(INTRO)
        gr.Markdown(CITE_MD)
        with gr.Accordion("Cite this work", open=False):
            gr.Markdown(CITE_TEXT)
        with gr.Row():
            preset = gr.Dropdown(choices=list(PRESETS), value=next(iter(PRESETS), None), label="Preset",
                                 info="Loads axes, exemplars and example collections", scale=3)
            load_btn = gr.Button("Load preset", scale=1)
        with gr.Tab("1 · Axes"):
            sep = gr.Dropdown(choices=[("auto", "auto"), ("line with ---", "dash"), ("blank lines", "blank"),
                                       ("one per line", "line")], value="auto", label="Units are separated by")
            axis_fields = axis_inputs("Primary axis")
            for i in range(N_SEC):
                with gr.Accordion(f"Secondary axis {i + 1}" + (" (extended)" if i >= 3 else ""), open=i < 3):
                    axis_fields += axis_inputs(f"Secondary axis {i + 1}", open_=False)
        with gr.Tab("2 · Collections"):
            mode = gr.Radio(["Paired", "Independent"], value="Independent", label="Design")
            with gr.Row():
                with gr.Column():
                    name_a = gr.Textbox(label="Collection A name (reference)", value="Collection A")
                    text_a = gr.Textbox(label="Collection A units", lines=12)
                with gr.Column():
                    name_b = gr.Textbox(label="Collection B name (compared with A)", value="Collection B")
                    text_b = gr.Textbox(label="Collection B units", lines=12)
            baseline = gr.Textbox(label="Neutral baseline (optional): factual, stance-free texts on the same subject", lines=3)
            unit_mode = gr.Radio(["As separated", "Split into sentences"], value="As separated",
                                 label="Unit of analysis (sentence splitting only in the independent design)")
            with gr.Accordion("Generate the two collections with language models (optional, uses GPU time)", open=False):
                prompts = gr.Textbox(label="Prompts (one per line)", lines=6)
                with gr.Row():
                    model_a = gr.Textbox(label="Model A (Hugging Face id)")
                    model_b = gr.Textbox(label="Model B (Hugging Face id)")
                    max_tok = gr.Slider(50, 400, value=300, step=10, label="Max new tokens (greedy decoding)")
                gen_btn = gr.Button("Generate responses into A and B")
        with gr.Tab("3 · Settings"):
            encoder = gr.Dropdown(choices=list(ENCODERS), value=DEFAULT_ENCODER, label="Sentence encoder",
                                  info="Re-run with other encoders to check that the findings do not depend on one model")
            with gr.Row():
                n_perm = gr.Slider(99, 4999, value=999, step=100, label="Permutations")
                n_boot = gr.Slider(99, 4999, value=999, step=100, label="Bootstrap resamples")
                seed = gr.Number(value=0, precision=0, label="Random seed")
        go_btn = gr.Button("Run analysis", variant="primary")
        report = gr.Markdown()
        p1 = gr.Plot(label="Primary axis")
        p2 = gr.Plot(label="Secondary axes")
        with gr.Group():
            p3 = gr.Plot(label="3D view")
            view = gr.Radio([VIEW_MAP, VIEW_AXES], value=VIEW_MAP, label="3D view",
                            info="Semantic map: the chosen x axis on its 0–100% scale, with its pole exemplars; "
                                 "y and z are the main directions in which all texts differ apart from that axis "
                                 "(as in SpatialSemanticAnalysis). Chosen axes: x, y and z are three axes of your choice.")
            with gr.Row():
                ax_x = gr.Dropdown(label="x axis", choices=[], scale=2)
                ax_y = gr.Dropdown(label="y axis (chosen axes only)", choices=[], scale=2)
                ax_z = gr.Dropdown(label="z axis (chosen axes only)", choices=[], scale=2)
                ortho = gr.Checkbox(label="Orthogonalise in the order x → y → z", value=True, scale=1)
                redraw = gr.Button("Redraw 3D view", scale=1)
        state = gr.State(None)
        p4 = gr.Plot(label="Per-unit shift")
        units_df = gr.Dataframe(label="Units (first 500; full table in the CSV)")
        files = gr.File(label="Downloads: per-unit CSV and all results (JSON)", file_count="multiple")

        preset_fields = axis_fields + [mode, name_a, text_a, name_b, text_b, baseline, prompts, model_a, model_b]
        load_btn.click(load_preset, inputs=preset, outputs=preset_fields)
        demo.load(load_preset, inputs=preset, outputs=preset_fields)
        gen_btn.click(generate, inputs=[prompts, model_a, model_b, max_tok], outputs=[text_a, text_b, mode, name_a, name_b])
        go_btn.click(run, inputs=[sep] + axis_fields + [mode, name_a, text_a, name_b, text_b, baseline, unit_mode,
                                                        encoder, n_perm, n_boot, seed],
                     outputs=[report, p1, p2, p3, p4, units_df, files, state, ax_x, ax_y, ax_z, ortho, view])
        redraw.click(draw_view, inputs=[state, view, ax_x, ax_y, ax_z, ortho], outputs=p3)
        view.change(draw_view, inputs=[state, view, ax_x, ax_y, ax_z, ortho], outputs=p3)
    return demo


if gr is not None:
    try:
        get_encoder(ENCODERS[DEFAULT_ENCODER][0])     # load once at start-up (recommended for ZeroGPU)
    except Exception as e:
        print("Encoder not preloaded:", e)
    demo = build_ui()

if __name__ == "__main__" and gr is not None:
    demo.launch()