#!/usr/bin/env python3
# Topic 06 — Language & Terms — Visuals (FINAL)
#
# Generates the standardized Topic 06 visuals from CSV exports under:
#   reports/06_language_and_terms/
#
# Outputs PNGs to:
#   reports/06_language_and_terms/visuals/
#
# Figures:
#   01  Exclusive-term concentration (Top 25 docs per cluster)
#   02  Top shared tracking-font hashes by doc coverage
#   03  Exclusive-terms presence heatmap (with term labels throttled)
#   04  Cluster ↔ Tracking Font Hash bipartite graph (margins + full-height cluster spacing)
#   05  Cluster ↔ Exclusive Term bipartite graph (margins + full-height cluster spacing)
#
# Notes:
# - Graphs (04/05) require networkx. Install:  python3 -m pip install networkx
# - Run from project root recommended (paths are relative to cwd).

from __future__ import annotations

import argparse
import os
import re
from typing import Callable, List

import pandas as pd
import matplotlib.pyplot as plt


# -----------------------------
# Configuration (edit as needed)
# -----------------------------

BASE = os.path.join("reports", "06_language_and_terms")
OUT  = os.path.join(BASE, "visuals")

# Current mapping for Topic 06 focal/peer clusters
CLUSTER_MAP = {
    1570: "MATTHEW GUERTIN",
    674:  "MUAD ABDULKADIR",
    290:  "ADRIAN WESLEY",
    696:  "PETER LEHMEYER",
}

# Used for most plots (left-to-right legend order)
CLUSTER_ORDER = [1570, 674, 290, 696]

# Used for graph plots (top-to-bottom on left; requested PETER at top)
CLUSTER_ORDER_TOPDOWN = [696, 290, 674, 1570]


# -----------------------------
# Helpers
# -----------------------------

def ensure_out() -> None:
    os.makedirs(OUT, exist_ok=True)

def read_csv(name: str) -> pd.DataFrame:
    path = os.path.join(BASE, name)
    if not os.path.exists(path):
        raise FileNotFoundError(f"Missing expected input: {path}")
    return pd.read_csv(path)

def save_fig(filename: str) -> None:
    path = os.path.join(OUT, filename)
    # bbox_inches prevents label clipping; works well for all figures
    plt.savefig(path, dpi=220, bbox_inches="tight")
    plt.close()
    print("wrote", path)

def short_hash(h: str, left: int = 12, right: int = 6) -> str:
    if not isinstance(h, str) or not h:
        return ""
    if len(h) <= left + right + 1:
        return h
    return f"{h[:left]}…{h[-right:]}"

def clean_term(s: str, maxlen: int = 56) -> str:
    if s is None:
        return ""
    s = str(s)
    s = re.sub(r"\s+", " ", s).strip()
    if len(s) <= maxlen:
        return s
    return s[:maxlen-1] + "…"


# -----------------------------
# Plot 01
# -----------------------------

def plot_doc_concentration(top_n: int = 25) -> None:
    df = read_csv("06_exclusive_terms_doc_concentration.csv")

    # rank within cluster
    df = df.sort_values(["cluster_id", "n_exclusive_terms_in_doc"], ascending=[True, False])
    df["rank"] = df.groupby("cluster_id").cumcount() + 1
    df_top = df[df["rank"] <= top_n].copy()

    plt.figure(figsize=(10, 6))
    for cid in CLUSTER_ORDER:
        g = df_top[df_top["cluster_id"] == cid].sort_values("rank")
        if len(g) == 0:
            continue
        plt.plot(
            g["rank"], g["n_exclusive_terms_in_doc"],
            marker="o", linestyle="-",
            label=f"{CLUSTER_MAP.get(cid, cid)} ({cid})"
        )

    plt.xlabel(f"Rank (top {top_n} docs per cluster)")
    plt.ylabel("Distinct exclusive terms in document")
    plt.title(f"Exclusive-term concentration (Top {top_n} docs per cluster)")
    plt.legend(title="Cluster")
    plt.tight_layout()
    save_fig("01_doc_concentration_top25.png")


# -----------------------------
# Plot 02
# -----------------------------

