src/jdm_agent/viz/subgraph.py

"""Construit un sous-graphe JDM autour d'un terme racine, sérialisable en
JSON ou en HTML interactif (vis-network).

Spécification : voir `subgraph_visualization_howto.md` à la racine du dépôt.
Fichier de référence (rendu cible) : `plat_asiatique_subgraph.html`.
"""
from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Literal, Optional

from jdm_agent.client import JDMClient
from jdm_agent.viz.template import HTML_TEMPLATE


# ---------- Paramètres par défaut (recette du howto) ----------

#: Jeu standard de relations explorées à la profondeur 1.
#: Sélection calibrée pour des sous-graphes lisibles (8 relations) — exclut
#: par défaut les inverses verbo-nominaux (patient-1, agent-1) et r_associated,
#: trop bruyants en exploration générale. Le LLM/utilisateur peut les
#: rajouter explicitement quand pertinent.
DEFAULT_RELATIONS: list[str] = [
    "r_isa", "r_hypo", "r_syn", "r_anto",
    "r_carac", "r_has_part", "r_lieu", "r_domain",
]

#: Sous-ensemble exploré à la profondeur 2 (limite l'explosion combinatoire).
#: Garde r_isa pour permettre la remontée catégorielle au 2nd niveau.
DEFAULT_DEPTH2_RELATIONS: list[str] = [
    "r_isa", "r_carac", "r_has_part", "r_lieu",
]

#: Profondeur 3 — encore plus étroit pour éviter le bruit.
DEFAULT_DEPTH3_RELATIONS: list[str] = [
    "r_isa", "r_has_part", "r_carac",
]

#: Profondeur 4 — uniquement la remontée catégorielle + caractérisation.
DEFAULT_DEPTH4_RELATIONS: list[str] = [
    "r_isa", "r_carac",
]

#: Opacité appliquée aux nœuds et aux arêtes par profondeur. Decroissance
#: linéaire bornée pour garder le niveau 4 encore visible.
OPACITY_BY_DEPTH: dict[int, float] = {0: 1.0, 1: 1.0, 2: 0.65, 3: 0.42, 4: 0.28}


# ---------- Mapping relation → couleur / "kind" CSS ----------

#: Type de nœud dérivé de la relation entrante (pour la palette de couleurs).
KIND_OF_REL: dict[str, str] = {
    "r_isa": "isa",
    "r_hypo": "hypo",
    "r_syn": "syn",
    "r_anto": "anto",
    "r_carac": "carac",
    "r_has_part": "part",
    "r_lieu": "lieu",
    "r_patient-1": "verb",
    "r_agent-1": "verb",
    "r_domain": "domain",
    "r_associated": "assoc",
}

#: Palette UNIQUE par kind. Fonds clairs (Material 100, ~10 % de teinte)
#: pour que les labels lisent facilement, bordures saturées (Material 700-800)
#: pour l'identité de la famille et la cohérence avec les arêtes (= même
#: couleur que la bordure). La profondeur n'influence plus la teinte :
#: seule l'opacité (OPACITY_BY_DEPTH) distingue les niveaux.
PALETTE: dict[str, dict[str, str]] = {
    "center": {"background": "#1a1a1a", "border": "#000"},
    "isa":    {"background": "#bbdefb", "border": "#1565c0"},  # bleu
    "hypo":   {"background": "#c8e6c9", "border": "#2e7d32"},  # vert
    "syn":    {"background": "#dcedc8", "border": "#558b2f"},  # vert-jaune
    "anto":   {"background": "#ffcdd2", "border": "#c62828"},  # rouge
    "carac":  {"background": "#e1bee7", "border": "#6a1b9a"},  # violet
    "part":   {"background": "#ffe0b2", "border": "#a04500"},  # terracotta
    "lieu":   {"background": "#b2dfdb", "border": "#00695c"},  # teal
    "verb":   {"background": "#f8bbd0", "border": "#ad1457"},  # rose
    "domain": {"background": "#d1c4e9", "border": "#4527a0"},  # indigo
    "assoc":  {"background": "#cfd8dc", "border": "#455a64"},  # blue-grey
    "d2":     {"background": "#e0e2e5", "border": "#9aa0a6"},  # fallback générique
}

