#!/usr/bin/env python3
"""What engineering application forms ask about pay, and whether the posting printed a range first.

Reads roles.json, companies.json, labels.json, pay-fields.json, laws.json (here)
and OUT_DIR/eng.jsonl, OUT_DIR/hits.jsonl. Prints every figure the article
uses; writes role-pay.json, one row per role, for the appendix.

Counting rules:
- A role is what a company is filling (see roles.py). A role asks a question if
  the form of any of its postings carries it; the question is required for the
  role if it is required on any of those forms. Roles whose forms were not read
  at all are left out of every share.
- Main sample: roles open in the United States at companies based in the
  United States (companies.json), as on 13 October.
- A role prints a range if any of its postings has one in the board's pay field
  (pay-fields.json) or an annual range in its text (../engineering-by-country/pay.py,
  the PAY pattern, any currency, 20 000 to 5 000 000).
- Comparison sample: roles not open in the United States, at any company.
Usage: python3 analyze.py OUT_DIR
"""
import collections
import json
import re
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
OUT = Path(sys.argv[1])
sys.path.insert(0, str(HERE.parent / 'engineering-by-country'))
sys.argv = [sys.argv[0], str(OUT)]
import pay  # noqa: E402
from states import states  # noqa: E402

US = 'United States'
KINDS = ('expectation', 'history', 'range_check', 'offers', 'pay_other')


def pct(a, b):
    return f'{100 * a / b:.1f}%' if b else 'n/a'


# What the posting itself says the job pays, in US dollars. pay.PAY (built for
# comparing two countries' bands) misses most US spellings, found by listing
# every US role with a dollar amount and no range on 10 November 2026:
#   "$130,000.00 - $160,000.00"        SpaceX, every posting (cents)
#   "$191,000/yr to $225,000/yr"       Gusto (unit between the amounts)
#   "163,000–204,000 in Denver"        Gusto (no currency at all)
#   "150,000 - 251,900 USD"            PagerDuty, Affirm (code after, no sign)
#   "$170000 - $195000"                Betterment (no thousands separator)
#   "$176,000 and $315,00"             Nuro, one posting (typo for $315,000)
#   "$30.00/hour", "$6,500 to $10,500" interns, paid by the hour, week or month
# So the check here is wider: a range of two annual amounts, or a single amount
# stated as the pay for a period ("$60/hour", "$2,400 per week", "Annual
# Salary: $185,000 USD"). Either counts as the posting printing its pay.
CENTS = re.compile(r'(\d)\.\d{1,2}\b(?=\s?(?:/|per|an|a|-|–|—|to|USD|\s|$))')
UNIT_IN = re.compile(r'(\d)\s?/\s?(?:yr|year|annum)\b|(\d)\s?(?:per year|annually|a year)\s?(?=-|–|—|to\b)', re.I)
FIXES = {'$176,000 and $315,00 ': '$176,000 and $315,000 '}
AMT = r'(?:USD\s?)?(?:US\$|\$)?\s?(?:\d{2,3}(?:,\d{3})+|\d{5,7}|\d{2,3}(?:\.\d)?\s?[kK])(?:\s?USD)?'
ANNUAL = re.compile(r'(?P<lo>' + AMT + r')\s?(?:-|–|—|to|and|up to)\s?(?P<hi>' + AMT + r')')
PERIOD = re.compile(r'(?:US\$|\$)\s?\d[\d,]*(?:\.\d+)?\s?(?:USD\s?)?(?:/\s?|per\s|an\s|a\s)(?:hour|hr\b|week|month|mo\b)'
                    r'|(?:US\$|\$)\s?\d[\d,]*(?:\.\d+)?\s?(?:-|–|—|to)\s?(?:US\$|\$)?\s?\d[\d,]*(?:\.\d+)?\s?(?:USD\s?)?'
                    r'(?:/\s?|per\s|an\s|a\s)?(?:hour|hr\b|week|month|mo\b|monthly|hourly)'
                    r'|(?:monthly|hourly|weekly) (?:base )?(?:salary|pay|rate)[^.$]{0,60}\$\s?\d', re.I)
SINGLE = re.compile(r'(?:annual|base) (?:salary|pay)\s*(?:is|:)\s*(?:US\$|\$)\s?\d{2,3},\d{3}', re.I)
PAYWORD = re.compile(r'salary|pay\b|compensation|wage|earn|OTE|base', re.I)
FOREIGN = re.compile(r'\bCAD\b|CA\$|C\$|Canad|Ontario|British Columbia|Toronto|Vancouver|Mexic|MXN|€|£|EUR\b|GBP\b|'
                     r'AUD\b|A\$|SGD\b|S\$|INR\b|₹|R\$|BRL\b|PLN\b|zł')


