#!/usr/bin/env python3
"""What each department uses Claude for: Anthropic Economic Index tasks mapped to departments.

Public data only: Anthropic Economic Index, release_2025_02_10, pinned to one commit.
  https://huggingface.co/datasets/Anthropic/EconomicIndex
No API key, no Claude calls. Needs:  pip install duckdb
Run:  python3 departments_claude_usage.py          (measure, run seven checks, write the two CSVs)
      python3 departments_claude_usage.py --fault  (plant faults in the input; checks must catch them)

Method in one line: every O*NET task has a share of Claude conversations. A task that several
occupations list is split equally across them. Each occupation code goes to a department by the
longest matching prefix in DEPARTMENT_BY_PREFIX below. Shares are summed per department.
"""
import sys
from pathlib import Path

import duckdb

COMMIT = "4e5f69fab13a9c3f7d0dcc31a6f893a9818738f6"  # dataset commit; the files below never change under it
BASE = f"https://huggingface.co/datasets/Anthropic/EconomicIndex/resolve/{COMMIT}/release_2025_02_10/"
OUT = Path(__file__).resolve().parent
FAULT = "--fault" in sys.argv

# The judgment in this script: SOC occupation-code prefix -> department. Longest prefix wins.
# Edit this table to disagree with it, then rerun.
DEPARTMENT_BY_PREFIX = {
    "11": "Executive management",
    "11-2011": "Marketing, media & creative",   # advertising and promotions managers
    "11-2021": "Marketing, media & creative",   # marketing managers
    "11-2022": "Sales",                         # sales managers
    "11-2031": "Marketing, media & creative",   # public relations and fundraising managers
    "11-3": "Operations & business support",
    "11-3021": "IT, software & data",           # computer and information systems managers
    "11-3031": "Finance & accounting",          # financial managers
    "11-3121": "Human resources",               # human resources managers
    "13": "Operations & business support",
    "13-1071": "Human resources",               # human resources specialists
    "13-1075": "Human resources",               # labor relations specialists
    "13-1161": "Marketing, media & creative",   # market research analysts and marketing specialists
    "13-2": "Finance & accounting",
    "15": "IT, software & data",
    "17": "Engineering",
    "19": "Research & science",
    "21": "Other occupations",
    "23": "Legal",
    "25": "Education & training",
    "27": "Marketing, media & creative",
    "29": "Healthcare",
    "31": "Healthcare",
    "33": "Other occupations",
    "35": "Other occupations",
    "37": "Other occupations",
    "39": "Other occupations",
    "41": "Sales",
    "43": "Customer support & administration",
    "43-3": "Finance & accounting",              # financial clerks
    "45": "Other occupations",
    "47": "Other occupations",
    "49": "Other occupations",
    "51": "Other occupations",
    "53": "Other occupations",
    # No "55" (military): the release has no military occupation codes.
}
NOT_ATTRIBUTABLE = "No occupation (task recorded as none)"

con = duckdb.connect()
con.execute(f"CREATE TABLE mappings AS SELECT * FROM read_csv('{BASE}onet_task_mappings.csv')")
con.execute(f"CREATE TABLE statements AS SELECT * FROM read_csv('{BASE}onet_task_statements.csv')")
con.execute(f"CREATE TABLE soc AS SELECT * FROM read_csv('{BASE}SOC_Structure.csv')")
con.execute(f"CREATE TABLE bls AS SELECT * FROM read_csv('{BASE}bls_employment_may_2023.csv')")

# Row counts straight from the source files, before anything is changed.
source_mapping_rows = con.execute("SELECT count(*) FROM mappings").fetchone()[0]
source_statement_rows = con.execute("SELECT count(*) FROM statements").fetchone()[0]

if FAULT:
    # Planted faults: five duplicated task rows, and one statement row with an occupation code
    # that no department prefix covers.
    con.execute("INSERT INTO mappings SELECT * FROM mappings ORDER BY task_name LIMIT 5")
    con.execute("""INSERT INTO statements
                   SELECT '99-9999.00', 'Planted Unknown Occupation', 0, Task, 'Core', 0, '07/2014', 'Incumbent'
                   FROM statements ORDER BY Task LIMIT 1""")

