#!/usr/bin/env python3
"""
survey-phase0.py — Deterministic, FREE (no LLM API) Phase-0 analytics for the
quantitative survey paper of the IETF AI/agent Internet-Draft landscape.

CORPUS: "clean IETF corpus" = drafts where source='ietf'
        AND (ratings.false_positive=0 OR NULL). Join drafts d with ratings r
        on d.name=r.draft_name. This is 524 drafts.

Computes:
  1. Author / working-group concentration
  2. Embedding overlap / redundancy (pairwise cosine, nomic-embed-text 768-d)
  3. Category coverage map (primary category from ratings.categories JSON)

Read-only on data/drafts.db. Writes data/reports/survey-phase0.md and prints a
headline summary to stdout. Pure stdlib + numpy only.
"""
import json
import sqlite3
import statistics
from collections import Counter, defaultdict
from datetime import datetime, timezone

import numpy as np

DB = "data/drafts.db"
OUT = "data/reports/survey-phase0.md"
CLEAN = "d.source='ietf' AND (r.false_positive=0 OR r.false_positive IS NULL)"

# Well-known IETF Datatracker group IDs. ID 1027 is the "none" pseudo-group used
# for individually-submitted drafts (not adopted by a WG). Others are real WGs/RGs.
GROUP_ID_NAMES = {
    "1027": "(individual / no WG)",
}


def conn():
    c = sqlite3.connect(DB)
    c.row_factory = sqlite3.Row
    return c


def gid(group_uri):
    """Extract numeric group id from a group_uri like /api/v1/group/group/1027/ ."""
    if not group_uri:
        return None
    parts = [p for p in group_uri.strip("/").split("/") if p]
    return parts[-1] if parts else None


# ---------------------------------------------------------------------------
# 1. AUTHOR / WORKING-GROUP CONCENTRATION
# ---------------------------------------------------------------------------
def section_authors(c):
    total_clean = c.execute(
        f"SELECT COUNT(*) FROM drafts d JOIN ratings r ON d.name=r.draft_name WHERE {CLEAN}"
    ).fetchone()[0]

    # Top authors by # of clean-IETF drafts. Join via draft_authors.person_id -> authors.person_id.
    author_counts = c.execute(
        f"""
        SELECT a.person_id, a.name,
               COUNT(DISTINCT d.name) AS n
        FROM drafts d
        JOIN ratings r ON d.name = r.draft_name
        JOIN draft_authors da ON da.draft_name = d.name
        JOIN authors a ON a.person_id = da.person_id
        WHERE {CLEAN}
        GROUP BY a.person_id
        ORDER BY n DESC, a.name ASC
        """
    ).fetchall()

    total_authors = len(author_counts)
    top15 = author_counts[:15]
    top10 = author_counts[:10]

    # Concentration: share of (author, draft) memberships covered by top-10 authors,
    # and share of distinct drafts touched by at least one top-10 author.
    total_memberships = sum(a["n"] for a in author_counts)
    top10_memberships = sum(a["n"] for a in top10)

    top10_ids = [a["person_id"] for a in top10]
    if top10_ids:
        placeholders = ",".join("?" * len(top10_ids))
        top10_drafts = c.execute(
            f"""
            SELECT COUNT(DISTINCT d.name)
            FROM drafts d
            JOIN ratings r ON d.name = r.draft_name
            JOIN draft_authors da ON da.draft_name = d.name
            WHERE {CLEAN} AND da.person_id IN ({placeholders})
            """,
            top10_ids,
        ).fetchone()[0]
    else:
        top10_drafts = 0

    drafts_with_authors = c.execute(
        f"""
        SELECT COUNT(DISTINCT d.name)
        FROM drafts d
        JOIN ratings r ON d.name = r.draft_name
        JOIN draft_authors da ON da.draft_name = d.name
        WHERE {CLEAN}
        """
    ).fetchone()[0]

    # Working groups. drafts.group text column is unpopulated for this corpus, so
    # report it (per spec) AND the richer group_uri-based breakdown.
    raw_group = c.execute(
        f"""
        SELECT d."group" AS g, COUNT(*) AS n
        FROM drafts d JOIN ratings r ON d.name=r.draft_name
        WHERE {CLEAN}
        GROUP BY d."group" ORDER BY n DESC
        """
    ).fetchall()
    no_group_raw = c.execute(
        f"""
        SELECT COUNT(*) FROM drafts d JOIN ratings r ON d.name=r.draft_name
        WHERE {CLEAN} AND (d."group" IS NULL OR d."group"='' OR d."group"='none')
        """
    ).fetchone()[0]

    uri_rows = c.execute(
        f"""
        SELECT d.group_uri AS uri, COUNT(*) AS n
        FROM drafts d JOIN ratings r ON d.name=r.draft_name
        WHERE {CLEAN}
        GROUP BY d.group_uri ORDER BY n DESC
        """
    ).fetchall()
    wg_breakdown = []
    no_wg_uri = 0  # drafts in the "individual / no WG" pseudo-group or null uri
    for row in uri_rows:
        g = gid(row["uri"])
        label = GROUP_ID_NAMES.get(g, f"group/{g}" if g else "(no group_uri)")
        if g == "1027" or g is None:
            no_wg_uri += row["n"]
        wg_breakdown.append((label, g, row["n"]))

    return {
        "total_clean": total_clean,
        "total_authors": total_authors,
        "top15": top15,
        "total_memberships": total_memberships,
        "top10_memberships": top10_memberships,
        "top10_drafts": top10_drafts,
        "drafts_with_authors": drafts_with_authors,
        "raw_group": raw_group,
        "no_group_raw": no_group_raw,
        "wg_breakdown": wg_breakdown,
        "no_wg_uri": no_wg_uri,
        "distinct_uris": len(uri_rows),
    }


