#!/usr/bin/env python3
"""
10_font_tracking__run01__visuals__one_shot_final_v1.py

One-shot visual pipeline for Run01 font-tracking bipartite graph:
- Reads edges CSV (cluster/person -> font hash/style/etc.)
- Builds bipartite NetworkX graph
- Writes:
    * static graph PNG (Graphviz if available, fallback spring-layout)
    * static graph SVG (Graphviz only)
    * interactive graph HTML (PyVis)
    * node and edge export CSVs
    * summary JSON

Designed for Ubuntu and large-ish graphs.
Assumes dependencies are installed:
    pandas, networkx, matplotlib, pyvis
Optional for best static layout:
    graphviz (system binary) + pydot

Usage examples:
    python3 10_font_tracking__run01__visuals__one_shot_final_v1.py
    python3 10_font_tracking__run01__visuals__one_shot_final_v1.py \
        --input reports/10_font_tracking/10r1_edges.csv \
        --out_dir reports/10_font_tracking \
        --min_edge_weight 2 \
        --max_nodes 1200 \
        --interactive \
        --static

Notes:
- This script expects an edge list CSV. If your column names differ,
  pass --left_col / --right_col / --weight_col or let auto-detect try.
- Graphviz label warnings ("size too small for label") are non-fatal.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Optional, Tuple

import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt


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

def _log(msg: str) -> None:
    print(msg, flush=True)


def _safe_mkdir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def _first_present(df: pd.DataFrame, candidates: List[str]) -> Optional[str]:
    lower_map = {c.lower(): c for c in df.columns}
    for c in candidates:
        if c.lower() in lower_map:
            return lower_map[c.lower()]
    return None


def detect_columns(df: pd.DataFrame,
                   left_col: Optional[str],
                   right_col: Optional[str],
                   weight_col: Optional[str]) -> Tuple[str, str, str]:
    """
    Detect reasonable defaults for left/right/weight columns.
    """
    if left_col and left_col not in df.columns:
        raise ValueError(f"--left_col '{left_col}' not in CSV columns")
    if right_col and right_col not in df.columns:
        raise ValueError(f"--right_col '{right_col}' not in CSV columns")
    if weight_col and weight_col not in df.columns:
        raise ValueError(f"--weight_col '{weight_col}' not in CSV columns")

    if left_col is None:
        left_col = _first_present(df, [
            "cluster", "cluster_id", "case_cluster", "person", "name",
            "left", "source", "src", "node_a"
        ])
    if right_col is None:
        right_col = _first_present(df, [
            "font_hash", "font", "font_id", "term", "right", "target",
            "dst", "node_b"
        ])
    if weight_col is None:
        weight_col = _first_present(df, [
            "weight", "edge_weight", "docs", "count", "n", "value"
        ])

    # fallback heuristics
    if left_col is None or right_col is None:
        cols = list(df.columns)
        if len(cols) < 2:
            raise ValueError("Need at least 2 columns in edge CSV")
        if left_col is None:
            left_col = cols[0]
        if right_col is None:
            right_col = cols[1]

    if weight_col is None:
        # create synthetic weight=1
        tmp_weight = "__synthetic_weight__"
        df[tmp_weight] = 1
        weight_col = tmp_weight

    return left_col, right_col, weight_col


def normalize_edge_df(df: pd.DataFrame, left_col: str, right_col: str, weight_col: str) -> pd.DataFrame:
    out = df[[left_col, right_col, weight_col]].copy()
    out.columns = ["left", "right", "weight"]

    # clean
    out["left"] = out["left"].astype(str).str.strip()
    out["right"] = out["right"].astype(str).str.strip()
    out["weight"] = pd.to_numeric(out["weight"], errors="coerce").fillna(1).astype(float)

    out = out[(out["left"] != "") & (out["right"] != "")]
    out = out[out["weight"] > 0]

    # aggregate duplicate edges
    out = out.groupby(["left", "right"], as_index=False)["weight"].sum()
    return out


def build_bipartite_graph(edges: pd.DataFrame, min_edge_weight: float = 1.0) -> nx.Graph:
    edges2 = edges[edges["weight"] >= min_edge_weight].copy()
    G = nx.Graph()

    # Add nodes and edges
    left_nodes = sorted(edges2["left"].unique().tolist())
    right_nodes = sorted(edges2["right"].unique().tolist())

    for n in left_nodes:
        G.add_node(f"L::{n}", label=n, bipartite="left", group="left")

    for n in right_nodes:
        G.add_node(f"R::{n}", label=n, bipartite="right", group="right")

    for row in edges2.itertuples(index=False):
        u = f"L::{row.left}"
        v = f"R::{row.right}"
        w = float(row.weight)
        G.add_edge(u, v, weight=w)

    return G


def trim_graph_nodes(G: nx.Graph, max_nodes: int) -> nx.Graph:
    """
    If graph is too large, keep highest-degree nodes (and induced subgraph).
    """
    if G.number_of_nodes() <= max_nodes:
        return G

    deg = dict(G.degree())
    keep = sorted(deg, key=lambda n: deg[n], reverse=True)[:max_nodes]
    return G.subgraph(keep).copy()


def compute_node_sizes(G: nx.Graph, min_size: float = 8.0, max_size: float = 40.0) -> Dict[str, float]:
    deg = dict(G.degree())
    if not deg:
        return {}
    dvals = list(deg.values())
    dmin, dmax = min(dvals), max(dvals)

    sizes = {}
    for n, d in deg.items():
        if dmax == dmin:
            s = (min_size + max_size) / 2.0
        else:
            t = (d - dmin) / (dmax - dmin)
            s = min_size + t * (max_size - min_size)
        sizes[n] = s
    return sizes


def short_label(s: str, max_len: int = 18) -> str:
    s = s.strip()
    if len(s) <= max_len:
        return s
    return s[: max_len - 1] + "…"


# ---------------------------
# Static rendering
# ---------------------------

def render_static_graphviz(G: nx.Graph, out_png: Path, out_svg: Path,
                           rankdir: str = "LR") -> bool:
    """
    Render with Graphviz via pydot. Returns True if success, else False.
    """
    try:
        import pydot  # noqa: F401
        from networkx.drawing.nx_pydot import to_pydot
    except Exception as e:
        _log(f"[static/graphviz] pydot unavailable: {e}")
        return False

    try:
        H = G.copy()

        # style attributes
        sizes = compute_node_sizes(H, min_size=0.35, max_size=1.1)  # inches-ish for graphviz width/height
        for n, d in H.nodes(data=True):
            label_full = str(d.get("label", n))
            lbl = short_label(label_full, 20)
            bip = d.get("bipartite", "left")

            d["label"] = lbl
            d["tooltip"] = label_full
            d["shape"] = "ellipse" if bip == "left" else "box"
            d["style"] = "filled"
            d["fillcolor"] = "lightblue" if bip == "left" else "lightgoldenrod1"
            d["fontsize"] = "10"
            d["fixedsize"] = "false"  # avoid label-too-small warnings as much as possible
            d["width"] = str(max(0.3, sizes.get(n, 0.4)))
            d["height"] = str(max(0.2, sizes.get(n, 0.25)))

        for u, v, d in H.edges(data=True):
            w = float(d.get("weight", 1.0))
            pen = 1.0 + math.log1p(w)
            d["penwidth"] = f"{pen:.2f}"
            d["label"] = ""  # keep uncluttered

        P = to_pydot(H)
        P.set_rankdir(rankdir)
        P.set_splines("true")
        P.set_overlap("false")
        P.set_outputorder("edgesfirst")
        P.set_concentrate("true")
        P.set_nodesep("0.25")
        P.set_ranksep("0.6")
        P.set_bgcolor("white")

        P.write_png(str(out_png))
        P.write_svg(str(out_svg))
        _log(f"[ok] {out_png}")
        _log(f"[ok] {out_svg}")
        return True

    except Exception as e:
        _log(f"[static/graphviz] failed: {e}")
        return False


def render_static_fallback_matplotlib(G: nx.Graph, out_png: Path) -> bool:
    """
    Fallback static PNG renderer (spring layout).
    """
    try:
        H = G.copy()
        sizes = compute_node_sizes(H, min_size=40, max_size=500)

        # deterministic layout
        pos = nx.spring_layout(H, k=0.9 / max(1, math.sqrt(H.number_of_nodes())), iterations=200, seed=42)

        left_nodes = [n for n, d in H.nodes(data=True) if d.get("bipartite") == "left"]
        right_nodes = [n for n, d in H.nodes(data=True) if d.get("bipartite") == "right"]

        plt.figure(figsize=(18, 12), dpi=180)
        # edges
        edge_widths = []
        for _, _, d in H.edges(data=True):
            w = float(d.get("weight", 1.0))
            edge_widths.append(0.2 + math.log1p(w) * 0.6)
        nx.draw_networkx_edges(H, pos, width=edge_widths, alpha=0.2)

        # nodes
        nx.draw_networkx_nodes(
            H, pos, nodelist=left_nodes,
            node_size=[sizes[n] for n in left_nodes],
            alpha=0.9
        )
        nx.draw_networkx_nodes(
            H, pos, nodelist=right_nodes,
            node_size=[sizes[n] for n in right_nodes],
            alpha=0.9
        )

        # selective labels: top degree nodes only
        deg = dict(H.degree())
        top_n = 120 if H.number_of_nodes() > 300 else H.number_of_nodes()
        top_nodes = set(sorted(deg, key=lambda n: deg[n], reverse=True)[:top_n])
        labels = {n: short_label(str(H.nodes[n].get("label", n)), 20) for n in top_nodes}
        nx.draw_networkx_labels(H, pos, labels=labels, font_size=7)

        plt.axis("off")
        plt.tight_layout()
        plt.savefig(out_png, bbox_inches="tight")
        plt.close()
        _log(f"[ok] {out_png}")
        return True

    except Exception as e:
        _log(f"[static/fallback] failed: {e}")
        return False


# ---------------------------
# Interactive rendering
# ---------------------------

def render_pyvis_html(G: nx.Graph, out_html: Path) -> bool:
    """
    Robust PyVis renderer (avoids net.show() notebook path).
    """
    try:
        from pyvis.network import Network
    except Exception as e:
        _log(f"[interactive] pyvis unavailable: {e}")
        return False

    try:
        H = G.copy()
        net = Network(
            height="920px",
            width="100%",
            bgcolor="#ffffff",
            font_color="#111111",
            notebook=False,
            directed=False,
        )

        # physics tuned for large bipartite graphs
        net.barnes_hut(
            gravity=-15000,
            central_gravity=0.12,
            spring_length=180,
            spring_strength=0.018,
            damping=0.11,
            overlap=0.15,
        )

        node_sizes = compute_node_sizes(H, min_size=8, max_size=36)

        for n, d in H.nodes(data=True):
            bip = d.get("bipartite", "left")
            label_full = str(d.get("label", n))
            label_short = short_label(label_full, 24)
            degree = H.degree(n)

            # Keep labels readable in browser; full text in tooltip
            title = (
                f"<b>{label_full}</b><br>"
                f"type: {bip}<br>"
                f"degree: {degree}"
            )

            net.add_node(
                n,
                label=label_short,
                title=title,
                group="Case/Cluster" if bip == "left" else "Font/Artifact",
                value=float(node_sizes.get(n, 10)),
                shape="dot" if bip == "left" else "square",
            )

        for u, v, d in H.edges(data=True):
            w = float(d.get("weight", 1.0))
            net.add_edge(
                u, v,
                value=max(1.0, math.log1p(w) * 2.0),
                title=f"weight: {w:g}",
            )

        # Optional control buttons in generated HTML
        net.show_buttons(filter_=["physics"])

        # CRITICAL: avoid net.show() (it can force notebook=True in some pyvis versions)
        net.write_html(str(out_html), open_browser=False, notebook=False, local=True)

        _log(f"[ok] {out_html}")
        return True

    except Exception as e:
        _log(f"[interactive] failed: {e}")
        return False


# ---------------------------
# Exports
# ---------------------------

def export_nodes_edges(G: nx.Graph, out_nodes_csv: Path, out_edges_csv: Path) -> None:
    nodes_rows = []
    for n, d in G.nodes(data=True):
        nodes_rows.append({
            "node_id": n,
            "label": d.get("label", n),
            "bipartite": d.get("bipartite", ""),
            "degree": G.degree(n),
        })
    pd.DataFrame(nodes_rows).sort_values(["bipartite", "degree"], ascending=[True, False]).to_csv(out_nodes_csv, index=False)

    edges_rows = []
    for u, v, d in G.edges(data=True):
        edges_rows.append({
            "source": u,
            "target": v,
            "weight": d.get("weight", 1),
        })
    pd.DataFrame(edges_rows).sort_values("weight", ascending=False).to_csv(out_edges_csv, index=False)

    _log(f"[ok] {out_nodes_csv}")
    _log(f"[ok] {out_edges_csv}")


def write_summary_json(path: Path, payload: dict) -> None:
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    _log(f"[ok] {path}")


# ---------------------------
# Main
# ---------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Run01 one-shot visual graph generator.")
    p.add_argument("--input", type=Path,
                   default=Path("reports/10_font_tracking/10r1_edges.csv"),
                   help="Input edge CSV path.")
    p.add_argument("--out_dir", type=Path,
                   default=Path("reports/10_font_tracking"),
                   help="Output directory.")
    p.add_argument("--left_col", type=str, default=None,
                   help="Left/source column (e.g., cluster/person).")
    p.add_argument("--right_col", type=str, default=None,
                   help="Right/target column (e.g., font_hash/term).")
    p.add_argument("--weight_col", type=str, default=None,
                   help="Edge weight column.")
    p.add_argument("--min_edge_weight", type=float, default=1.0,
                   help="Minimum edge weight to keep.")
    p.add_argument("--max_nodes", type=int, default=1400,
                   help="Max nodes retained (highest-degree trim if exceeded).")
    p.add_argument("--interactive", action="store_true",
                   help="Produce interactive HTML.")
    p.add_argument("--static", action="store_true",
                   help="Produce static PNG/SVG.")
    p.add_argument("--rankdir", type=str, default="LR", choices=["LR", "TB", "RL", "BT"],
                   help="Graphviz rank direction for static outputs.")
    p.add_argument("--prefix", type=str, default="10r1_graph_bipartite",
                   help="Output file prefix.")
    return p.parse_args()


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

    # If neither flag provided, do both
    do_static = args.static or (not args.static and not args.interactive)
    do_interactive = args.interactive or (not args.static and not args.interactive)

    _safe_mkdir(args.out_dir)

    if not args.input.exists():
        raise FileNotFoundError(f"Input CSV not found: {args.input}")

    _log(f"[read] {args.input}")
    df = pd.read_csv(args.input)

    left_col, right_col, weight_col = detect_columns(df, args.left_col, args.right_col, args.weight_col)
    _log(f"[cols] left={left_col} right={right_col} weight={weight_col}")

    edges = normalize_edge_df(df, left_col, right_col, weight_col)
    _log(f"[edges] normalized rows={len(edges):,}")

    G = build_bipartite_graph(edges, min_edge_weight=args.min_edge_weight)
    _log(f"[graph] nodes={G.number_of_nodes():,} edges={G.number_of_edges():,}")

    # Trim if too large
    G = trim_graph_nodes(G, args.max_nodes)
    _log(f"[graph] after trim nodes={G.number_of_nodes():,} edges={G.number_of_edges():,}")

    out_png = args.out_dir / f"{args.prefix}.png"
    out_svg = args.out_dir / f"{args.prefix}.svg"
    out_html = args.out_dir / f"{args.prefix}_interactive.html"
    out_nodes = args.out_dir / f"{args.prefix}__nodes.csv"
    out_edges = args.out_dir / f"{args.prefix}__edges.csv"
    out_summary = args.out_dir / f"{args.prefix}__summary.json"

    static_ok = False
    interactive_ok = False

    if do_static:
        _log("[static] attempting graphviz...")
        static_ok = render_static_graphviz(G, out_png, out_svg, rankdir=args.rankdir)
        if not static_ok:
            _log("[static] graphviz failed/unavailable; falling back to matplotlib PNG.")
            fallback_ok = render_static_fallback_matplotlib(G, out_png)
            static_ok = fallback_ok
            # no svg in fallback

    if do_interactive:
        _log("[interactive] rendering pyvis html...")
        interactive_ok = render_pyvis_html(G, out_html)

    export_nodes_edges(G, out_nodes, out_edges)

    summary = {
        "script": Path(__file__).name,
        "timestamp_utc": datetime.utcnow().isoformat() + "Z",
        "input_csv": str(args.input),
        "output_dir": str(args.out_dir),
        "columns": {
            "left_col": left_col,
            "right_col": right_col,
            "weight_col": weight_col,
        },
        "params": {
            "min_edge_weight": args.min_edge_weight,
            "max_nodes": args.max_nodes,
            "rankdir": args.rankdir,
            "do_static": do_static,
            "do_interactive": do_interactive,
        },
        "graph": {
            "nodes": G.number_of_nodes(),
            "edges": G.number_of_edges(),
            "left_nodes": sum(1 for _, d in G.nodes(data=True) if d.get("bipartite") == "left"),
            "right_nodes": sum(1 for _, d in G.nodes(data=True) if d.get("bipartite") == "right"),
            "connected_components": nx.number_connected_components(G) if G.number_of_nodes() else 0,
            "density": nx.density(G) if G.number_of_nodes() > 1 else 0.0,
        },
        "outputs": {
            "static_png": str(out_png) if do_static else None,
            "static_svg": str(out_svg) if do_static else None,
            "interactive_html": str(out_html) if do_interactive else None,
            "nodes_csv": str(out_nodes),
            "edges_csv": str(out_edges),
        },
        "status": {
            "static_ok": static_ok if do_static else None,
            "interactive_ok": interactive_ok if do_interactive else None,
        }
    }
    write_summary_json(out_summary, summary)

    _log("\n=== DONE ===")
    if do_static:
        _log(f"Static PNG: {out_png}")
        _log(f"Static SVG: {out_svg} (Graphviz only)")
    if do_interactive:
        _log(f"Interactive HTML: {out_html}")
    _log(f"Node export: {out_nodes}")
    _log(f"Edge export: {out_edges}")
    _log(f"Summary JSON: {out_summary}")

    # exit non-zero only if requested outputs failed
    if do_static and not static_ok:
        _log("[warn] static output failed.")
    if do_interactive and not interactive_ok:
        _log("[warn] interactive output failed.")

    if (do_static and not static_ok) or (do_interactive and not interactive_ok):
        sys.exit(2)


if __name__ == "__main__":
    main()