con.execute("CREATE TABLE dept_map(prefix VARCHAR, department VARCHAR)")
con.executemany("INSERT INTO dept_map VALUES (?, ?)", list(DEPARTMENT_BY_PREFIX.items()))

con.execute("""
CREATE TABLE occ_dept AS
SELECT DISTINCT code,
       coalesce(first(department ORDER BY length(prefix) DESC), 'UNMAPPED') AS department
FROM (SELECT DISTINCT "O*NET-SOC Code" AS code FROM statements) o
LEFT JOIN dept_map d ON o.code LIKE d.prefix || '%'
GROUP BY code""")

con.execute("""
CREATE TABLE task_occ AS
SELECT DISTINCT lower(trim(Task)) AS task, "O*NET-SOC Code" AS code FROM statements""")

con.execute("""
CREATE TABLE weighted AS
SELECT m.task_name, m.pct, t.code, d.department, t.task AS matched,
       m.pct / greatest(count(t.code) OVER (PARTITION BY m.task_name), 1) AS weight
FROM mappings m
LEFT JOIN task_occ t ON lower(trim(m.task_name)) = t.task
LEFT JOIN occ_dept d ON t.code = d.code""")

con.execute(f"""
CREATE TABLE by_department AS
SELECT coalesce(department, '{NOT_ATTRIBUTABLE}') AS department,
       sum(weight) AS share_pct,
       count(DISTINCT task_name) AS tasks_with_usage,
       count(DISTINCT code) AS occupations
FROM weighted GROUP BY 1""")

con.execute(f"""
CREATE TABLE top_task AS
SELECT department, task_name, task_pct FROM (
  SELECT coalesce(department, '{NOT_ATTRIBUTABLE}') AS department, task_name,
         sum(weight) AS task_pct,
         row_number() OVER (PARTITION BY coalesce(department, '{NOT_ATTRIBUTABLE}')
                            ORDER BY sum(weight) DESC, task_name) AS rk
  FROM weighted GROUP BY 1, 2) WHERE rk = 1""")

result = con.execute(f"""
SELECT b.department, round(b.share_pct, 2) AS share_pct, b.tasks_with_usage, b.occupations,
       t.task_name AS top_task, round(t.task_pct, 3) AS top_task_share_pct
FROM by_department b JOIN top_task t USING (department)
ORDER BY (b.department = '{NOT_ATTRIBUTABLE}'), b.share_pct DESC, b.department""").fetchall()
cols = ["department", "share_pct", "tasks_with_usage", "occupations", "top_task", "top_task_share_pct"]

major = con.execute("""
WITH usage AS (
  SELECT substr(code, 1, 2) AS grp, sum(weight) AS usage_pct FROM weighted WHERE code IS NOT NULL GROUP BY 1),
titles AS (
  SELECT substr("Major Group", 1, 2) AS grp, "SOC or O*NET-SOC 2019 Title" AS title
  FROM soc WHERE "Major Group" IS NOT NULL),
emp AS (SELECT "SOC or O*NET-SOC 2019 Title" AS title,
               100.0 * bls_distribution / sum(bls_distribution) OVER () AS emp_pct FROM bls)
SELECT u.grp AS soc_major_group, t.title, round(u.usage_pct, 2) AS claude_share_pct,
       round(e.emp_pct, 2) AS us_employment_share_pct,
       round(u.usage_pct / e.emp_pct, 2) AS ratio
FROM usage u LEFT JOIN titles t USING (grp) LEFT JOIN emp e USING (title)
ORDER BY u.usage_pct DESC, u.grp""").fetchall()
major_cols = ["soc_major_group", "title", "claude_share_pct", "us_employment_share_pct", "ratio"]

# ---- the script checks its own result -------------------------------------------------------
checks = []


