#!/usr/bin/env python3
"""WP44 — RQ4/RQ5 category-level citation flows and co-classification (after taxonomy/classify.py).

Inputs: cit_edges_internal (knowledge_graph/01), classification tables (work_categories, work_primary), corpus_works.
Outputs (results/graph/): category_flows_yearly.csv (citing year × citing primary category × cited primary category, counts and
row shares), area_flows_yearly.csv (same at area level), cross_area_share_yearly.csv (share of citations leaving the citing
work's area, by citing area), category_cooccurrence_yearly.csv (works carrying both categories), category_counts_yearly.csv
(works per category per year, multi-label and primary), modularity_yearly.csv (modularity of the yearly category citation
network with categories as communities and with areas as communities — Newman–Girvan Q on the directed weighted graph)."""
import argparse, os, json, duckdb, yaml, numpy as np, pandas as pd


def modularity(W, groups):
    """Newman–Girvan modularity for a weighted directed graph with adjacency W (i→j) and community labels."""
    m = W.sum(); kout = W.sum(1); kin = W.sum(0)
    g = np.asarray(groups); same = g[:, None] == g[None, :]
    return float(((W - np.outer(kout, kin) / m) * same).sum() / m)


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--ontology", required=True); a = ap.parse_args(); R = a.root
    out = os.path.join(R, "results/graph"); os.makedirs(out, exist_ok=True)
    ont = yaml.safe_load(open(a.ontology)); cat_area = {c["id"]: c["area"] for c in ont["categories"]}; cats = [c["id"] for c in ont["categories"]]
    con = duckdb.connect(os.path.join(R, "data/processed/wp44.duckdb")); con.execute("SET threads TO 24; SET memory_limit='70GB';")
    con.execute(f"CREATE OR REPLACE VIEW work_primary AS SELECT * FROM read_parquet('{R}/data/processed/classification/primary/*.parquet')")
    con.execute(f"CREATE OR REPLACE VIEW work_categories AS SELECT * FROM read_parquet('{R}/data/processed/classification/cats/*.parquet')")
    con.execute("CREATE OR REPLACE TEMP TABLE area_map AS SELECT * FROM (VALUES " + ",".join(f"('{c}','{cat_area[c]}')" for c in cats) + ") t(category_id, area)")
    con.execute("""CREATE OR REPLACE TABLE category_flows_yearly AS
        SELECT e.citing_year, pa.primary_category AS citing_cat, pb.primary_category AS cited_cat, count(*) n
        FROM cit_edges_internal e JOIN work_primary pa ON pa.work_id = e.citing_id JOIN work_primary pb ON pb.work_id = e.cited_id GROUP BY 1,2,3""")
    con.execute(f"COPY (SELECT *, n*1.0/sum(n) OVER (PARTITION BY citing_year, citing_cat) AS row_share FROM category_flows_yearly ORDER BY 1,2,4 DESC) TO '{out}/category_flows_yearly.csv' (HEADER)")
    con.execute(f"""COPY (SELECT f.citing_year, a1.area AS citing_area, a2.area AS cited_area, sum(n) n FROM category_flows_yearly f JOIN area_map a1 ON a1.category_id = f.citing_cat JOIN area_map a2 ON a2.category_id = f.cited_cat
                 GROUP BY 1,2,3 ORDER BY 1,2,4 DESC) TO '{out}/area_flows_yearly.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT citing_year, citing_area, sum(n) n, sum(CASE WHEN citing_area <> cited_area THEN n ELSE 0 END) n_cross, round(sum(CASE WHEN citing_area <> cited_area THEN n ELSE 0 END)*1.0/sum(n),4) share_cross
                 FROM read_csv_auto('{out}/area_flows_yearly.csv') GROUP BY 1,2 ORDER BY 1,2) TO '{out}/cross_area_share_yearly.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT c.publication_year AS year, a.category_id AS cat_a, b.category_id AS cat_b, count(*) n FROM work_categories a JOIN work_categories b ON a.work_id = b.work_id AND a.category_id < b.category_id
                 JOIN corpus_works c ON c.id = a.work_id WHERE c.eligible GROUP BY 1,2,3 ORDER BY 1,4 DESC) TO '{out}/category_cooccurrence_yearly.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT c.publication_year AS year, c.period, k.category_id, count(*) n_multilabel, sum((p.primary_category = k.category_id)::int) n_primary
                 FROM work_categories k JOIN corpus_works c ON c.id = k.work_id JOIN work_primary p ON p.work_id = k.work_id WHERE c.eligible GROUP BY 1,2,3 ORDER BY 1,3) TO '{out}/category_counts_yearly.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT c.first_month, c.period, k.category_id, count(*) n_multilabel, sum((p.primary_category = k.category_id)::int) n_primary, sum((c.date_precision='month')::int) n_month_precision
                 FROM work_categories k JOIN corpus_works c ON c.id = k.work_id JOIN work_primary p ON p.work_id = k.work_id WHERE c.eligible GROUP BY 1,2,3 ORDER BY 1,3) TO '{out}/category_counts_monthly.csv' (HEADER)""")
    rows = []
    df = pd.read_csv(f"{out}/category_flows_yearly.csv")
    idx = {c: i for i, c in enumerate(cats)}
    for y, g in df.groupby("citing_year"):
        W = np.zeros((len(cats), len(cats)))
        for r in g.itertuples():
            if r.citing_cat in idx and r.cited_cat in idx:
                W[idx[r.citing_cat], idx[r.cited_cat]] += r.n
        if W.sum() == 0:
            continue
        rows.append({"year": int(y), "edges": float(W.sum()), "Q_categories": modularity(W, cats), "Q_areas": modularity(W, [cat_area[c] for c in cats]),
                     "self_share": float(np.trace(W) / W.sum())})
    pd.DataFrame(rows).to_csv(f"{out}/modularity_yearly.csv", index=False); print(pd.DataFrame(rows))


if __name__ == "__main__":
    main()