#: Couleur d'arête par kind = bordure de la palette (assortie au nœud).
EDGE_COLOR: dict[str, str] = {
    k: v["border"] for k, v in PALETTE.items() if k != "center"
}


def _palette_for(kind: str, depth: int) -> dict[str, str]:
    """Renvoie la palette du kind. La profondeur n'altère plus la teinte
    (gérée via l'opacité du nœud) — on garde l'argument pour compat API."""
    _ = depth  # unused, kept for backward signature
    return PALETTE.get(kind, PALETTE["assoc"])


def _edge_color(kind: str, depth: int, negative: bool) -> str:
    """Couleur d'arête : rouge saturé pour négation, teinte de la famille
    sinon. La profondeur ne change pas la teinte — l'opacité s'en charge."""
    _ = depth  # unused
    if negative:
        return "#c62828"
    return EDGE_COLOR.get(kind, "#9aa0a6")


def _opacity_for(depth: int) -> float:
    return OPACITY_BY_DEPTH.get(depth, OPACITY_BY_DEPTH[4])


# ---------- Slug pour nom de fichier ----------

_SLUG_RE = re.compile(r"[^a-z0-9]+")


def _slugify(text: str) -> str:
    s = text.lower().strip()
    s = _SLUG_RE.sub("_", s)
    return s.strip("_") or "subgraph"


# ---------- Coeur : construction du graphe ----------

def _fetch_relation(
    client: JDMClient,
    term: str,
    rel: str,
    top_k: int,
    min_weight: Optional[float],
) -> list[dict[str, Any]]:
    """Récupère les triplets pour une relation donnée, depuis le terme `term`.

    Renvoie la liste tronquée à `top_k` (par |w| décroissant), avec
    `{target_display, target_id, w, polarity}`.
    """
    rid = client.relation_type_id(rel)
    if rid is None:
        return []
    try:
        result = client.relations_from(
            term, types_ids=[rid],
            min_weight=min_weight,
            limit=max(top_k * 2, 20),  # marge pour pouvoir filtrer ensuite
        )
    except Exception:
        return []

    idx = result.node_index()
    out: list[dict[str, Any]] = []
    for r in sorted(result.relations, key=lambda x: -abs(x.w)):
        node = idx.get(r.node2)
        if node is None:
            try:
                node = client.node_by_id(r.node2)
            except Exception:
                continue
        if node.type == 8:  # chunks (agrégats syntaxiques) — exclus
            continue
        dec = client.decode_node_name(node.name, local_nodes=idx)
        out.append({
            "target_display": dec["decoded"],
            "target_id": node.name,
            "w": r.w,
            "polarity": "négation" if r.w < 0 else "affirmation",
        })
        if len(out) >= top_k:
            break
    return out


def _build_node(
    node_id: str,
    label: str,
    kind: str,
    depth: int,
    fixed_center: bool = False,
) -> dict[str, Any]:
    """Construit un nœud vis-network. Le centre garde sa couleur dédiée ;
    les autres nœuds prennent une palette plus claire à partir du niveau 2
    (même teinte que le niveau 1, mais visiblement adoucie). L'opacité du
    nœud décroît avec la profondeur (cf. OPACITY_BY_DEPTH)."""
    if fixed_center:
        color = PALETTE["center"]
    else:
        color = _palette_for(kind, depth)
    opacity = _opacity_for(depth)
    # Taille de police : 28 au centre, 22 au niveau 1, 17 ensuite ;
    # +1 pour le niveau 2 pour ne pas trop pénaliser le 1er anneau d2.
    if fixed_center:
        font_size = 28
    elif depth >= 3:
        font_size = 16
    elif depth == 2:
        font_size = 18
    else:
        font_size = 22
    node: dict[str, Any] = {
        "id": node_id,
        "label": label,
        "color": color,
        "opacity": opacity,
        "shape": "ellipse" if fixed_center else "box",
        "font": {
            "size": font_size,
            "color": "#fff" if fixed_center else "#222",
            "face": "system-ui",
        },
        "margin": 12,
        "widthConstraint": {"maximum": 220},
        "_depth": depth,
        "_kind": kind,
    }
    if fixed_center:
        node["x"] = 0
        node["y"] = 0
        node["fixed"] = {"x": True, "y": True}
        node["mass"] = 5
    return node


