#!/usr/bin/env python3
"""Carroway Research, note 7: how many clinicians practice in each specialty, from Medicare's clinician file.

Reproduces every CSV in this folder from the public CMS Doctors and Clinicians National Downloadable File
(Provider Data Catalog dataset mj5m-pzi6, release of 18 August 2026):
  https://data.cms.gov/provider-data/dataset/mj5m-pzi6

  python3 count_specialties.py DAC_NationalDownloadableFile.csv [--out DIR]

Needs Python 3.10+ and pandas. Reads 12 of the file's 31 columns and never writes a name, address, phone or NPI.

Rules:
- The file has one line per clinician, enrollment, group and address. A clinician is one NPI.
- A specialty group is the set of CMS primary specialty strings in GROUPS below. A clinician counts in a group
  when any of their lines carries one of those primary specialties, so a clinician whose enrollments list
  specialties from different groups counts in each of them (coverage.csv counts them).
- Person-level fields (gender, credential, graduation year, telehealth, Medicare assignment, group size, state) come
  from one line per clinician within the group: the line with the largest group (num_org_mem, blank = 0), then the
  lowest address ID, then the newest enrollment ID. coverage.csv counts the clinicians whose lines differ on each of
  these fields (a handful in this release).
- Addresses count across all of the clinician's lines in the group. CMS gives offices in the same building the
  same address ID apart from the last 2 characters, so "more than one practice location" uses the ID without them.
"""
import argparse
import os
import re

import pandas as pd

# slug: (label for tables, CMS primary specialty strings)
GROUPS = {
    "cardiologist": ("Cardiologists", ["CARDIOVASCULAR DISEASE (CARDIOLOGY)", "INTERVENTIONAL CARDIOLOGY", "CARDIAC ELECTROPHYSIOLOGY", "ADVANCED HEART FAILURE AND TRANSPLANT CARDIOLOGY"]),
    "dermatologist": ("Dermatologists", ["DERMATOLOGY"]),
    "oncologist": ("Oncologists", ["HEMATOLOGY/ONCOLOGY", "MEDICAL ONCOLOGY", "SURGICAL ONCOLOGY", "GYNECOLOGICAL ONCOLOGY"]),
    "neurologist": ("Neurologists", ["NEUROLOGY"]),
    "urologist": ("Urologists", ["UROLOGY"]),
    "general-surgeon": ("General surgeons", ["GENERAL SURGERY"]),
    "orthopedic-surgeon": ("Orthopedic surgeons", ["ORTHOPEDIC SURGERY", "HAND SURGERY"]),
    "plastic-surgeon": ("Plastic surgeons", ["PLASTIC AND RECONSTRUCTIVE SURGERY"]),
    "pediatrician": ("Pediatricians", ["PEDIATRIC MEDICINE"]),
    "anesthesiologist": ("Anesthesiologists", ["ANESTHESIOLOGY"]),
    "gastroenterologist": ("Gastroenterologists", ["GASTROENTEROLOGY"]),
    "neurosurgeon": ("Neurosurgeons", ["NEUROSURGERY"]),
    "ophthalmologist": ("Ophthalmologists", ["OPHTHALMOLOGY"]),
    "ent": ("Otolaryngologists (ENT)", ["OTOLARYNGOLOGY"]),
    "psychiatrist": ("Psychiatrists", ["PSYCHIATRY", "GERIATRIC PSYCHIATRY"]),
    "radiologist": ("Radiologists", ["DIAGNOSTIC RADIOLOGY", "INTERVENTIONAL RADIOLOGY"]),
    "emergency-physician": ("Emergency physicians", ["EMERGENCY MEDICINE"]),
    "obgyn": ("OB/GYNs", ["OBSTETRICS/GYNECOLOGY"]),
    "pulmonologist": ("Pulmonologists", ["PULMONARY DISEASE"]),
    "endocrinologist": ("Endocrinologists", ["ENDOCRINOLOGY"]),
    "rheumatologist": ("Rheumatologists", ["RHEUMATOLOGY"]),
    "nephrologist": ("Nephrologists", ["NEPHROLOGY"]),
    "vascular-surgeon": ("Vascular surgeons", ["VASCULAR SURGERY"]),
    "cardiothoracic-surgeon": ("Cardiothoracic surgeons", ["THORACIC SURGERY", "CARDIAC SURGERY"]),
    "infectious-disease-physician": ("Infectious disease physicians", ["INFECTIOUS DISEASE"]),
    "allergist": ("Allergists", ["ALLERGY/IMMUNOLOGY"]),
    "pain-management-physician": ("Pain management physicians", ["PAIN MANAGEMENT", "INTERVENTIONAL PAIN MANAGEMENT"]),
    "physiatrist": ("Physiatrists (PM&R)", ["PHYSICAL MEDICINE AND REHABILITATION"]),
    "podiatrist": ("Podiatrists", ["PODIATRY"]),
    "radiation-oncologist": ("Radiation oncologists", ["RADIATION ONCOLOGY"]),
}
# primary care and other large groups, counted for context in other-groups.csv
OTHER = {
    "internal-medicine": ("Internal medicine physicians", ["INTERNAL MEDICINE"]),
    "family-physician": ("Family physicians", ["FAMILY PRACTICE"]),
    "hospitalist": ("Hospitalists", ["HOSPITALIST"]),
    "nurse-practitioner": ("Nurse practitioners", ["NURSE PRACTITIONER"]),
    "physician-assistant": ("Physician assistants", ["PHYSICIAN ASSISTANT"]),
    "crna": ("Nurse anesthetists (CRNA)", ["CERTIFIED REGISTERED NURSE ANESTHETIST (CRNA)"]),
    "chiropractor": ("Chiropractors", ["CHIROPRACTIC"]),
    "optometrist": ("Optometrists", ["OPTOMETRY"]),
    "colorectal-surgeon": ("Colorectal surgeons", ["COLORECTAL SURGERY (PROCTOLOGY)"]),
    "oral-surgeon": ("Oral and maxillofacial surgeons", ["ORAL SURGERY", "MAXILLOFACIAL SURGERY"]),
}
COLS = ["NPI", "Ind_enrl_ID", "gndr", "Cred", "Grd_yr", "pri_spec", "Telehlth", "num_org_mem", "State", "ind_assgn", "adrs_id", "org_pac_id"]
STATES = set("AL AK AZ AR CA CO CT DE DC FL GA HI ID IL IN IA KS KY LA ME MD MA MI MN MS MO MT NE NV NH NJ NM NY NC ND OH OK OR PA RI SC SD TN TX UT VT VA WA WV WI WY".split())


