#!/usr/bin/env python3
"""
extract.py

Rebuilds the two Careery CSVs from the raw ONS workbook, so the extraction
step is checkable rather than taken on trust.

Source
    ONS, "Labour demand volumes by Standard Occupation Classification
    (SOC 2020), UK", January 2017 to September 2026 edition, released
    23 October 2026.

    Table 3 holds new online job adverts split by region and by four-digit
    SOC 2020 code. Rows where Region is "Total UK" are the UK totals.

What this does and does not do
    It selects thirteen occupations and reshapes wide month columns into
    long format. It does not adjust, smooth, interpolate or recompute any
    figure. ONS suppression markers are carried through verbatim.

Usage
    python3 extract.py --workbook refv1.xlsx --outdir .
    python3 extract.py --workbook refv1.xlsx --verify-against ./existing
"""

import argparse
import csv
import os
import sys

try:
    import openpyxl
except ImportError:
    sys.exit("openpyxl is required: pip install openpyxl")

SOC_CODES = [
    2121, 2122, 2133, 2134, 2142, 2421, 2422, 2431,
    3543, 3544, 3554, 3556, 3571,
]

SHEET = "Table 3"
HEADER_ROW = 5
UK_LABEL = "Total UK"
UK_OUTPUT = "United Kingdom"
TREND_MONTHS = 12


def load_table(workbook_path):
    wb = openpyxl.load_workbook(workbook_path, read_only=True, data_only=True)
    if SHEET not in wb.sheetnames:
        sys.exit("sheet not found: " + SHEET)
    ws = wb[SHEET]
    rows = ws.iter_rows(min_row=HEADER_ROW, values_only=True)
    header = next(rows)
    if str(header[0]).strip() != "Region":
        sys.exit("unexpected header, ONS may have changed the layout, stop and check")
    months = [(i, str(h)) for i, h in enumerate(header) if i >= 3 and h]
    for r in rows:
        if r[0] is None or r[1] is None:
            continue
        try:
            code = int(r[1])
        except (TypeError, ValueError):
            continue
        yield str(r[0]).strip(), code, str(r[2]).strip(), {m: r[i] for i, m in months}


def cell(value):
    if value is None:
        return ""
    return str(value).strip()


def build(workbook_path, outdir):
    wanted = set(SOC_CODES)
    by_region = {}
    labels = {}
    all_months = []
    for region, code, label, vals in load_table(workbook_path):
        if code not in wanted:
            continue
        labels[code] = label
        by_region.setdefault(region, {})[code] = vals
        if not all_months:
            all_months = list(vals.keys())
    missing = wanted - set(by_region.get(UK_LABEL, {}))
    if missing:
        sys.exit("SOC codes not found in UK rows: " + str(sorted(missing)))
    latest = all_months[-1]
    trend_months = all_months[-TREND_MONTHS:]

    def as_int(v):
        try:
            return int(str(v).replace(",", ""))
        except (TypeError, ValueError):
            return -1

    ordered_codes = sorted(
        SOC_CODES,
        key=lambda c: as_int(by_region[UK_LABEL][c].get(latest)),
        reverse=True,
    )
    os.makedirs(outdir, exist_ok=True)
    p1 = os.path.join(outdir, "adverts-by-occupation-and-region.csv")
    with open(p1, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["soc_2020", "occupation", "area", "month", "new_online_job_adverts"])
        top = ordered_codes[0]
        regions = sorted(
            [r for r in by_region if r != UK_LABEL],
            key=lambda r: as_int(by_region[r].get(top, {}).get(latest)),
            reverse=True,
        )
        for code in ordered_codes:
            w.writerow([code, labels[code], UK_OUTPUT, latest,
                        cell(by_region[UK_LABEL][code].get(latest))])
            for region in regions:
                if code not in by_region[region]:
                    continue
                w.writerow([code, labels[code], region, latest,
                            cell(by_region[region][code].get(latest))])
    p2 = os.path.join(outdir, "adverts-uk-monthly-trend.csv")
    with open(p2, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["soc_2020", "occupation", "month", "new_online_job_adverts"])
        for code in ordered_codes:
            vals = by_region[UK_LABEL][code]
            for m in trend_months:
                w.writerow([code, labels[code], m, cell(vals.get(m))])
    print("latest month: " + str(latest))
    print("trend window: " + str(trend_months[0]) + " to " + str(trend_months[-1]))
    return p1, p2


def verify(generated, existing_dir):
    ok = True
    for path in generated:
        name = os.path.basename(path)
        other = os.path.join(existing_dir, name)
        if not os.path.exists(other):
            print("MISSING: " + name)
            ok = False
            continue
        a = list(csv.reader(open(path, encoding="utf-8")))
        b = list(csv.reader(open(other, encoding="utf-8")))
        if a == b:
            print("IDENTICAL: " + name)
        else:
            ok = False
            print("DIFFERS: " + name)
    return ok


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workbook", required=True)
    ap.add_argument("--outdir", default=".")
    ap.add_argument("--verify-against")
    a = ap.parse_args()
    generated = build(a.workbook, a.outdir)
    if a.verify_against:
        sys.exit(0 if verify(generated, a.verify_against) else 1)


if __name__ == "__main__":
    main()