def _build_edge(
    from_id: str,
    to_id: str,
    relation: str,
    w: float,
    polarity: str,
    depth: int,
    from_negation: bool = False,
) -> dict[str, Any]:
    """Construit une arête vis-network.

    - Couleur du trait ET de l'étiquette = teinte de la famille de relation
      (r_isa → bleu, r_hypo → vert, ...), avec une variante plus claire au
      niveau 2.
    - Opacité dépend du niveau : 1.0 au niveau 1, 0.55 au niveau 2 — d'où
      un effet "voisinage immédiat saturé / extensions estompées".
    - Négation = rouge (saturé au niveau 1, pâle au niveau 2), label
      préfixé « NON ».
    - Niveau 2 = pointillés (en plus de la teinte adoucie).
    - `from_negation` : l'arête part d'un nœud lui-même introduit par une
      relation négative. Côté HTML, le bouton « Approfondir négations »
      masque/affiche ces arêtes (OFF par défaut → nœud non approfondi).
    """
    is_neg = polarity == "négation"
    is_d2 = depth >= 2
    kind = KIND_OF_REL.get(relation, "assoc")
    label = f"NON {relation} {int(round(w))}" if is_neg else f"{relation} {int(round(w))}"
    color = _edge_color(kind, depth, is_neg)
    opacity = _opacity_for(depth)
    # Couleur du label : même teinte que le trait pour lier visuellement,
    # avec un léger fond blanc pour rester lisible par-dessus les nœuds.
    font_color = color
    return {
        "from": from_id,
        "to": to_id,
        "label": label,
        "arrows": "to",
        "font": {
            "size": 13 if is_d2 else 14,
            "color": font_color,
            "background": "#d4d7daee",
            "strokeWidth": 0,
        },
        # vis-network supporte color: { color, opacity }
        "color": {"color": color, "opacity": opacity},
        "dashes": is_d2,
        "smooth": {"type": "dynamic"},
        "_relation": relation,
        "_kind": kind,
        "_weight": w,
        "_polarity": polarity,
        "_negative": is_neg,
        "_depth": depth,
        "_from_negation": from_negation,
    }


