#!/usr/bin/env python3
"""Appendix Table A4 — changes in relative representation by occupation group.

Why this file exists:
    The paper cites this table for its headline occupational claim ("the
    prize share of technical professionals and managers falls from 84% to
    73%"). This module builds it.

    Read the ACCURACY section before shipping these numbers: the prize-share
    half is sound, the census-share half is approximate.

What the table is (per the paper's own note):
    For each 1-digit occ1950 group: the share of laureate fathers in that group
    in 1850 and 1950, the share of census workers in the same group and years,
    the ratio of the two (relative representation), and the change in prize
    share across the century.

    Prize shares are *predicted*, not observed: for each group, an indicator for
    "father is in this group" is regressed on the laureate's birth year (OLS),
    and the fitted line is evaluated at 1850 and 1950. This smooths sampling
    noise, and is why some published shares are negative (e.g. Operatives, 1850,
    -0.5%) — a linear fit is not bounded to [0, 1]. Reproducing that sign is the
    strongest confirmation the method is right.

ACCURACY vs. the published table (the published paper, p. 35):
    Prize shares       — method confirmed; 7 of 10 groups land within 0.6pp on
                         both years (mean error 0.5pp, max 2.4pp). The residual
                         concentrates in Technical professional, Managers and
                         Craftsmen, and is driven by the aristocrat handling
                         described below.
    Census shares 1850 — within ~0.2pp of published.
    Census shares 1950 — off by ~0.8pp per group. NOT publication-accurate.

    The cause is weighting. `sample_count` in occ_rank_timeseries.csv is a raw
    unweighted IPUMS record count. The paper's shares are person-weighted
    (PERWT). The 1850 IPUMS sample is very nearly self-weighting, so unweighted
    counts reproduce it; later samples are not, so 1950 drifts. No committed file
    carries the weights — `occ_pop_share` is country-specific and covers only
    occupations laureates' fathers actually held (15 of 269 in 1850), so it
    cannot substitute. Fixing this needs PERWT-weighted shares exported from the
    underlying IPUMS extract, which is not included in this package.

    Because relative representation divides by the census share, its 1950 column
    inherits that error and overstates ratios by roughly 10-15%. Every affected
    field is flagged `"approximate": true` in the output, and the meta block
    records the caveat. Treat the 1950 census and relative-rep columns as
    provisional until weighted shares land.

Aristocrats — the one judgement call in this file:
    9 laureate fathers carry the sentinel occ1950 == 9999 ("aristocrat"), which
    is not a census occupation and has no 1-digit group. They are excluded here,
    matching occ_shares.py, whose 10- and 25-group tables already divide by 704
    rather than 713.

    The published A4 instead keeps them in the denominator AND assigns them into
    groups, by a rule the committed data does not record. Evidence: holding them
    in the sample as non-members reproduces 8 of 10 groups (Managers lands on
    40.8 exactly) but leaves the 1850 column summing to 94%, with the missing 6pp
    falling entirely in Technical professional and Craftsmen — the two groups an
    aristocrat could plausibly be coded into. Since those 9 are heavily
    early-born (median 1874 vs 1921 overall), where they go visibly tilts the
    1850 column.

    Excluding them is the defensible choice: shares sum to 100%, and the table
    stays consistent with neighbouring Tables A2/A3. The cost is that
    Managers/Technical professional/Craftsmen sit 1.5-2.4pp off the published
    column. If the aristocrat coding rule is recovered, revisit this.

Headline stat — flagging a discrepancy in the paper itself:
    the paper's text says the prize share of technical professionals and
    managers "falls from 84% to 73%", and the abstract says 84% to 74%. Table A4
    itself gives 84.6% -> 75.7%; this module gives 85.6% -> 75.4%. So the prose
    disagrees with its own table by ~2.7pp on the endpoint, independent of
    anything done here. The 1850 endpoint (~84-86%) is robust across all three.

Upstream dependencies:
    data/nobelists_analysis.csv        (laureate birth year, occ1950)
    data/occ_shares_occ1950_table.csv  (occ1950 -> 1-digit group crosswalk)
    data/occ_rank_timeseries.csv       (census counts by group and year)

Output:
    code/occ_representation_change/output/occ_representation_change.json

Run:
    python code/occ_representation_change/occ_representation_change.py
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd
import statsmodels.formula.api as smf

PACKAGE_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = PACKAGE_ROOT / "data"
OUT_DIR = Path(__file__).resolve().parent / "output"

EARLY_YEAR = 1850
LATE_YEAR = 1950
OCC_ARISTOCRAT = 9999  # sentinel: not a census occupation, has no 1-digit group

CENSUS_NOTE = (
    "Census shares are computed from unweighted IPUMS record counts "
    "(occ_rank_timeseries.sample_count). The published table uses person-weighted "
    "(PERWT) shares. 1850 is very nearly self-weighting and matches to ~0.2pp; "
    f"{LATE_YEAR} is off by ~0.8pp per group and is provisional."
)


def load_laureates() -> pd.DataFrame:
    """Laureates with a father occupation mapped to a 1-digit occ1950 group."""
    df = pd.read_csv(DATA_DIR / "nobelists_analysis.csv", low_memory=False)
    crosswalk = (
        pd.read_csv(DATA_DIR / "occ_shares_occ1950_table.csv")[
            ["occ1950", "occ1950_1", "occ1950_1_str"]
        ]
        .drop_duplicates()
    )

    df = df[df["occ1950"].notna()]
    n_aristocrat = int((df["occ1950"] == OCC_ARISTOCRAT).sum())
    df = df[df["occ1950"] != OCC_ARISTOCRAT]

    merged = df.merge(crosswalk, on="occ1950", how="left")
    unmapped = int(merged["occ1950_1_str"].isna().sum())
    if unmapped:
        codes = sorted(merged.loc[merged["occ1950_1_str"].isna(), "occ1950"].unique())
        raise ValueError(f"{unmapped} laureate(s) have occ1950 outside the crosswalk: {codes}")

    print(f"  laureates with occ1950: {len(merged)} (excluded {n_aristocrat} aristocrats)")
    return merged


def prize_shares(df: pd.DataFrame) -> dict[str, dict]:
    """Per group: OLS of a group indicator on birth year, evaluated at 1850/1950.

    Returns share percentages, which may be negative — the published table
    reports them that way, and clipping would misrepresent the fitted line.
    """
    out: dict[str, dict] = {}
    for group in sorted(df["occ1950_1_str"].dropna().unique()):
        frame = pd.DataFrame(
            {
                "birth_year": df["birth_year"],
                "ind": (df["occ1950_1_str"] == group).astype(float),
            }
        ).dropna()
        model = smf.ols("ind ~ birth_year", data=frame).fit()
        predicted = model.predict(pd.DataFrame({"birth_year": [EARLY_YEAR, LATE_YEAR]}))
        early, late = 100 * float(predicted.iloc[0]), 100 * float(predicted.iloc[1])
        out[group] = {
            "prize_share_early": round(early, 2),
            "prize_share_late": round(late, 2),
            "prize_share_change_pp": round(late - early, 2),
            "n": int(frame["ind"].sum()),
        }
    return out


def census_shares(groups: list[str]) -> dict[str, dict]:
    """Per group: unweighted census share in 1850 and 1950. See CENSUS_NOTE."""
    ts = pd.read_csv(DATA_DIR / "occ_rank_timeseries.csv")
    ts = ts[ts["level"] == "occgroup_10"]

    out: dict[str, dict] = {g: {} for g in groups}
    for year, key in ((EARLY_YEAR, "census_share_early"), (LATE_YEAR, "census_share_late")):
        counts = ts[ts["year"] == year][["occ_label", "sample_count"]].dropna()
        total = float(counts["sample_count"].sum())
        if total <= 0:
            raise ValueError(f"No census counts for {year}")
        for _, row in counts.iterrows():
            if row.occ_label in out:
                out[row.occ_label][key] = round(100 * float(row.sample_count) / total, 2)

    missing = [g for g in groups if len(out[g]) < 2]
    if missing:
        raise ValueError(f"Groups missing a census share: {missing}")
    return out


def main() -> None:
    print("[1/3] Loading laureates")
    df = load_laureates()

    print("[2/3] Prize shares (OLS of group indicator on birth year)")
    prize = prize_shares(df)

    print("[3/3] Census shares (unweighted — see CENSUS_NOTE)")
    census = census_shares(list(prize))

    rows = []
    for group, p in prize.items():
        c = census[group]
        row = {"group": group, **p, **c}
        # Relative representation: prize share / census share. Guard against a
        # zero census share rather than emitting an infinity into the JSON.
        for era in ("early", "late"):
            cs = c[f"census_share_{era}"]
            row[f"relative_rep_{era}"] = round(p[f"prize_share_{era}"] / cs, 2) if cs else None
        # The 1950 census share is unweighted, so anything derived from it is too.
        row["approximate"] = True
        rows.append(row)

    rows.sort(key=lambda r: r["prize_share_late"], reverse=True)
    for r in rows:
        print(
            f"  {r['group'][:34]:36} prize {r['prize_share_early']:6.1f} -> {r['prize_share_late']:5.1f} "
            f"({r['prize_share_change_pp']:+5.1f}pp)  census {r['census_share_early']:5.1f} -> "
            f"{r['census_share_late']:4.1f}"
        )

    payload = {
        "meta": {
            "script": "occ_representation_change",
            "generated": datetime.now(timezone.utc).isoformat(),
            "early_year": EARLY_YEAR,
            "late_year": LATE_YEAR,
            "n_laureates": int(len(df)),
            "prize_share_method": (
                "Per group, OLS of a 'father in group' indicator on laureate birth year, "
                f"evaluated at {EARLY_YEAR} and {LATE_YEAR}. Shares may be negative."
            ),
            "census_share_method": CENSUS_NOTE,
            "census_share_approximate": True,
            "aristocrats_excluded": True,
            "paper_table": "Appendix Table A4",
            "source_dofile": "a/table_occ_representation_change.do",
        },
        "data": {"rows": rows},
    }

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    out = OUT_DIR / "occ_representation_change.json"
    with out.open("w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
        f.write("\n")
    print(f"\nWrote {out.relative_to(PACKAGE_ROOT)}")


if __name__ == "__main__":
    main()