def plot_font_bridge_bars(top_k: int = 20) -> None:
    df = read_csv("06x_font_hash_cluster_counts.csv")

    df = df.sort_values(
        ["n_clusters", "n_docs", "n_object_rows"],
        ascending=[False, False, False]
    ).head(top_k).copy()

    df["hash_short"] = df["object_sha256"].map(lambda x: short_hash(str(x), 12, 6))

    plt.figure(figsize=(12, 6))
    plt.bar(range(len(df)), df["n_docs"])
    plt.xticks(range(len(df)), df["hash_short"], rotation=70, ha="right")
    plt.xlabel("Tracking font object_sha256 (short)")
    plt.ylabel("Distinct docs containing font hash")
    plt.title(f"Top shared tracking-font hashes by doc coverage (Top {top_k})")
    plt.tight_layout()
    save_fig("02_top_tracking_font_hashes_by_docs.png")


# -----------------------------
# Plot 03
# -----------------------------

def plot_term_presence_heatmap(max_rows: int = 140, label_every: int = 5) -> None:
    df = read_csv("06_terms_exclusive_to_focal_plus_3.csv")

    cols = ["in_cluster_1570", "in_cluster_674", "in_cluster_290", "in_cluster_696"]
    for c in cols:
        if c not in df.columns:
            raise ValueError(f"Missing expected column in term table: {c}")

    # presence pattern bitstring: 1101 etc.
    df["presence_bits"] = df[cols].astype(int).astype(str).agg("".join, axis=1)
    df["presence_sum"] = df[cols].sum(axis=1)

    # sort: most shared first, then by pattern, then stable
    df = df.sort_values(
        ["presence_sum", "presence_bits", "string_group_n", "string_norm"],
        ascending=[False, True, True, True]
    ).head(max_rows).copy()

    mat = df[cols].astype(int).values
    xticklabels = [f"{CLUSTER_MAP[int(cid)]}\n({cid})" for cid in CLUSTER_ORDER]

    # throttle y-labels
    ylabels: List[str] = []
    for i, s in enumerate(df["string_norm"].astype(str).tolist()):
        ylabels.append(clean_term(s, 42) if i % label_every == 0 else "")

    plt.figure(figsize=(10, 10))
    plt.imshow(mat, aspect="auto")
    plt.xticks(range(4), xticklabels)
    plt.yticks(range(len(df)), ylabels, fontsize=7)
    plt.xlabel("Cluster")
    plt.title(
        "Exclusive terms: presence across clusters\n"
        f"(Top terms by cross-cluster sharing; labels every {label_every} rows)"
    )
    plt.tight_layout()
    save_fig("03_exclusive_terms_presence_heatmap.png")


# -----------------------------
# Plot 04/05 — Bipartite graphs
# -----------------------------

def plot_bipartite_graph(
    edge_csv: str,
    left_col: str,
    right_col: str,
    weight_col: str,
    out_png: str,
    title: str,
    right_label_fn: Callable[[str], str],
    top_right: int,
    top_edges: int,
    node_size: int = 90,
    edge_width: float = 0.7,
) -> None:
    """
    Graph layout improvements:
    - prevents label clipping
    - spreads left cluster labels across full height (PETER at top)
    - spreads right labels across full height
    """
    df = read_csv(edge_csv)

    # aggregate edges
    df = df.groupby([left_col, right_col], as_index=False)[weight_col].sum()
    df[left_col] = df[left_col].astype(str)
    df[right_col] = df[right_col].astype(str)

    # choose top right nodes
    right_strength = (
        df.groupby(right_col, as_index=False)[weight_col]
          .sum()
          .sort_values(weight_col, ascending=False)
    )
    keep_right = set(right_strength.head(top_right)[right_col].astype(str).tolist())
    df = df[df[right_col].isin(keep_right)].copy()
    df = df.sort_values(weight_col, ascending=False).head(top_edges).copy()

    try:
        import networkx as nx
    except Exception as e:
        raise RuntimeError(
            "networkx is required for graphs 04/05. Install with: python3 -m pip install networkx"
        ) from e

    G = nx.Graph()
    for _, r in df.iterrows():
        G.add_edge(r[left_col], r[right_col], weight=float(r[weight_col]))

    present_left = set(df[left_col].unique().tolist())
    left_nodes = [str(cid) for cid in CLUSTER_ORDER_TOPDOWN if str(cid) in present_left]
    left_nodes += [n for n in df[left_col].unique().tolist() if n not in left_nodes]
    right_nodes = df[right_col].unique().tolist()

    # Layout parameters
    x_left, x_right = 0.10, 0.90

    # Spread nodes evenly across full height
    if len(left_nodes) == 1:
        left_y = {left_nodes[0]: 0.5}
    else:
        left_y = {n: i/(len(left_nodes)-1) for i, n in enumerate(left_nodes)}

    if len(right_nodes) == 1:
        right_y = {right_nodes[0]: 0.5}
    else:
        right_y = {n: j/(len(right_nodes)-1) for j, n in enumerate(right_nodes)}

    pos = {**{n: (x_left, left_y[n]) for n in left_nodes},
           **{n: (x_right, right_y[n]) for n in right_nodes}}

    plt.figure(figsize=(22, 12))
    ax = plt.gca()
    nx.draw(G, pos, with_labels=False, node_size=node_size, width=edge_width, ax=ax)

    # Pad canvas & margins so labels don't clip
    ax.set_xlim(0.0, 1.0)
    ax.set_ylim(-0.05, 1.05)
    plt.subplots_adjust(left=0.08, right=0.92, top=0.92, bottom=0.06)

    # Left labels (shift inward)
    for n in left_nodes:
        cid = int(n) if n.isdigit() else n
        label = CLUSTER_MAP.get(cid, str(cid))
        x, y = pos[n]
        ax.text(x - 0.02, y, label, ha="right", va="center", fontsize=14)

    # Right labels (shift inward)
    for n in right_nodes:
        x, y = pos[n]
        ax.text(x + 0.02, y, right_label_fn(n), ha="left", va="center", fontsize=10)

    ax.set_title(title, fontsize=18)
    ax.axis("off")
    save_fig(out_png)


