Update conversation_storyline/plots.py

conversation_storyline/plots.py

@@ -1,83 +1,277 @@
-from __future__ import annotations
-
-from typing import Dict, List, Tuple
-import numpy as np
-import pandas as pd
-import plotly.graph_objects as go
-import networkx as nx
-
-
-def plot_reply_distance_hist(metrics: pd.DataFrame) -> go.Figure:
-    d = metrics["reply_distance"].dropna().astype(int)
-    fig = go.Figure()
-    fig.add_histogram(x=d, nbinsx=40)
-    fig.update_layout(title="Distribución distancia reply_to (message_id - reply_to_id)", xaxis_title="distancia", yaxis_title="conteo")
-    return fig
-
-
-def plot_interaction_heatmap(inter_matrix: pd.DataFrame) -> go.Figure:
-    fig = go.Figure(data=go.Heatmap(z=inter_matrix.values, x=inter_matrix.columns, y=inter_matrix.index))
-    fig.update_layout(title="Heatmap interacciones (conteo respuestas)", xaxis_title="to", yaxis_title="from")
-    return fig
-
-
-def plot_bump_activity(metrics: pd.DataFrame) -> go.Figure:
-    # actividad por topic_id y speaker
-    if "topic_id" not in metrics.columns:
-        fig = go.Figure()
-        fig.update_layout(title="Bump actividad (no topic_id)")
-        return fig
-
-    piv = metrics.pivot_table(index="topic_id", columns="speaker", values="message_id", aggfunc="count", fill_value=0)
-    # rank por segmento (mayor actividad = rank 1)
-    ranks = piv.rank(axis=1, method="average", ascending=False)
-
-    fig = go.Figure()
-    for sp in piv.columns:
-        fig.add_trace(go.Scatter(x=piv.index, y=ranks[sp], mode="lines+markers", name=sp))
-    fig.update_layout(
-        title="Bump chart: ranking actividad por segmento",
-        xaxis_title="topic_id",
-        yaxis_title="rank (1 = más activo)",
-        yaxis_autorange="reversed",
-    )
-    return fig
-
-
-def plot_sankey_speaker_to_topic(metrics: pd.DataFrame) -> go.Figure:
-    if "topic_label" not in metrics.columns:
-        fig = go.Figure()
-        fig.update_layout(title="Sankey (no topic_label)")
-        return fig
-
-    speakers = metrics["speaker"].unique().tolist()
-    topics = metrics["topic_label"].fillna("Tema").unique().tolist()
-
-    s_idx = {s: i for i, s in enumerate(speakers)}
-    t_idx = {t: i + len(speakers) for i, t in enumerate(topics)}
-
-    links = metrics.groupby(["speaker", "topic_label"])["message_id"].count().reset_index()
-    source = [s_idx[r["speaker"]] for _, r in links.iterrows()]
-    target = [t_idx[r["topic_label"]] for _, r in links.iterrows()]
-    value = links["message_id"].tolist()
-
-    labels = speakers + topics
-
-    fig = go.Figure(
-        data=[
-            go.Sankey(
-                node=dict(label=labels, pad=10, thickness=12),
-                link=dict(source=source, target=target, value=value),
-            )
-        ]
-    )
-    fig.update_layout(title="Sankey: Speaker → Topic (volumen de mensajes)")
-    return fig
-
-
-def plot_sentiment_placeholder(metrics: pd.DataFrame) -> go.Figure:
-    fig = go.Figure()
-    if "sentiment" in metrics.columns and metrics["sentiment"].notna().any():
-        fig.add_trace(go.Scatter(x=metrics["message_id"], y=metrics["sentiment"], mode="lines+markers"))
-    fig.update_layout(title="Sentiment (si disponible)", xaxis_title="message_id", yaxis_title="sentiment")
-    return fig
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+import networkx as nx
+
+
+def plot_reply_distance_hist(metrics: pd.DataFrame) -> go.Figure:
+    d = metrics["reply_distance"].dropna().astype(int)
+    fig = go.Figure()
+    fig.add_histogram(x=d, nbinsx=40)
+    fig.update_layout(title="Distribución distancia reply_to (message_id - reply_to_id)", xaxis_title="distancia", yaxis_title="conteo")
+    return fig
+
+
+def plot_interaction_heatmap(inter_matrix: pd.DataFrame) -> go.Figure:
+    fig = go.Figure(data=go.Heatmap(z=inter_matrix.values, x=inter_matrix.columns, y=inter_matrix.index))
+    fig.update_layout(title="Heatmap interacciones (conteo respuestas)", xaxis_title="to", yaxis_title="from")
+    return fig
+
+
+def plot_bump_activity(metrics: pd.DataFrame) -> go.Figure:
+    # actividad por topic_id y speaker
+    if "topic_id" not in metrics.columns:
+        fig = go.Figure()
+        fig.update_layout(title="Bump actividad (no topic_id)")
+        return fig
+
+    piv = metrics.pivot_table(index="topic_id", columns="speaker", values="message_id", aggfunc="count", fill_value=0)
+    # rank por segmento (mayor actividad = rank 1)
+    ranks = piv.rank(axis=1, method="average", ascending=False)
+
+    fig = go.Figure()
+    for sp in piv.columns:
+        fig.add_trace(go.Scatter(x=piv.index, y=ranks[sp], mode="lines+markers", name=sp))
+    fig.update_layout(
+        title="Bump chart: ranking actividad por segmento",
+        xaxis_title="topic_id",
+        yaxis_title="rank (1 = más activo)",
+        yaxis_autorange="reversed",
+    )
+    return fig
+
+
+def plot_sankey_speaker_to_topic(metrics: pd.DataFrame) -> go.Figure:
+    if "topic_label" not in metrics.columns:
+        fig = go.Figure()
+        fig.update_layout(title="Sankey (no topic_label)")
+        return fig
+
+    speakers = metrics["speaker"].unique().tolist()
+    topics = metrics["topic_label"].fillna("Tema").unique().tolist()
+
+    s_idx = {s: i for i, s in enumerate(speakers)}
+    t_idx = {t: i + len(speakers) for i, t in enumerate(topics)}
+
+    links = metrics.groupby(["speaker", "topic_label"])["message_id"].count().reset_index()
+    source = [s_idx[r["speaker"]] for _, r in links.iterrows()]
+    target = [t_idx[r["topic_label"]] for _, r in links.iterrows()]
+    value = links["message_id"].tolist()
+
+    labels = speakers + topics
+
+    fig = go.Figure(
+        data=[
+            go.Sankey(
+                node=dict(label=labels, pad=10, thickness=12),
+                link=dict(source=source, target=target, value=value),
+            )
+        ]
+    )
+    fig.update_layout(title="Sankey: Speaker → Topic (volumen de mensajes)")
+    return fig
+
+
+def plot_sentiment_placeholder(metrics: pd.DataFrame) -> go.Figure:
+    fig = go.Figure()
+    if "sentiment" in metrics.columns and metrics["sentiment"].notna().any():
+        fig.add_trace(go.Scatter(x=metrics["message_id"], y=metrics["sentiment"], mode="lines+markers"))
+    fig.update_layout(title="Sentiment (si disponible)", xaxis_title="message_id", yaxis_title="sentiment")
+    return fig
+# conversation_storyline/plots.py
+"""
+Plot helpers for Conversation Storyline outputs.
+
+Derived from:
+- interactions.jsonl
+- graph.json
+
+Returns Plotly figures (Gradio-friendly).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+
+INTERACTIONS_FILENAME = "interactions.jsonl"
+GRAPH_FILENAME = "graph.json"
+
+
+def _read_jsonl(path: Path) -> List[dict]:
+    rows: List[dict] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            rows.append(json.loads(line))
+    return rows
+
+
+def load_interactions_df(out_dir: str | Path) -> pd.DataFrame:
+    out_dir = Path(out_dir)
+    path = out_dir / INTERACTIONS_FILENAME
+    if not path.exists():
+        raise FileNotFoundError(f"Missing {INTERACTIONS_FILENAME} in {out_dir}")
+    df = pd.DataFrame(_read_jsonl(path))
+    if "id" in df.columns:
+        df = df.sort_values("id").reset_index(drop=True)
+    if "speaker" in df.columns:
+        df["speaker"] = df["speaker"].astype(str).fillna("?")
+    if "topic_label" in df.columns:
+        df["topic_label"] = df["topic_label"].astype(str).fillna("?")
+    if "sentiment_score" in df.columns:
+        df["sentiment_score"] = pd.to_numeric(df["sentiment_score"], errors="coerce").fillna(0.0)
+    if "reply_to_id" in df.columns:
+        df["reply_to_id"] = pd.to_numeric(df["reply_to_id"], errors="coerce")
+    return df
+
+
+def load_graph_json(out_dir: str | Path) -> dict:
+    out_dir = Path(out_dir)
+    path = out_dir / GRAPH_FILENAME
+    if not path.exists():
+        raise FileNotFoundError(f"Missing {GRAPH_FILENAME} in {out_dir}")
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def plot_sentiment_over_time(df: pd.DataFrame, rolling_window: int | None = None) -> go.Figure:
+    if df.empty:
+        return go.Figure()
+
+    n = len(df)
+    if rolling_window is None:
+        rolling_window = max(3, min(25, int(n * 0.05)))
+
+    dfx = df[["id", "speaker", "sentiment_score"]].copy()
+    dfx["rolling_sentiment"] = (
+        dfx.groupby("speaker")["sentiment_score"]
+        .transform(lambda s: s.rolling(rolling_window, min_periods=1).mean())
+    )
+
+    fig = px.line(dfx, x="id", y="rolling_sentiment", color="speaker")
+    fig.update_layout(
+        title=f"Sentiment (media móvil {rolling_window} mensajes)",
+        xaxis_title="Turno (id)",
+        yaxis_title="Sentiment",
+        height=420,
+    )
+    return fig
+
+
+def plot_sentiment_histogram(df: pd.DataFrame) -> go.Figure:
+    if df.empty:
+        return go.Figure()
+    fig = px.histogram(df, x="sentiment_score", nbins=30)
+    fig.update_layout(
+        title="Distribución de sentiment",
+        xaxis_title="Sentiment",
+        yaxis_title="Nº de mensajes",
+        height=380,
+    )
+    return fig
+
+
+def plot_speaker_topic_heatmap(df: pd.DataFrame, max_topics: int = 20) -> go.Figure:
+    if df.empty:
+        return go.Figure()
+
+    top_topics = df["topic_label"].value_counts().head(max_topics).index.tolist()
+    sub = df[df["topic_label"].isin(top_topics)]
+    pivot = pd.pivot_table(sub, index="speaker", columns="topic_label", values="id", aggfunc="count", fill_value=0)
+
+    fig = px.imshow(pivot, aspect="auto", color_continuous_scale="Blues")
+    fig.update_layout(
+        title="Heatmap: speaker × topic (conteo de mensajes)",
+        xaxis_title="Topic",
+        yaxis_title="Speaker",
+        height=520,
+    )
+    return fig
+
+
+def plot_speaker_activity_heatmap(df: pd.DataFrame, bins: int = 25) -> go.Figure:
+    """
+    Heatmap: speakers vs time buckets. Useful for large chats.
+    """
+    if df.empty:
+        return go.Figure()
+
+    n = len(df)
+    bins = max(10, min(bins, 60))
+    df = df.copy()
+    df["bucket"] = pd.cut(df["id"], bins=bins, labels=False, include_lowest=True)
+    pivot = pd.pivot_table(df, index="speaker", columns="bucket", values="id", aggfunc="count", fill_value=0)
+
+    fig = px.imshow(pivot, aspect="auto", color_continuous_scale="Viridis")
+    fig.update_layout(
+        title=f"Actividad por speaker (buckets={bins})",
+        xaxis_title="Bucket temporal",
+        yaxis_title="Speaker",
+        height=520,
+    )
+    return fig
+
+
+def plot_topic_shift_timeline(df: pd.DataFrame) -> go.Figure:
+    if df.empty or "is_topic_shift" not in df.columns:
+        return go.Figure()
+
+    df = df.copy()
+    df["shift"] = df["is_topic_shift"].astype(int)
+
+    fig = px.scatter(
+        df[df["shift"] == 1],
+        x="id",
+        y=["shift"] * len(df[df["shift"] == 1]),
+        color="speaker",
+        hover_data=["topic_label"],
+    )
+    fig.update_layout(
+        title="Timeline de Topic Shifts",
+        xaxis_title="Turno (id)",
+        yaxis_title="Shift (1)",
+        height=250,
+        showlegend=True,
+    )
+    return fig
+
+
+def plot_reply_sankey(graph: dict, min_weight: int = 1) -> go.Figure:
+    nodes = graph.get("nodes", [])
+    links = graph.get("links", [])
+
+    speakers = [n["id"] for n in nodes if "id" in n]
+    idx = {s: i for i, s in enumerate(speakers)}
+
+    sources, targets, values = [], [], []
+    for l in links:
+        w = int(l.get("weight", 1))
+        if w < min_weight:
+            continue
+        s = l.get("source")
+        t = l.get("target")
+        if s in idx and t in idx:
+            sources.append(idx[s])
+            targets.append(idx[t])
+            values.append(w)
+
+    fig = go.Figure(
+        data=[
+            go.Sankey(
+                node=dict(label=speakers, pad=15, thickness=15),
+                link=dict(source=sources, target=targets, value=values),
+            )
+        ]
+    )
+    fig.update_layout(title="Sankey: quién responde a quién", height=520)
+    return fig