#!/usr/bin/env python3
"""WP44 — RQ6 collaboration, institutions, countries (eligible corpus), yearly and by period.
Outputs (results/bibliometrics/): team_size_yearly.csv, countries_yearly.csv (top-N country output + share, with
fractional counting), international_collab_yearly.csv, institution_types_yearly.csv (education/company/government/…),
top_institutions_by_period.csv (full and fractional counts), country_pairs_by_period.csv (edges for the collaboration
network, ≥ 50 joint works), top_sources_by_period.csv (venue concentration inputs), venue_hhi_yearly.csv."""
import argparse, os, duckdb


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); a = ap.parse_args(); R = a.root
    out = os.path.join(R, "results/bibliometrics"); os.makedirs(out, exist_ok=True)
    con = duckdb.connect(os.path.join(R, "data/processed/wp44.duckdb"), read_only=True); con.execute("SET threads TO 20; SET memory_limit='60GB';")
    con.execute("CREATE OR REPLACE TEMP VIEW E AS SELECT * FROM corpus_works WHERE eligible")
    con.execute(f"""COPY (SELECT publication_year AS year, period, count(*) n, round(avg(n_authors),3) mean_authors, quantile_cont(n_authors,0.5) med_authors,
        sum((n_authors=1)::int) single, sum((n_authors>=5)::int) five_plus, sum((n_authors>=10)::int) ten_plus, sum((n_countries>=2)::int) intl,
        sum((n_institutions>=2)::int) multi_inst, sum((n_countries=0)::int) no_country FROM E GROUP BY 1,2 ORDER BY 1) TO '{out}/team_size_yearly.csv' (HEADER)""")
    # country output: full counting (work counted once per country) and fractional (1/n_countries)
    con.execute("""CREATE OR REPLACE TEMP TABLE wc AS SELECT w.id, w.publication_year AS year, w.period, unnest(w.countries) AS country, 1.0/len(w.countries) AS frac FROM E w WHERE len(w.countries) > 0""")
    con.execute(f"""COPY (SELECT year, country, count(*) n_full, round(sum(frac),1) n_frac FROM wc GROUP BY 1,2 ORDER BY 1, 3 DESC) TO '{out}/countries_yearly.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT year, count(*) n, sum((n_countries>=2)::int) n_intl, round(sum((n_countries>=2)::int)*1.0/count(*),4) share_intl FROM E WHERE n_countries>0 GROUP BY 1 ORDER BY 1) TO '{out}/international_collab_yearly.csv' (HEADER)""")
    con.execute("""CREATE OR REPLACE TEMP TABLE wi AS SELECT w.id, w.publication_year AS year, w.period, unnest(w.institution_types) AS itype FROM E w WHERE len(w.institution_types) > 0""")
    con.execute(f"""COPY (SELECT year, itype, count(*) n FROM wi GROUP BY 1,2 ORDER BY 1,3 DESC) TO '{out}/institution_types_yearly.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT year, count(DISTINCT id) n_with_inst, count(DISTINCT CASE WHEN itype='company' THEN id END) n_company,
        round(count(DISTINCT CASE WHEN itype='company' THEN id END)*1.0/count(DISTINCT id),4) share_company FROM wi GROUP BY 1 ORDER BY 1) TO '{out}/industry_share_yearly.csv' (HEADER)""")
    # institutions (from authorships): full and fractional (by distinct institutions per work)
    con.execute("""CREATE OR REPLACE TEMP TABLE winst AS
        SELECT DISTINCT a.work_id, unnest(a.institution_ids) AS inst_id FROM oa_authorships a WHERE a.work_id IN (SELECT id FROM E) AND len(a.institution_ids) > 0""")
    con.execute("""CREATE OR REPLACE TEMP TABLE winst2 AS SELECT w.work_id, w.inst_id, 1.0/count(*) OVER (PARTITION BY w.work_id) AS frac, e.period, e.publication_year AS year FROM winst w JOIN E e ON e.id = w.work_id""")
    con.execute(f"""COPY (SELECT period, inst_id, count(*) n_full, round(sum(frac),1) n_frac FROM winst2 GROUP BY 1,2 HAVING count(*) >= 200 ORDER BY 1, 3 DESC) TO '{out}/top_institutions_by_period.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT a.period, a.country AS c1, b.country AS c2, count(DISTINCT a.id) n FROM wc a JOIN wc b ON a.id = b.id AND a.country < b.country
        GROUP BY 1,2,3 HAVING count(DISTINCT a.id) >= 50 ORDER BY 1, 4 DESC) TO '{out}/country_pairs_by_period.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT period, source_id, any_value(source_name) source_name, any_value(source_type) source_type, count(*) n FROM E WHERE source_id IS NOT NULL
        GROUP BY 1,2 HAVING count(*) >= 100 ORDER BY 1, 5 DESC) TO '{out}/top_sources_by_period.csv' (HEADER)""")
    con.execute(f"""COPY (SELECT year, count(*) n, count(DISTINCT source_id) n_sources, round(sum(s*s),6) hhi, round(max(s),5) top1_share FROM
        (SELECT publication_year AS year, source_id, count(*)*1.0/sum(count(*)) OVER (PARTITION BY publication_year) AS s FROM E WHERE source_id IS NOT NULL GROUP BY 1,2)
        GROUP BY 1 ORDER BY 1) TO '{out}/venue_hhi_yearly.csv' (HEADER)""")
    print(open(f"{out}/team_size_yearly.csv").read()[:800]); print(open(f"{out}/industry_share_yearly.csv").read())


if __name__ == "__main__":
    main()