def plot_graph_cluster_to_font(top_right: int = 28, top_edges: int = 180) -> None:
    plot_bipartite_graph(
        edge_csv="06x_edges_cluster_font.csv",
        left_col="cluster_id",
        right_col="object_sha256",
        weight_col="doc_weight",
        out_png="04_graph_cluster_to_font_edges.png",
        title="Cluster ↔ Tracking Font Hash graph\nRight nodes = tracking font object_sha256 (short)",
        right_label_fn=lambda h: short_hash(h, 12, 6),
        top_right=top_right,
        top_edges=top_edges,
    )

def plot_graph_cluster_to_term(top_right: int = 45, top_edges: int = 240) -> None:
    plot_bipartite_graph(
        edge_csv="06x_edges_cluster_term.csv",
        left_col="cluster_id",
        right_col="string_norm",
        weight_col="doc_weight",
        out_png="05_graph_cluster_to_term_edges.png",
        title="Cluster ↔ Exclusive Term graph\nRight nodes = exclusive term string_norm (truncated)",
        right_label_fn=lambda s: clean_term(s, 56),
        top_right=top_right,
        top_edges=top_edges,
    )


# -----------------------------
# CLI
# -----------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Generate Topic 06 visuals from exported CSVs under reports/06_language_and_terms/."
    )
    p.add_argument(
        "--only",
        default="",
        help="Comma-separated list of figure numbers to run (e.g. 01,03,05). Default runs all.",
    )
    p.add_argument("--top-docs", type=int, default=25, help="Top docs per cluster for plot 01.")
    p.add_argument("--top-fonts", type=int, default=20, help="Top font hashes for plot 02.")
    p.add_argument("--heatmap-rows", type=int, default=140, help="Max rows to show in heatmap (03).")
    p.add_argument("--heatmap-label-every", type=int, default=5, help="Label every Nth row in heatmap (03).")
    p.add_argument("--graph-font-top-right", type=int, default=28, help="Right nodes to keep in graph 04.")
    p.add_argument("--graph-font-top-edges", type=int, default=180, help="Edges to keep in graph 04.")
    p.add_argument("--graph-term-top-right", type=int, default=45, help="Right nodes to keep in graph 05.")
    p.add_argument("--graph-term-top-edges", type=int, default=240, help="Edges to keep in graph 05.")
    return p.parse_args()

def main() -> None:
    args = parse_args()
    ensure_out()

    only = [s.strip() for s in args.only.split(",") if s.strip()]

    def run(fig: str) -> bool:
        return (not only) or (fig in only)

    if run("01"):
        plot_doc_concentration(top_n=args.top_docs)
    if run("02"):
        plot_font_bridge_bars(top_k=args.top_fonts)
    if run("03"):
        plot_term_presence_heatmap(max_rows=args.heatmap_rows, label_every=args.heatmap_label_every)
    if run("04"):
        plot_graph_cluster_to_font(top_right=args.graph_font_top_right, top_edges=args.graph_font_top_edges)
    if run("05"):
        plot_graph_cluster_to_term(top_right=args.graph_term_top_right, top_edges=args.graph_term_top_edges)

    print("Done. Outputs in:", OUT)

if __name__ == "__main__":
    main()