#!/usr/bin/env python3
"""Does Frey and Osborne's 2013 map predict employment change beyond routine-task intensity?

Extends Coelli and Borland (2019), who tested 2013-2018, to later BLS years. Reproduces the
regression results in the Save The 1 Percent video "It Took Twelve Years to Show"
(savethe1percent.com/videos/002). Needs fo_vs_oews.py in the same folder.

Chain: 2010 SOC code -> 2010 Census occupation code (Census Bureau code list, sheet
"2010summ"; aggregated SOC entries such as 15-113X match by prefix) -> occ1990dd (Dorn
crosswalk [A8]) -> 1977 Dictionary of Occupational Titles task scores (Dorn [B1]).

Routine Task Intensity (RTI) = ln(routine) - ln(manual) - ln(abstract), after Autor and Dorn
(2013). Task scores of zero are floored at the smallest positive value of that task; this is
our choice, not theirs.

Second benchmark, as in Coelli and Borland's Table 6 columns 3-4: four task groups by SOC
major group, after Acemoglu and Autor (2011): non-routine cognitive (11-29), routine
cognitive (41, 43), routine manual (47, 49, 51, 53), non-routine manual (31-39); farming (45)
is left out.

Regression: percentage employment change from May 2013, weighted by May 2013 employment,
heteroscedasticity-robust (HC1) standard errors. Robustness rows vary the measure and weights.

The BLS site refuses downloads without a contact address, so set your own first:
    export CONTACT_EMAIL=you@example.org
    uv run --with pandas --with openpyxl --with xlrd fo_vs_rti.py

Copyright (c) 2026 Stephen Angeloni. MIT License: see LICENSE.
"""
import argparse
import re
import sys
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd
import xlrd

sys.path.insert(0, str(Path(__file__).parent))
from fo_vs_oews import fetch, frey_osborne, oews  # noqa: E402

DORN = "https://www.ddorn.net/data/{name}.zip"
CENSUS = ("https://www2.census.gov/programs-surveys/demo/guidance/industry-occupation/"
          "2010-occ-codes-with-crosswalk-from-2002-2011.xls")


def dorn(work, name):
    z = fetch(DORN.format(name=name), work / f"{name}.zip")
    with zipfile.ZipFile(z) as zf:
        dta = next(n for n in zf.namelist() if n.endswith(".dta"))
        zf.extract(dta, work)
    return pd.read_stata(work / dta)


def census_codes(work):
    sheet = xlrd.open_workbook(fetch(CENSUS, work / "census2010.xls")).sheet_by_name("2010summ")
    exact, prefixes = {}, []
    for r in range(3, sheet.nrows):
        soc = str(sheet.cell_value(r, 1)).strip()
        v = sheet.cell_value(r, 2)
        cc = str(int(v)).zfill(4) if isinstance(v, float) else str(v).strip().zfill(4)
        if not re.match(r"^\d{2}-\d{3}[\dX]$", soc):
            continue
        exact.setdefault(soc, cc)
        prefixes.append((re.sub(r"[X0]+$", "", soc.replace("-", "")), cc))
    prefixes.sort(key=lambda x: -len(x[0]))

    def lookup(code):
        if code in exact:
            return exact[code]
        digits = code.replace("-", "")
        return next((cc for pre, cc in prefixes if len(pre) >= 4 and digits.startswith(pre)), None)
    return lookup


GROUPS = {"routine cognitive": {41, 43}, "routine manual": {47, 49, 51, 53},
          "non-routine manual": {31, 33, 35, 37, 39}}  # base: non-routine cognitive


def dummies(codes):
    major = codes.str[:2].astype(int)
    return [major.isin(g).astype(float).to_numpy() for g in GROUPS.values()]