def check(name, ok, detail):
    checks.append(ok)
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}: {detail}")


total = con.execute("SELECT sum(share_pct) FROM by_department").fetchone()[0]
rows_now = con.execute("SELECT count(*) FROM mappings").fetchone()[0]
dup = con.execute("SELECT count(*) - count(DISTINCT task_name) FROM mappings").fetchone()[0]
unmapped_codes = con.execute("SELECT count(*) FROM occ_dept WHERE department = 'UNMAPPED'").fetchone()[0]
n_codes = con.execute("SELECT count(*) FROM occ_dept").fetchone()[0]
unmatched = con.execute("SELECT count(DISTINCT task_name) FROM weighted WHERE code IS NULL").fetchone()[0]
# Each task's share must be handed out in full across its occupations, no more and no less.
spread_err = con.execute("""SELECT coalesce(max(abs(s - pct)), 0) FROM (
    SELECT task_name, pct, sum(weight) AS s FROM weighted GROUP BY task_name, pct)""").fetchone()[0]
# Every prefix in DEPARTMENT_BY_PREFIX must match at least one real occupation code, so a typo
# cannot silently send occupations to a shorter prefix's department.
dead = [r[0] for r in con.execute("""SELECT prefix FROM dept_map d WHERE NOT EXISTS (
    SELECT 1 FROM occ_dept o WHERE o.code LIKE d.prefix || '%') ORDER BY prefix""").fetchall()]
none_pct = con.execute("SELECT coalesce(sum(pct), 0) FROM (SELECT DISTINCT task_name, pct FROM weighted WHERE code IS NULL)").fetchone()[0]

print("Checks")
check("shares sum to 100%", abs(total - 100.0) < 1e-6, f"departments total {total:.6f}%")
check("mapping rows match the source file", rows_now == source_mapping_rows,
      f"{rows_now} rows loaded, {source_mapping_rows} in the source file")
check("no duplicate task in the mapping file", dup == 0, f"{dup} duplicated task name(s)")
check("every task's share is spread in full", spread_err < 1e-9,
      f"largest gap between a task's share and what its occupations received: {spread_err:.2e}")
check("every department prefix matches a real code", not dead,
      "all prefixes used" if not dead else "no code starts with: " + ", ".join(dead))
check("every occupation code maps to a department", unmapped_codes == 0,
      f"{n_codes - unmapped_codes} of {n_codes} codes mapped, {unmapped_codes} unmapped")
check("tasks with no occupation are only the recorded 'none'", unmatched == 1,
      f"{unmatched} task(s) without an occupation, carrying {none_pct:.4f}% of conversations")
ok = all(checks)

print()
print("Share of Claude conversations by department (Economic Index release_2025_02_10)")
print(f"{'department':<42}{'share %':>9}{'tasks':>7}{'occ.':>6}")
for r in result:
    print(f"{r[0]:<42}{r[1]:>9.2f}{r[2]:>7}{r[3]:>6}")
print()
print("Top task per department")
for r in result:
    print(f"  {r[0]}: {r[5]:.3f}% - {r[4]}")
print()
print("SOC major group: Claude share vs US employment share (BLS, May 2023)")
print(f"{'group':<7}{'claude %':>9}{'jobs %':>8}{'ratio':>7}  title")
for r in major:
    print(f"{r[0]:<7}{r[2]:>9.2f}{(r[3] if r[3] is not None else float('nan')):>8.2f}"
          f"{(r[4] if r[4] is not None else float('nan')):>7.2f}  {r[1]}")
print()

if ok:
    import csv
    for name, header, data in (("departments_claude_usage.csv", cols, result),
                               ("major_groups_claude_usage.csv", major_cols, major)):
        with open(OUT / name, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(header)
            w.writerows(data)
        print(f"wrote {name} ({len(data)} rows)")
    print("ALL CHECKS PASSED")
else:
    print(f"{checks.count(False)} CHECK(S) FAILED: no CSV written")
    sys.exit(1)