# ---------------------------------------------------------------------------
# 2. EMBEDDING OVERLAP / REDUNDANCY
# ---------------------------------------------------------------------------
def load_embeddings(c):
    """Return (names, matrix L2-normalized, titles dict) for clean corpus."""
    rows = c.execute(
        f"""
        SELECT d.name AS name, d.title AS title, em.vector AS vec
        FROM drafts d
        JOIN ratings r ON d.name = r.draft_name
        JOIN embeddings em ON em.draft_name = d.name
        WHERE {CLEAN}
        ORDER BY d.name
        """
    ).fetchall()
    names, titles, vectors = [], {}, []
    for row in rows:
        v = row["vec"]
        # Vector serialization: float32 BLOB (verified). Fall back to JSON if needed.
        try:
            arr = np.frombuffer(v, dtype=np.float32)
            if arr.size == 0:
                raise ValueError
        except Exception:
            arr = np.asarray(json.loads(v), dtype=np.float32)
        names.append(row["name"])
        titles[row["name"]] = row["title"]
        vectors.append(arr)
    mat = np.vstack(vectors).astype(np.float64)
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    mat = mat / norms
    return names, mat, titles


def section_embeddings(c):
    names, mat, titles = load_embeddings(c)
    n = len(names)
    sim = mat @ mat.T  # cosine since rows are unit-normalized
    iu, ju = np.triu_indices(n, k=1)
    off = sim[iu, ju]

    dist = {
        "n_drafts": n,
        "n_pairs": int(off.size),
        "dim": mat.shape[1],
        "mean": float(np.mean(off)),
        "median": float(np.median(off)),
        "p90": float(np.percentile(off, 90)),
        "p99": float(np.percentile(off, 99)),
        "max": float(np.max(off)),
    }

    # Top 20 most-similar pairs
    order = np.argsort(off)[::-1][:20]
    top_pairs = []
    for k in order:
        a, b = names[iu[k]], names[ju[k]]
        top_pairs.append((a, b, float(off[k]), titles.get(a, ""), titles.get(b, "")))

    # Drafts with at least one near-duplicate (cosine > 0.9)
    thr = 0.9
    near = sim > thr
    np.fill_diagonal(near, False)
    has_near = int(np.sum(near.any(axis=1)))
    n_near_pairs = int(np.sum(off > thr))

    dist["near_dup_drafts"] = has_near
    dist["near_dup_pairs"] = n_near_pairs
    return dist, top_pairs