def pct(n, d):
    return round(100.0 * n / d, 1) if d else ""


def load(path):
    df = pd.read_csv(path, usecols=COLS, dtype=str, keep_default_na=False, na_filter=False, encoding="utf-8", encoding_errors="replace")
    df.columns = [re.sub(r"[^a-z0-9]+", "_", c.strip().lower()).strip("_") for c in df.columns]
    df = df[df["npi"].str.strip() != ""].copy()
    df["npi"] = df["npi"].str.strip()
    df["_num"] = pd.to_numeric(df["num_org_mem"], errors="coerce").fillna(0).astype("int64")
    df["_bldg"] = df["adrs_id"].str.strip().str[:-2]
    return df


def one_per_npi(rows):
    rows = rows.sort_values(["npi", "_num", "adrs_id", "ind_enrl_id"], ascending=[True, False, True, False], kind="stable")
    return rows.drop_duplicates("npi", keep="first")


def profile(label, strings, df):
    rows = df[df["pri_spec"].isin(strings)]
    p = one_per_npi(rows)
    n = len(p)
    cred = p["cred"].str.strip().str.upper()
    md, do = int((cred == "MD").sum()), int((cred == "DO").sum())
    yrs = pd.to_numeric(p["grd_yr"], errors="coerce")
    yrs = yrs[(yrs >= 1900) & (yrs <= 2100)]
    bldg = rows.loc[rows["_bldg"] != "", ["npi", "_bldg"]].drop_duplicates().groupby("npi").size()
    addr = rows.loc[rows["adrs_id"].str.strip() != "", ["npi", "adrs_id"]].drop_duplicates().groupby("npi").size()
    grp = p["_num"]
    return {
        "specialty": label,
        "cms_primary_specialties": "; ".join(strings),
        "clinicians": n,
        "lines_in_file": len(rows),
        "women": int((p["gndr"] == "F").sum()),
        "women_pct": pct(int((p["gndr"] == "F").sum()), n),
        "md": md,
        "do": do,
        "md_or_do": md + do,
        # DO share only where most of the group holds an MD or DO (podiatrists hold a DPM)
        "do_pct_of_md_and_do": pct(do, md + do) if md + do >= n / 2 else "",
        "other_or_blank_credential": n - md - do,
        "grad_year_known": int(len(yrs)),
        "grad_year_median": int(yrs.median()) if len(yrs) else "",
        "grad_year_p25": int(yrs.quantile(0.25)) if len(yrs) else "",
        "grad_year_p75": int(yrs.quantile(0.75)) if len(yrs) else "",
        "grad_2015_or_later_pct": pct(int((yrs >= 2015).sum()), len(yrs)),
        "grad_before_1990_pct": pct(int((yrs < 1990).sum()), len(yrs)),
        "telehealth_pct": pct(int((p["telehlth"].str.strip().str.upper() == "Y").sum()), n),
        "multi_location_pct": pct(int((p["npi"].map(bldg).fillna(1) > 1).sum()), n),
        "multi_address_id_pct": pct(int((p["npi"].map(addr).fillna(1) > 1).sum()), n),
        "group_100_plus_pct": pct(int((grp >= 100).sum()), n),
        "no_group_pct": pct(n - int(rows.loc[rows["org_pac_id"].str.strip() != "", "npi"].nunique()), n),
        "accepts_assignment_pct": pct(int((p["ind_assgn"].str.strip().str.upper() == "Y").sum()), n),
    }, p


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cms_file")
    ap.add_argument("--out", default=os.path.dirname(os.path.abspath(__file__)))
    a = ap.parse_args()
    df = load(a.cms_file)
    out = lambda name: os.path.join(a.out, name)

    # 1. the 30 specialty groups
    prof, members, states = [], {}, []
    for slug, (label, strings) in GROUPS.items():
        row, p = profile(label, strings, df)
        prof.append(row)
        members[slug] = set(p["npi"])
        st = p["state"].str.strip().str.upper()
        vc = st[st != ""].value_counts()
        vc = vc.sort_index().sort_values(ascending=False, kind="stable")  # ties in alphabetical order
        for rank, (s, c) in enumerate(vc.items(), 1):
            states.append({"specialty": label, "state": s, "clinicians": int(c), "share_pct": pct(int(c), len(p)), "rank_in_field": rank})
    P = pd.DataFrame(prof).sort_values("clinicians", ascending=False)
    P[["specialty", "clinicians", "lines_in_file", "cms_primary_specialties"]].to_csv(out("clinicians-by-specialty.csv"), index=False)
    P[["specialty", "clinicians", "women", "women_pct", "md", "do", "md_or_do", "do_pct_of_md_and_do", "other_or_blank_credential",
       "grad_year_known", "grad_year_median", "grad_year_p25", "grad_year_p75", "grad_2015_or_later_pct", "grad_before_1990_pct",
       "telehealth_pct", "multi_location_pct", "multi_address_id_pct", "group_100_plus_pct", "no_group_pct", "accepts_assignment_pct"]
      ].to_csv(out("specialty-profile.csv"), index=False)
    pd.DataFrame(states).to_csv(out("clinicians-by-state.csv"), index=False)

    # 2. context groups
    other = []
    for slug, (label, strings) in OTHER.items():
        row, _ = profile(label, strings, df)
        other.append({"specialty": row["specialty"], "clinicians": row["clinicians"], "women_pct": row["women_pct"], "telehealth_pct": row["telehealth_pct"], "counted_as": row["cms_primary_specialties"]})
    pd.DataFrame(other).sort_values("clinicians", ascending=False).to_csv(out("other-groups.csv"), index=False)

    # 3. every primary specialty string in the file
    allspec = (df.assign(pri_spec=df["pri_spec"].str.strip()).groupby("pri_spec")
                 .agg(clinicians=("npi", "nunique"), lines_in_file=("npi", "size")).reset_index()
                 .rename(columns={"pri_spec": "cms_primary_specialty"}))
    allspec = allspec[allspec["cms_primary_specialty"].str.match(r"^[A-Z]")].sort_values(["clinicians", "cms_primary_specialty"], ascending=[False, True])
    allspec.to_csv(out("all-primary-specialties.csv"), index=False)

    # 4. context from the AAMC's 2025 key findings (active physicians in 2024), next to the same measures in this file
    aamc_url = "https://www.aamc.org/data-reports/data/2025-key-findings"
    byname = {r["specialty"]: r for r in prof}
    md_do = int(df.loc[df["cred"].str.strip().str.upper().isin(["MD", "DO"]), "npi"].nunique())
    ortho_lines = df[df["pri_spec"].isin(["ORTHOPEDIC SURGERY", "HAND SURGERY"])]
    hand_npis = set(ortho_lines.loc[ortho_lines["pri_spec"] == "HAND SURGERY", "npi"])
    ortho_only = one_per_npi(ortho_lines[(ortho_lines["pri_spec"] == "ORTHOPEDIC SURGERY") & (~ortho_lines["npi"].isin(hand_npis))])
    hand_any = one_per_npi(ortho_lines[ortho_lines["npi"].isin(hand_npis)])
    bench = [
        ("physicians or clinicians counted", md_do, 1032365, "Medicare file: clinicians with an MD or DO credential; AAMC: active physicians"),
        ("direct patient care physicians", "", 866460, "AAMC only; the Medicare file has no equivalent measure"),
        ("orthopedic surgeons with no hand surgery line (clinicians)", len(ortho_only), "", "Medicare file only"),
        ("women among orthopedic surgeons with no hand surgery line (%)", pct(int((ortho_only["gndr"] == "F").sum()), len(ortho_only)), 6.8, "closest match to the AAMC orthopedic surgery group"),
        ("clinicians with a hand surgery line in the orthopedic group", len(hand_any), "", "Medicare file only"),
        ("women among clinicians with a hand surgery line (%)", pct(int((hand_any["gndr"] == "F").sum()), len(hand_any)), "", "Medicare file only"),
        ("women among pediatricians (%)", byname["Pediatricians"]["women_pct"], 66.7, ""),
        ("women among OB/GYNs (%)", byname["OB/GYNs"]["women_pct"], 64.1, ""),
        ("women among orthopedic surgeons (%)", byname["Orthopedic surgeons"]["women_pct"], 6.8, "the Medicare file group includes hand surgery"),
    ]
    pd.DataFrame([{"measure": m, "medicare_file_aug_2026": a_, "aamc_2024": b_, "note": n_, "aamc_source": aamc_url if b_ != "" else ""} for m, a_, b_, n_ in bench], dtype=object).to_csv(out("benchmark-aamc.csv"), index=False)

    # 5. subspecialty codes counted outside the field a reader would expect them in
    related = [
        ("MICROGRAPHIC DERMATOLOGIC SURGERY (MDS)", "dermatologist"),
        ("EPILEPTOLOGISTS", "neurologist"),
        ("ADULT CONGENITAL HEART DISEASE (ACHD)", "cardiologist"),
        ("HEMATOPOIETIC CELL TRANSPLANTATION AND CELLULAR THERAPY", "oncologist"),
    ]
    sub = []
    for code, slug in related:
        npis = set(df.loc[df["pri_spec"] == code, "npi"])
        sub.append({"cms_primary_specialty": code, "related_specialty": GROUPS[slug][0], "clinicians": len(npis),
                    "not_counted_in_related_specialty": len(npis - members[slug])})
    pd.DataFrame(sub).to_csv(out("subspecialties-outside.csv"), index=False)

    # 6. coverage
    union = set().union(*members.values())
    counts = {}
    for s in members.values():
        for n in s:
            counts[n] = counts.get(n, 0) + 1
    in_two = sum(1 for v in counts.values() if v > 1)
    nspec = df.groupby("npi")["pri_spec"].nunique()
    st = df["state"].str.strip().str.upper()
    cov = [
        ("lines in the file", len(df), "one line per clinician, enrollment, group and address"),
        ("clinicians in the file", df["npi"].nunique(), "distinct NPIs"),
        ("clinicians with two or more primary specialties", int((nspec > 1).sum()), "different enrollments list different primary specialties"),
        ("clinicians with an MD or DO credential", int(df.loc[df["cred"].str.strip().str.upper().isin(["MD", "DO"]), "npi"].nunique()), "on at least one line, across every specialty"),
        ("clinicians with no credential on any line", int(df["npi"].nunique() - df.loc[df["cred"].str.strip() != "", "npi"].nunique()), "the credential field is blank"),
        ("primary specialty strings", int(allspec.shape[0]), "distinct values that start with a letter; %d lines carry a blank or numeric value" % int((~df["pri_spec"].str.strip().str.match(r"^[A-Z]")).sum())),
        ("lines in the 30 specialty groups", int(df["pri_spec"].isin([s for _, ss in GROUPS.values() for s in ss]).sum()), ""),
        ("clinicians summed over the 30 groups", int(P["clinicians"].sum()), "a clinician in two or three groups counts once in each"),
        ("distinct clinicians in the 30 groups", len(union), ""),
        ("clinicians in two or more of the 30 groups", in_two, ""),
        ("clinicians in exactly two of the 30 groups", sum(1 for v in counts.values() if v == 2), ""),
        ("clinicians in three of the 30 groups", sum(1 for v in counts.values() if v == 3), ""),
        ("clinicians in four or more of the 30 groups", sum(1 for v in counts.values() if v >= 4), ""),
        ("lines outside the 50 states and DC", int((~st.isin(STATES)).sum()), "they stay in every count; values: " + ", ".join(f"{k or 'blank'} {v}" for k, v in st[~st.isin(STATES)].value_counts().items())),
    ]
    g = df.groupby("npi")
    for col, label in [("gndr", "gender"), ("grd_yr", "graduation year"), ("telehlth", "telehealth flag"), ("cred", "credential"), ("ind_assgn", "Medicare assignment")]:
        cov.append((f"clinicians whose lines differ on {label}", int((g[col].nunique() > 1).sum()), "across the whole file"))
    pd.DataFrame(cov, columns=["measure", "value", "note"]).to_csv(out("coverage.csv"), index=False)
    print(P[["specialty", "clinicians", "women_pct", "do_pct_of_md_and_do", "grad_year_median", "telehealth_pct", "multi_location_pct", "multi_address_id_pct", "group_100_plus_pct", "accepts_assignment_pct"]].to_string(index=False))
    print(pd.DataFrame(cov).to_string(index=False))


if __name__ == "__main__":
    main()