def build_subgraph(
    term: str,
    *,
    client: Optional[JDMClient] = None,
    depth: int = 1,
    top_k_per_relation: int = 3,
    top_k_depth2: Optional[int] = None,
    top_k_depth3: Optional[int] = None,
    top_k_depth4: Optional[int] = None,
    min_weight: Optional[float] = None,
    relations: Optional[list[str]] = None,
    depth2_relations: Optional[list[str]] = None,
    depth3_relations: Optional[list[str]] = None,
    depth4_relations: Optional[list[str]] = None,
    output: Literal["json", "html"] = "html",
    output_path: Optional[str] = None,
) -> dict[str, Any]:
    """Construit un sous-graphe JDM centré sur `term`.

    Args:
        term: terme racine (en français, accentué si besoin).
        client: JDMClient injecté ; un client par défaut sera créé sinon.
        depth: profondeur d'exploration (1..4 ; au-delà = illisible).
        top_k_per_relation: nombre max de cibles retenues par relation et par nœud.
        min_weight: poids minimum (None = pas de filtre, JDM décide).
        relations: relations explorées à la profondeur 1 (défaut = `DEFAULT_RELATIONS`).
        depth2_relations: relations à la profondeur 2 (défaut = `DEFAULT_DEPTH2_RELATIONS`).
        depth3_relations: relations à la profondeur 3 (défaut = `DEFAULT_DEPTH3_RELATIONS`).
        depth4_relations: relations à la profondeur 4 (défaut = `DEFAULT_DEPTH4_RELATIONS`).
        output: "json" → dict `{nodes, edges, ...}` ; "html" → écrit un fichier autonome.
        output_path: chemin d'écriture (si None et output="html", utilise `<slug>_subgraph.html` dans le CWD).

    Returns:
        Dict avec les clés :
        - `root`: terme racine
        - `nodes`: liste de nœuds vis-network (présent si output="json")
        - `edges`: liste d'arêtes vis-network (présent si output="json")
        - `stats`: {n_nodes, n_edges, n_negative, relations_used, depth}
        - `html_path`: chemin du fichier écrit (si output="html")
    """
    c = client or JDMClient()
    depth = max(1, min(int(depth), 4))  # garde-fou : 1..4

    # Si l'agent passe un raffinement brut (`guitare>91594`), on garde la
    # forme brute pour les requêtes HTTP (clé interne JDM) mais on affiche
    # la forme décodée partout (titre, légende, nœud central).
    raw_term = term
    try:
        _dec = c.decode_node_name(term)
        term_display = _dec.get("decoded") or term
    except Exception:
        term_display = term

    # Sélection effective par profondeur.
    rels_by_depth: dict[int, list[str]] = {
        1: list(relations) if relations is not None else list(DEFAULT_RELATIONS),
        2: list(depth2_relations) if depth2_relations is not None else list(DEFAULT_DEPTH2_RELATIONS),
        3: list(depth3_relations) if depth3_relations is not None else list(DEFAULT_DEPTH3_RELATIONS),
        4: list(depth4_relations) if depth4_relations is not None else list(DEFAULT_DEPTH4_RELATIONS),
    }
    # Top-K effectif par profondeur : si non précisé, on retombe sur le
    # top_k_per_relation global (compat ascendante).
    top_k_by_depth: dict[int, int] = {
        1: int(top_k_per_relation),
        2: int(top_k_depth2) if top_k_depth2 is not None else int(top_k_per_relation),
        3: int(top_k_depth3) if top_k_depth3 is not None else int(top_k_per_relation),
        4: int(top_k_depth4) if top_k_depth4 is not None else int(top_k_per_relation),
    }

    # 1) Nœud central — affichage en forme DÉCODÉE (cf. raffinements
    # `guitare>91594` → `guitare (instrument de musique)`), HTTP queries
    # restent sur `raw_term` pour que JDM accepte la clé interne.
    root_node = _build_node("ROOT", term_display, "center", depth=0, fixed_center=True)
    nodes: list[dict[str, Any]] = [root_node]
    edges: list[dict[str, Any]] = []
    # Index pour dédupliquer les nœuds. On enregistre les DEUX formes
    # (raw + display) → "ROOT", car les rows enfants peuvent référencer
    # la racine via l'une ou l'autre.
    label_to_id: dict[str, str] = {term_display: "ROOT", raw_term: "ROOT"}
    next_uid = [0]

    def _ensure_node(label: str, kind: str, depth_lv: int) -> str:
        if label in label_to_id:
            return label_to_id[label]
        next_uid[0] += 1
        nid = f"N{next_uid[0]}"
        label_to_id[label] = nid
        nodes.append(_build_node(
            nid, label,
            kind,  # garde la teinte de la famille ; _palette_for éclaircit au niveau 2
            depth=depth_lv,
        ))
        return nid

    # Ensemble des labels de nœuds INTRODUITS par une arête négative.
    # Les arêtes qui PARTENT de ces nœuds (= leur approfondissement vers
    # le niveau suivant) sont taguées `_from_negation` pour pouvoir être
    # masquées côté HTML (bouton « Approfondir négations », OFF par défaut).
    via_negation: set[str] = set()

    # Boucle généralisée sur les profondeurs 1..depth. À chaque tour on
    # garde la liste des nœuds tout juste ajoutés pour itérer dessus au
    # tour suivant.
    current_layer: list[tuple[str, str]] = [(term, "center")]
    for d in range(1, depth + 1):
        rels_for_d = rels_by_depth.get(d, [])
        # Top-K par niveau choisi par l'utilisateur (défaut = top_k_per_relation).
        top_k_d = max(1, top_k_by_depth.get(d, top_k_per_relation))
        next_layer: list[tuple[str, str]] = []
        for parent_label, _parent_kind in current_layer:
            parent_id = label_to_id.get(parent_label, "ROOT")
            parent_is_neg = parent_label in via_negation
            for rel in rels_for_d:
                kind = KIND_OF_REL.get(rel, "assoc")
                rows = _fetch_relation(c, parent_label, rel, top_k_d, min_weight)
                for row in rows:
                    tgt = row["target_display"]
                    is_neg_edge = row["polarity"] == "négation"
                    if tgt == term_display or tgt == raw_term:
                        # Lien retour vers la racine : on matérialise l'arête
                        # mais on ne ré-ajoute pas le nœud.
                        edges.append(_build_edge(
                            parent_id, "ROOT", rel, row["w"], row["polarity"],
                            depth=d, from_negation=parent_is_neg,
                        ))
                        continue
                    is_new = tgt not in label_to_id
                    nid_to = _ensure_node(tgt, kind, depth_lv=d)
                    if parent_id == nid_to:
                        continue
                    edges.append(_build_edge(
                        parent_id, nid_to, rel, row["w"], row["polarity"],
                        depth=d, from_negation=parent_is_neg,
                    ))
                    if is_new:
                        next_layer.append((tgt, kind))
                        # Nœud nouveau atteint par une arête négative → sa
                        # descendance sera masquable via le bouton dédié.
                        if is_neg_edge:
                            via_negation.add(tgt)
        current_layer = next_layer
        if not current_layer:
            break

    n_negative = sum(1 for e in edges if e.get("_negative"))
    n_from_negation = sum(1 for e in edges if e.get("_from_negation"))
    stats = {
        "n_nodes": len(nodes),
        "n_edges": len(edges),
        "n_negative": n_negative,
        "n_from_negation": n_from_negation,
        "relations_used": rels_by_depth[1],
        "relations_by_depth": {d: rels_by_depth[d] for d in range(1, depth + 1)},
        "depth": depth,
    }
    rels = rels_by_depth[1]  # legacy alias pour la légende ci-dessous

    result: dict[str, Any] = {"root": term, "stats": stats}

    if output == "json":
        result["nodes"] = nodes
        result["edges"] = edges
        return result

    # output == "html"
    legend_chips = [
        f'<span style="background:#212121;color:#fff;">{term_display}</span>',
    ]
    for rel in rels:
        kind = KIND_OF_REL.get(rel, "assoc")
        bg = PALETTE.get(kind, PALETTE["assoc"])["background"]
        legend_chips.append(f'<span style="background:{bg};">{rel}</span>')
    if depth >= 2:
        legend_chips.append(
            '<span style="background:#f5f5f5;border:1px dashed #9e9e9e;">profondeur ≥ 2 (pointillés, opacité décroissante)</span>'
        )
    legend_chips.append(
        '<span style="background:#ffebee;color:#c62828;border:1px solid #c62828;">négation</span>'
    )

    html = (
        HTML_TEMPLATE
        .replace("{{TITLE}}", f"« {term_display} » — sous-graphe JDM (profondeur {depth})")
        .replace(
            "{{SUBTITLE}}",
            "Couleur par type de relation, opacité décroissante avec la profondeur, "
            "négations en rouge. Molette = zoom · glisser = déplacer · "
            "clic sur un nœud = le cadrer avec ses connexions · "
            "double-clic = en faire le centre de gravité.",
        )
        .replace("{{LEGEND}}", "".join(legend_chips))
        .replace("{{NODES_JSON}}", json.dumps(nodes, ensure_ascii=False))
        .replace("{{EDGES_JSON}}", json.dumps(edges, ensure_ascii=False))
    )

    if output_path is None:
        output_path = f"{_slugify(term)}_subgraph.html"
    p = Path(output_path)
    p.write_text(html, encoding="utf-8")
    result["html_path"] = str(p.resolve())
    return result