"""Three-axis spatial-geometric profile of a single text.

Implements a simplified, single-text reading of Vallego (2026),
"A Spatial-Geometric Framework for Discourse Comparison".

Method
------
1. Each axis is defined by two poles, each a set of exemplar sentences.
   Its raw direction is the difference between the two pole centroids
   (high-pole centroid minus low-pole centroid) in embedding space.
2. The three raw directions overlap (they are not orthogonal). They are
   orthogonalised with symmetric (Loewdin) orthogonalisation, which finds
   the orthonormal set closest to the raw directions and, unlike
   Gram-Schmidt, treats all three axes equally (no order dependence).
3. A text is split into sentences; each sentence and the text centroid
   (mean of sentence embeddings) are projected onto the three orthonormal
   directions. These three numbers are the 3D coordinates.
4. Each coordinate is rescaled so that 0% = the low-pole centroid and
   100% = the high-pole centroid on that axis. Values below 0 or above 100
   lie further out than the typical exemplar of that pole.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

import numpy as np
import pandas as pd

AXES = [
    {"id": "eco", "low": "Anthropocentric", "high": "Ecocentric"},
    {"id": "change", "low": "Conservative", "high": "Transformative"},
    {"id": "outlook", "low": "Pessimistic", "high": "Optimistic"},
]


def split_sentences(text: str, min_words: int = 3) -> list[str]:
    parts = re.split(r"(?<=[.!?;])\s+|\n+", text.strip())
    sents = [p.strip() for p in parts if len(p.strip().split()) >= min_words]
    return sents or ([text.strip()] if text.strip() else [])


def _normalise(x: np.ndarray) -> np.ndarray:
    return x / np.linalg.norm(x, axis=-1, keepdims=True).clip(min=1e-12)


def loewdin(D: np.ndarray) -> np.ndarray:
    """Symmetric orthogonalisation of the rows of D: Q = (D D^T)^(-1/2) D."""
    S = D @ D.T
    vals, vecs = np.linalg.eigh(S)
    S_inv_sqrt = vecs @ np.diag(1.0 / np.sqrt(vals)) @ vecs.T
    return S_inv_sqrt @ D


@dataclass
class Space:
    axes: list[dict]
    Q: np.ndarray            # 3 x d, orthonormal axis directions
    raw_dirs: np.ndarray     # 3 x d, unit raw directions (before orthogonalisation)
    lo: np.ndarray           # 3, projection of each low-pole centroid on its own axis
    hi: np.ndarray           # 3, projection of each high-pole centroid on its own axis
    raw_lo: np.ndarray
    raw_hi: np.ndarray
    pole_points: dict        # (axis_id, pole) -> (sentences, coords n x 3 in %)
    pole_centroids: dict     # (axis_id, pole) -> coords (3,) in %
    overlap: pd.DataFrame    # cosine between raw directions

    def to_percent(self, proj: np.ndarray) -> np.ndarray:
        return 100.0 * (proj - self.lo) / (self.hi - self.lo)

    def coords(self, E: np.ndarray) -> np.ndarray:
        return self.to_percent(E @ self.Q.T)

    def raw_percent(self, E: np.ndarray) -> np.ndarray:
        return 100.0 * (E @ self.raw_dirs.T - self.raw_lo) / (self.raw_hi - self.raw_lo)


def build_space(poles: pd.DataFrame, encode, axes=AXES) -> Space:
    """poles: DataFrame with columns axis, pole ('low'/'high'), text.
    encode: function list[str] -> np.ndarray (n x d)."""
    emb, cents = {}, {}
    for ax in axes:
        for pole in ("low", "high"):
            texts = poles[(poles.axis == ax["id"]) & (poles.pole == pole)].text.tolist()
            if len(texts) < 2:
                raise ValueError(f"Axis '{ax['id']}' pole '{pole}' needs at least 2 exemplars.")
            E = _normalise(np.asarray(encode(texts), dtype=float))
            emb[(ax["id"], pole)] = (texts, E)
            cents[(ax["id"], pole)] = E.mean(axis=0)

    raw = np.stack([cents[(a["id"], "high")] - cents[(a["id"], "low")] for a in axes])
    raw_unit = _normalise(raw)
    Q = loewdin(raw)

    lo = np.array([cents[(a["id"], "low")] @ Q[k] for k, a in enumerate(axes)])
    hi = np.array([cents[(a["id"], "high")] @ Q[k] for k, a in enumerate(axes)])
    raw_lo = np.array([cents[(a["id"], "low")] @ raw_unit[k] for k, a in enumerate(axes)])
    raw_hi = np.array([cents[(a["id"], "high")] @ raw_unit[k] for k, a in enumerate(axes)])

    names = [f"{a['low']}–{a['high']}" for a in axes]
    overlap = pd.DataFrame(raw_unit @ raw_unit.T, index=names, columns=names).round(2)

    space = Space(axes, Q, raw_unit, lo, hi, raw_lo, raw_hi, {}, {}, overlap)
    for key, (texts, E) in emb.items():
        space.pole_points[key] = (texts, space.coords(E))
        space.pole_centroids[key] = space.coords(cents[key][None, :])[0]
    return space


def reading(pct: float, low: str, high: str) -> str:
    if pct < 10:
        return f"Strongly {low.lower()}"
    if pct < 40:
        return f"Leans {low.lower()}"
    if pct <= 60:
        return "Between the poles"
    if pct <= 90:
        return f"Leans {high.lower()}"
    return f"Strongly {high.lower()}"


def analyse(text: str, space: Space, encode) -> dict:
    sents = split_sentences(text)
    if not sents:
        raise ValueError("Please enter some text.")
    E = _normalise(np.asarray(encode(sents), dtype=float))
    centroid = E.mean(axis=0)
    sent_xyz = space.coords(E)
    text_xyz = space.coords(centroid[None, :])[0]
    raw_pct = space.raw_percent(centroid[None, :])[0]

    rows = []
    for k, ax in enumerate(space.axes):
        rows.append({
            "Axis (0% → 100%)": f"{ax['low']} → {ax['high']}",
            "Position": round(float(text_xyz[k]), 1),
            "Reading": reading(float(text_xyz[k]), ax["low"], ax["high"]),
            "Spread across sentences (±)": round(float(sent_xyz[:, k].std()), 1) if len(sents) > 1 else None,
            "Before orthogonalisation": round(float(raw_pct[k]), 1),
        })
    return {
        "sentences": sents,
        "sent_xyz": sent_xyz,
        "text_xyz": text_xyz,
        "table": pd.DataFrame(rows),
    }