def money(s):
    s = re.sub(r'US\$|\$|USD|\s', '', s)
    if s[-1] in 'kK':
        return float(s[:-1]) * 1000
    return float(s.replace(',', ''))


def clean(text):
    for a, b in FIXES.items():
        text = text.replace(a, b)
    return UNIT_IN.sub(lambda m: m.group(1) or m.group(2), CENTS.sub(r'\1', text))


def text_range(text):
    """Any currency, pay.PAY as published, after cleaning; for the comparison sample."""
    text = clean(text)
    for m in pay.PAY.finditer(text):
        lo, hi = pay.amount(m.group('lo')), pay.amount(m.group('hi'))
        if 20000 <= lo < hi <= 5000000:
            return re.sub(r'\s+', ' ', m.group(0)).strip()
    return usd_text_range(text)


def usd_text_range(text):
    """The pay in US dollars, as described above; the matched words, or None."""
    text = clean(text)
    for m in ANNUAL.finditer(text):
        lo, hi = money(m.group('lo')), money(m.group('hi'))
        if not 20000 <= lo < hi <= 5000000:
            continue
        near = text[max(0, m.start() - 120):m.end() + 12]
        dollars = '$' in m.group(0) or 'USD' in m.group(0)
        if FOREIGN.search(text[max(0, m.start() - 80):m.end() + 8]):
            continue
        if not dollars and not PAYWORD.search(near):
            continue
        if re.search(r'(?i)raised|valuation|funding|invest|revenue|customers|billion|million', near[:120]):
            continue
        return re.sub(r'\s+', ' ', m.group(0)).strip()
    for rx in (PERIOD, SINGLE):
        for m in rx.finditer(text):
            if not FOREIGN.search(text[max(0, m.start() - 80):m.end() + 8]):
                return re.sub(r'\s+', ' ', m.group(0)).strip()
    return None


def usd_field(rng):
    if rng.get('currency'):
        return rng['currency'] == 'USD'
    lab = rng.get('label') or ''
    return '$' in lab and not re.search(r'CA\$|C\$|A\$|S\$|R\$|CAD|AUD|SGD', lab)