# ---------------------------------------------------------------------------
# 3. CATEGORY COVERAGE MAP
# ---------------------------------------------------------------------------
def section_categories(c):
    rows = c.execute(
        f"""
        SELECT d.name AS name, r.categories AS cats,
               r.novelty, r.maturity, r.overlap, r.momentum, r.relevance
        FROM drafts d JOIN ratings r ON d.name = r.draft_name
        WHERE {CLEAN}
        """
    ).fetchall()

    primary_counts = Counter()
    multi_cat = 0
    cat_dims = defaultdict(lambda: {"novelty": [], "maturity": [], "overlap": [],
                                    "momentum": [], "relevance": [], "composite": []})
    n_total = 0
    n_no_cat = 0

    for row in rows:
        n_total += 1
        cats = []
        if row["cats"]:
            try:
                cats = [x for x in json.loads(row["cats"]) if x]
            except Exception:
                cats = []
        if not cats:
            primary = "(uncategorized)"
            n_no_cat += 1
        else:
            primary = cats[0]
            if len(cats) > 1:
                multi_cat += 1
        primary_counts[primary] += 1

        dims = cat_dims[primary]
        vals = []
        for d in ("novelty", "maturity", "overlap", "momentum", "relevance"):
            if row[d] is not None:
                dims[d].append(row[d])
                vals.append(row[d])
        if vals:
            dims["composite"].append(statistics.mean(vals))

    # Build per-category summary
    cat_summary = []
    for cat, cnt in primary_counts.most_common():
        dims = cat_dims[cat]
        def m(key):
            return round(statistics.mean(dims[key]), 2) if dims[key] else None
        cat_summary.append({
            "cat": cat, "n": cnt,
            "relevance": m("relevance"),
            "composite": m("composite"),
            "novelty": m("novelty"), "maturity": m("maturity"),
            "overlap": m("overlap"), "momentum": m("momentum"),
        })

    sparsest = sorted([s for s in cat_summary if s["cat"] != "(uncategorized)"],
                      key=lambda s: s["n"])[:3]

    return {
        "n_total": n_total,
        "primary_counts": primary_counts,
        "cat_summary": cat_summary,
        "multi_cat": multi_cat,
        "n_no_cat": n_no_cat,
        "sparsest": sparsest,
        "sum_check": sum(primary_counts.values()),
    }


# ---------------------------------------------------------------------------
# REPORT
# ---------------------------------------------------------------------------
def md_table(headers, rows):
    out = ["| " + " | ".join(headers) + " |",
           "| " + " | ".join("---" for _ in headers) + " |"]
    for r in rows:
        out.append("| " + " | ".join(str(x) for x in r) + " |")
    return "\n".join(out)