def wls(y, xs, w):
    X = np.column_stack([np.ones(len(y))] + xs)
    b, *_ = np.linalg.lstsq(X * np.sqrt(w)[:, None], y * np.sqrt(w), rcond=None)
    e = y - X @ b
    bread = np.linalg.inv((X * w[:, None]).T @ X)
    s = X * (w * e)[:, None]
    n, k = X.shape
    se = np.sqrt(np.diag(bread @ (s.T @ s) @ bread * n / (n - k)))
    r2 = 1 - np.sum(w * e**2) / np.sum(w * (y - np.average(y, weights=w)) ** 2)
    return b, se, r2


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workdir", default="/tmp/fo_vs_oews")
    ap.add_argument("--years", default="18,22,25", help="two-digit OEWS end years")
    args = ap.parse_args()
    work = Path(args.workdir)
    work.mkdir(parents=True, exist_ok=True)

    fo = frey_osborne(work)
    lookup = census_codes(work)
    cw = dorn(work, "occ2010_occ1990dd")
    cw["occ"] = cw.occ.astype(str).str.zfill(4)
    tasks = dorn(work, "occ1990dd_task_alm")
    fo["cc"] = fo.code.map(lookup)
    fo = fo.merge(cw, left_on="cc", right_on="occ", how="left").merge(tasks, on="occ1990dd", how="left")
    floor = lambda s: s.clip(lower=s[s > 0].min())  # noqa: E731
    fo["rti"] = (np.log(floor(fo.task_routine)) - np.log(floor(fo.task_manual))
                 - np.log(floor(fo.task_abstract)))
    e13 = oews(work, "13")

    for yy in args.years.split(","):
        eN = oews(work, yy)
        d = fo.assign(e13=fo.code.map(e13), eN=fo.code.map(eN)).dropna(subset=["e13", "eN", "rti"])
        d = d[(d.e13 > 0) & (d.eN > 0)]
        pct = ((d.eN / d.e13 - 1) * 100).to_numpy()
        logc = (np.log(d.eN / d.e13) * 100).to_numpy()
        w = d.e13.to_numpy()
        p, rti = d.p.to_numpy(), d.rti.to_numpy()
        print(f"\nMay 2013 -> May 20{yy}: {len(d)} occupations with identical SOC code and an RTI score")
        for label, xs in [("FO alone", [p]), ("RTI alone", [rti]), ("FO + RTI", [p, rti])]:
            b, se, r2 = wls(pct, xs, w)
            coefs = "  ".join(f"b={b[i]:.2f} (se {se[i]:.2f}, t {b[i] / se[i]:.2f})" for i in range(1, len(b)))
            print(f"  {label:10s} {coefs}  R2={r2:.3f}")
        for label, y, ww in [("log, weighted", logc, w), ("pct, unweighted", pct, np.ones(len(d)))]:
            b, se, r2 = wls(y, [p, rti], ww)
            print(f"  FO + RTI, {label:16s} FO t {b[1] / se[1]:.2f}  RTI t {b[2] / se[2]:.2f}  R2={r2:.3f}")
        g = d[d.code.str[:2] != "45"]
        gp, gw = ((g.eN / g.e13 - 1) * 100).to_numpy(), g.e13.to_numpy()
        b, se, r2 = wls(gp, [g.p.to_numpy()] + dummies(g.code), gw)
        print(f"  FO + four task groups  FO t {b[1] / se[1]:.2f}  R2={r2:.3f}  "
              f"(groups: {', '.join(f'{k} t {b[i + 2] / se[i + 2]:.2f}' for i, k in enumerate(GROUPS))})")
        b, se, r2 = wls(gp, [g.p.to_numpy(), g.rti.to_numpy()] + dummies(g.code), gw)
        print(f"  FO + RTI + groups      FO t {b[1] / se[1]:.2f}  RTI t {b[2] / se[2]:.2f}  R2={r2:.3f}")
        z = lambda a: (a - a.mean()) / a.std()  # noqa: E731
        b, _, _ = wls(logc, [z(p), z(rti)], w)
        print(f"  Per standard deviation, log points, weighted: FO {b[1]:.1f}  RTI {b[2]:.1f}")


if __name__ == "__main__":
    main()