def main():
    roles = json.load(open(HERE / 'roles.json'))
    comp = json.load(open(HERE / 'companies.json'))
    labels = {v['text']: v for v in json.load(open(HERE / 'labels.json')).values()}
    fields = json.load(open(HERE / 'pay-fields.json'))
    laws = json.load(open(HERE / 'laws.json'))
    eng = {r['id']: r for r in map(json.loads, open(OUT / 'eng.jsonl'))}
    hits = {h['id']: h for h in map(json.loads, open(OUT / 'hits.jsonl'))}

    rows = []
    for r in roles:
        ps = r['posting_ids']
        read = any(hits[p]['form_read'] for p in ps)
        asked = {}
        for p in ps:
            for q in hits[p]['questions']:
                lab = labels[q['text']]
                if lab['label'] not in KINDS:
                    continue
                k = (lab['label'], q['text'])
                asked[k] = asked.get(k, False) or q['required']
        field = next(({'posting': p, **fields[p]['ranges'][0]} for p in ps if fields[p]['ranges']), None)
        text = next(((p, t) for p in ps for t in [text_range(eng[p]['text'])] if t), None)
        usd = next((f"field: {x.get('label') or ''} {x['low']}-{x['high']} {x.get('period') or ''}".strip()
                    for p in ps for x in fields[p]['ranges'] if usd_field(x)), None) or next(
            (f'text: {t}' for p in ps for t in [usd_text_range(eng[p]['text'])] if t), '')
        st = set()
        for p in ps:
            st |= states(eng[p]['location'])
            for s in eng[p]['more_locations']:
                st |= states(s)
        rows.append({
            'id': r['id'], 'company': r['company'], 'title': r['title'], 'url': r['url'],
            'api': eng[ps[0]]['api'], 'postings': len(ps), 'places': r['places'], 'countries': r['countries'],
            'hq': comp[r['company']]['hq'], 'seniority': r['seniority'], 'mode': r['mode'],
            'us_states': sorted(st), 'form_read': read,
            'prints_range': bool(field or text),
            'prints_usd_range': bool(usd),
            'usd_pay': usd,
            'range_in_field': field['label'] if field and field.get('label') else (
                f"{field['currency']} {field['low']}-{field['high']}" if field else ''),
            'range_in_text': text[1] if text else '',
            'questions': [{'label': k[0], 'text': k[1], 'required': req,
                           **({f: labels[k[1]][f] for f in ('answer', 'scope', 'currency', 'no_history_warning')}
                              if k[0] in ('expectation', 'history') else {})}
                          for k, req in sorted(asked.items())],
        })

    def has(r, kind, required=None):
        return any(q['label'] == kind and (required is None or q['required'] == required) for q in r['questions'])

    read = [r for r in rows if r['form_read']]
    print(f'roles {len(rows)}, forms read for {len(read)} roles; postings {len(eng)}, '
          f'forms read {sum(h["form_read"] for h in hits.values())}')
    print(f'companies {len({r["company"] for r in rows})}; based in the US '
          f'{sum(1 for c in {r["company"] for r in rows} if comp[c]["hq"] == US)}')
    main_ = [r for r in read if r['hq'] == US and US in r['countries']]
    abroad = [r for r in read if r['countries'] and US not in r['countries']]

    def block(name, rs, key):
        cos = sorted({r['company'] for r in rs})
        print(f'\n######## {name}: roles {len(rs)}, companies {len(cos)}')
        pr = [r for r in rs if r['prints_range']]
        pu = [r for r in rs if r['prints_usd_range']]
        print(f'  prints a range in US dollars {len(pu)} {pct(len(pu), len(rs))}')
        print(f'  prints a pay range {len(pr)} {pct(len(pr), len(rs))}; in a board field '
              f'{sum(1 for r in rs if r["range_in_field"])}, only in the text '
              f'{sum(1 for r in rs if r["range_in_text"] and not r["range_in_field"])}')
        for kind in KINDS:
            k = [r for r in rs if has(r, kind)]
            kc = sorted({r['company'] for r in k})
            req = [r for r in k if has(r, kind, True)]
            wr = [r for r in k if r[key]]
            print(f'  {kind:<12} roles {len(k):>4} {pct(len(k), len(rs)):>6}  required {len(req):>4} '
                  f'{pct(len(req), len(k)):>6} of askers  companies {len(kc):>3}  '
                  f'with a printed range {len(wr)} {pct(len(wr), len(k))}')
            if kind != 'pay_other' and kc:
                c = collections.Counter(r['company'] for r in k)
                print('     by company:', ', '.join(f'{x} {n}' for x, n in c.most_common()))
        print(f'  (below, "printed range" means {key})')
        e_r = [r for r in rs if r[key]]
        e_n = [r for r in rs if not r[key]]
        print(f'  expectation where a range is printed {sum(has(r, "expectation") for r in e_r)} of {len(e_r)} '
              f'{pct(sum(has(r, "expectation") for r in e_r), len(e_r))}; where none is printed '
              f'{sum(has(r, "expectation") for r in e_n)} of {len(e_n)} '
              f'{pct(sum(has(r, "expectation") for r in e_n), len(e_n))}')
        return cos

    cos = block('MAIN: roles open in the US at US-based companies', main_, 'prints_usd_range')
    block('COMPARISON: roles not open in the US, any company', abroad, 'prints_range')
    block('COMPARISON, US-based companies only: their roles not open in the US', [r for r in abroad if r['hq'] == US], 'prints_range')
    block('COMPARISON, companies based elsewhere: roles not open in the US', [r for r in abroad if r['hq'] != US], 'prints_range')

    print('\n######## main sample: the expectation questions themselves')
    qs = collections.Counter()
    attrs = collections.defaultdict(collections.Counter)
    for r in main_:
        for q in r['questions']:
            if q['label'] == 'expectation':
                qs[(r['company'], q['text'], q['required'])] += 1
                for f in ('answer', 'scope', 'currency', 'no_history_warning'):
                    attrs[f][q[f]] += 1
    for (c, t, req), n in sorted(qs.items(), key=lambda kv: -kv[1]):
        print(f'  {n:>3} {c:<14} {"required" if req else "optional":<8} {t[:150]}')
    for f, c in attrs.items():
        print(f'  question-in-role counts by {f}:', dict(c))
    ask_roles = [r for r in main_ if has(r, 'expectation')]
    sc = collections.Counter()
    for r in ask_roles:
        s = {q['scope'] for q in r['questions'] if q['label'] == 'expectation'}
        sc['total' if 'total' in s else 'base' if 'base' in s else 'unspecified'] += 1
    print('  roles by widest scope asked:', dict(sc))

    print('\n######## main sample: company level')
    by = collections.defaultdict(list)
    for r in main_:
        by[r['company']].append(r)
    every = [c for c, rs in by.items() if all(has(r, 'expectation') for r in rs)]
    some = [c for c, rs in by.items() if any(has(r, 'expectation') for r in rs) and c not in every]
    print(f'  companies: ask expectations in every US role {len(every)} {sorted(every)}; in some {len(some)} '
          f'{sorted(some)}; in none {len(by) - len(every) - len(some)}; of {len(by)}')
    print(f'  companies printing a dollar range in every US role {sum(all(r["prints_usd_range"] for r in rs) for rs in by.values())},'
          f' in none {sum(not any(r["prints_usd_range"] for r in rs) for rs in by.values())}')
    big = sorted(by.items(), key=lambda kv: -len(kv[1]))[:10]
    print('  ten largest employers in the sample:', ', '.join(f'{c} {len(rs)}' for c, rs in big))
    top = sum(len(rs) for _, rs in big)
    print(f'  their share of roles {top} {pct(top, len(main_))}')
    rest = [r for c, rs in by.items() if c not in dict(big) for r in rs]
    print(f'  without the ten largest: roles {len(rest)}, expectation {sum(has(r, "expectation") for r in rest)} '
          f'{pct(sum(has(r, "expectation") for r in rest), len(rest))}, prints a dollar range '
          f'{sum(r["prints_usd_range"] for r in rest)} {pct(sum(r["prints_usd_range"] for r in rest), len(rest))}')

    print('\n######## main sample: by board system')
    for api in ('greenhouse', 'ashby', 'lever'):
        rs = [r for r in main_ if r['api'] == api]
        print(f'  {api:<10} roles {len(rs):>5} companies {len({r["company"] for r in rs}):>3}  expectation '
              f'{sum(has(r, "expectation") for r in rs):>4} {pct(sum(has(r, "expectation") for r in rs), len(rs)):>6}  '
              f'prints a dollar range {pct(sum(r["prints_usd_range"] for r in rs), len(rs))}')

    print('\n######## main sample: by where the role is, against laws.json')
    covered = {j['code'] for j in laws['posting_laws'] if j['level'] == 'state'}
    groups = collections.defaultdict(list)
    for r in main_:
        s = set(r['us_states'])
        named = s - {'REMOTE_US'}
        if 'REMOTE_US' in s:
            g = 'remote anywhere in the US offered'
        elif not named:
            g = 'no state named'
        elif named <= covered:
            g = 'only states with a posting law'
        elif named & covered:
            g = 'some states with a posting law'
        else:
            g = 'only states without a posting law'
        groups[g].append(r)
    for g, rs in sorted(groups.items()):
        print(f'  {g:<36} roles {len(rs):>5} companies {len({r["company"] for r in rs}):>3}  prints a dollar range '
              f'{sum(r["prints_usd_range"] for r in rs):>5} {pct(sum(r["prints_usd_range"] for r in rs), len(rs)):>6}  '
              f'expectation {sum(has(r, "expectation") for r in rs):>4} {pct(sum(has(r, "expectation") for r in rs), len(rs))}')
    no_law = groups['only states without a posting law']
    st = collections.Counter(s for r in no_law for s in r['us_states'])
    print('   states in "only states without a posting law":', dict(st.most_common()))
    print('   companies there:', dict(collections.Counter(r['company'] for r in no_law).most_common()))

    # The same company on both sides of the line: companies with roles in
    # "only states with a posting law" and in "only states without one".
    print('   same company, law states against no-law states (dollar range printed / roles):')
    for c in sorted({r['company'] for r in no_law}):
        w = [r for r in groups['only states with a posting law'] if r['company'] == c]
        n = [r for r in no_law if r['company'] == c]
        print(f'     {c:<12} with a law {sum(r["prints_usd_range"] for r in w)}/{len(w)}   '
              f'without {sum(r["prints_usd_range"] for r in n)}/{len(n)}   '
              f'no-law states {sorted({s for r in n for s in r["us_states"]})}')

    bans = {j['code'] for j in laws['history_bans'] if j['level'] == 'state'}
    print('\n######## history questions, anywhere')
    for r in read:
        for q in r['questions']:
            if q['label'] == 'history':
                s = set(r['us_states']) - {'REMOTE_US'}
                print(f'  {r["company"]} | {r["title"]} | {r["places"]} | countries {r["countries"]} | '
                      f'states with a ban named {sorted(s & bans)} | {q["text"]} | required {q["required"]}')
    only_foreign = [r for r in main_ if r['prints_range'] and not r['prints_usd_range']]
    print(f'\n  main-sample roles printing a range only in another currency: {len(only_foreign)} at '
          f'{dict(collections.Counter(r["company"] for r in only_foreign))}')
    warn = [r for r in main_ if any(q.get('no_history_warning') for q in r['questions'])]
    print(f'  main-sample roles whose expectation question warns not to disclose past pay: {len(warn)} '
          f'at {sorted({r["company"] for r in warn})}')

    json.dump(rows, open(HERE / 'role-pay.json', 'w'), ensure_ascii=False, indent=0)
    print(f'\nrole-pay.json: {len(rows)} roles')


if __name__ == '__main__':
    main()