def build_report(au, em, cat):
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    L = []
    A = L.append
    A("# Survey Phase 0 — Quantitative IETF AI/Agent Draft Landscape\n")
    A(f"_Generated {ts} by `scripts/survey-phase0.py` (deterministic, no LLM API calls)._\n")
    A("**Corpus definition (\"clean IETF corpus\"):** `source='ietf' AND (ratings.false_positive=0 OR NULL)`, "
      f"joining `drafts d` with `ratings r` on `d.name=r.draft_name`. **N = {au['total_clean']} drafts.**\n")

    # 1
    A("## 1. Author / Working-Group Concentration\n")
    A(f"- Distinct authors across the clean corpus: **{au['total_authors']}**")
    A(f"- Clean drafts with at least one resolved author in `draft_authors`: **{au['drafts_with_authors']}** "
      f"of {au['total_clean']} "
      f"({100*au['drafts_with_authors']/au['total_clean']:.1f}%); the remainder have no rows in `draft_authors`.")
    A(f"- Total (author, draft) authorship memberships: **{au['total_memberships']}**\n")

    A("### Top 15 authors by clean-IETF draft count\n")
    A(md_table(["#", "Author", "Drafts"],
               [(i + 1, a["name"], a["n"]) for i, a in enumerate(au["top15"])]))
    A("")
    share_mem = 100 * au["top10_memberships"] / au["total_memberships"] if au["total_memberships"] else 0
    share_drafts = 100 * au["top10_drafts"] / au["total_clean"]
    A("### Concentration (top 10 authors)\n")
    A(f"- Top-10 authors account for **{au['top10_memberships']} of {au['total_memberships']} authorship "
      f"memberships ({share_mem:.1f}%)**.")
    A(f"- Top-10 authors appear on **{au['top10_drafts']} of {au['total_clean']} distinct clean drafts "
      f"({share_drafts:.1f}%)** (a draft is counted once even if multiple top-10 authors appear on it).\n")

    A("### Working groups\n")
    A("The `drafts.group` text column is effectively unpopulated for this corpus "
      f"(only {au['no_group_raw']} drafts resolve to no/empty/'none' group, and the column has no real WG names). "
      "The meaningful grouping signal is `drafts.group_uri` (Datatracker group IDs).\n")
    A(f"- Drafts with no working group (group/1027 \"individual\" pseudo-group or null `group_uri`): "
      f"**{au['no_wg_uri']}** of {au['total_clean']}.")
    A(f"- Distinct `group_uri` values: **{au['distinct_uris']}**.\n")
    A("Top 15 groups by `group_uri` (clean-IETF draft count):\n")
    A(md_table(["#", "Group (Datatracker ID / label)", "Drafts"],
               [(i + 1, lbl, n) for i, (lbl, g, n) in enumerate(au["wg_breakdown"][:15])]))
    A("")
    A("> Note: ID 1027 is the Datatracker \"none\" pseudo-group for individually submitted drafts; "
      "it is not a real working group. Real WG/RG IDs are the remaining rows.\n")

    # 2
    A("## 2. Embedding Overlap / Redundancy\n")
    A(f"- Embedding model: `nomic-embed-text`, dimension **{em['dim']}**, stored as float32 BLOB.")
    A(f"- Drafts embedded (clean corpus): **{em['n_drafts']}** (100% coverage); "
      f"pairwise comparisons: **{em['n_pairs']:,}**.\n")
    A("### Pairwise cosine similarity (off-diagonal)\n")
    A(md_table(["Mean", "Median", "p90", "p99", "Max"],
               [(f"{em['mean']:.4f}", f"{em['median']:.4f}", f"{em['p90']:.4f}",
                 f"{em['p99']:.4f}", f"{em['max']:.4f}")]))
    A("")
    A(f"- Near-duplicate pairs (cosine > 0.9): **{em['near_dup_pairs']}**.")
    A(f"- Drafts with at least one near-duplicate (cosine > 0.9): **{em['near_dup_drafts']}** "
      f"of {em['n_drafts']} ({100*em['near_dup_drafts']/em['n_drafts']:.1f}%).\n")
    A("### Top 20 most-similar draft pairs\n")
    rows = []
    for i, (a, b, cos, ta, tb) in enumerate(em["pairs"]):
        rows.append((i + 1, f"{cos:.4f}", f"`{a}`<br>`{b}`",
                     f"{(ta or '')[:70]}<br>{(tb or '')[:70]}"))
    A(md_table(["#", "Cosine", "Draft A / Draft B", "Title A / Title B"], rows))
    A("")

    # 3
    A("## 3. Category Coverage Map\n")
    A(f"Primary category = first element of the `ratings.categories` JSON array. "
      f"Sum check: **{cat['sum_check']} = {au['total_clean']}** "
      f"({'OK' if cat['sum_check'] == au['total_clean'] else 'MISMATCH'}).\n")
    A(f"- Drafts carrying more than one category: **{cat['multi_cat']}** of {cat['n_total']} "
      f"({100*cat['multi_cat']/cat['n_total']:.1f}%).")
    A(f"- Drafts with no category (uncategorized): **{cat['n_no_cat']}**.\n")
    A("### Primary-category distribution + mean scores\n")
    A("Scores are 1-5 integers from `ratings`. `composite` = mean of the 5 dimensions "
      "(novelty, maturity, overlap, momentum, relevance) per draft, averaged over the category.\n")
    rows = []
    for s in cat["cat_summary"]:
        rows.append((s["cat"], s["n"], f"{100*s['n']/cat['n_total']:.1f}%",
                     s["relevance"], s["composite"], s["novelty"], s["maturity"],
                     s["overlap"], s["momentum"]))
    A(md_table(["Category", "N", "Share", "Relevance", "Composite",
                "Novelty", "Maturity", "Overlap", "Momentum"], rows))
    A("")
    A("### Three sparsest categories (descriptive)\n")
    A("Listed neutrally as the categories with the fewest drafts in this corpus:\n")
    A(md_table(["Category", "N", "Relevance", "Composite"],
               [(s["cat"], s["n"], s["relevance"], s["composite"]) for s in cat["sparsest"]]))
    A("")

    return "\n".join(L) + "\n"


def main():
    c = conn()
    try:
        au = section_authors(c)
        dist, pairs = section_embeddings(c)
        dist["pairs"] = pairs
        cat = section_categories(c)
    finally:
        c.close()

    report = build_report(au, dist, cat)
    import os
    os.makedirs("data/reports", exist_ok=True)
    with open(OUT, "w") as f:
        f.write(report)

    # Internal consistency checks
    assert cat["sum_check"] == au["total_clean"], "category sum != 524"
    share_drafts = 100 * au["top10_drafts"] / au["total_clean"]
    assert share_drafts <= 100, "author draft share > 100%"

    # Headline summary
    print("=" * 60)
    print("SURVEY PHASE 0 — HEADLINE NUMBERS")
    print("=" * 60)
    print(f"Clean IETF corpus:           {au['total_clean']} drafts")
    print(f"Distinct authors:            {au['total_authors']}")
    print(f"Drafts with authors:         {au['drafts_with_authors']} ({100*au['drafts_with_authors']/au['total_clean']:.1f}%)")
    print(f"Top-10 author draft share:   {share_drafts:.1f}%")
    print(f"Drafts with no WG:           {au['no_wg_uri']} (group_uri-based)")
    print(f"Distinct group_uris:         {au['distinct_uris']}")
    print("-" * 60)
    print(f"Embeddings: {dist['n_drafts']} drafts, dim {dist['dim']}, {dist['n_pairs']:,} pairs")
    print(f"Cosine mean/median:          {dist['mean']:.4f} / {dist['median']:.4f}")
    print(f"Cosine p90/p99/max:          {dist['p90']:.4f} / {dist['p99']:.4f} / {dist['max']:.4f}")
    print(f"Near-dup pairs (>0.9):       {dist['near_dup_pairs']}")
    print(f"Drafts w/ near-dup (>0.9):   {dist['near_dup_drafts']}")
    print("-" * 60)
    print(f"Primary categories:          {len(cat['primary_counts'])} (sum={cat['sum_check']})")
    print(f"Multi-category drafts:       {cat['multi_cat']} ({100*cat['multi_cat']/cat['n_total']:.1f}%)")
    print(f"Uncategorized drafts:        {cat['n_no_cat']}")
    top3 = cat["cat_summary"][:3]
    print("Top 3 categories:            " + ", ".join(f"{s['cat']} ({s['n']})" for s in top3))
    print("Sparsest 3 categories:       " + ", ".join(f"{s['cat']} ({s['n']})" for s in cat["sparsest"]))
    print("=" * 60)
    print(f"Report written to {OUT}")


if __name__ == "__main__":
    main()
