#!/usr/bin/env python3
"""What engineering postings at US tech companies say, and ask, about visa sponsorship.

Reads roles.json, companies.json, labels.json, and OUT_DIR/hits.jsonl (which
sentences and form questions each posting carries). Prints every number the
article states; the output is kept as results.txt.

Main sample: roles open in the United States at companies based in the United
States (companies.json), the same split as the 29 September study.

A role is the unit (roles.py). What a role says is everything its postings say:
- in the text, a statement counts for the US if it names the US or no country;
  a sentence about Canada does not speak for the US opening of the same role.
- stance, when a role says several things: a refusal outranks a transfer-only
  rule, which outranks case by case, which outranks a plain yes. A role that
  says both yes and no is counted as no and listed under conflicts.
- the form (questions.py reads it on all three boards) asks about sponsorship
  if any posting of the role has a question labeled ask_sponsorship or ask_both
  that is about the US, the job's location or no country.
Usage: python3 analyze.py OUT_DIR > results.txt
"""
import collections
import json
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
OUT = Path(sys.argv[1])
US = 'United States'
FOR_US = {US, 'unspecified', 'job location'}
RANK = ['sponsor_no', 'sponsor_no_new', 'sponsor_case', 'sponsor_yes']
STANCE = {'sponsor_no': 'no', 'sponsor_no_new': 'no new', 'sponsor_case': 'case by case', 'sponsor_yes': 'yes'}
ASK = {'ask_sponsorship', 'ask_both'}


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


def main():
    roles = json.load(open(HERE / 'roles.json'))
    comp = json.load(open(HERE / 'companies.json'))
    labels = json.load(open(HERE / 'labels.json'))
    by_text = {('s', v['text']): k for k, v in labels.items() if k.startswith('s')}
    by_text.update({('q', v['text']): k for k, v in labels.items() if k.startswith('q')})
    hits = {h['id']: h for h in map(json.loads, open(OUT / 'hits.jsonl'))}

    def role_view(role):
        text, form, form_read, cites, ask_keys = set(), set(), False, [], set()
        for pid in role['posting_ids']:
            h = hits[pid]
            form_read |= h['form_read']
            for s in h['sentences']:
                k = by_text[('s', s)]
                lab = labels[k]
                if lab['label'] in STANCE and lab['country'] in FOR_US:
                    text.add(lab['label'])
                    cites.append(k)
                elif lab['label'] == 'auth_required' and lab['country'] in FOR_US:
                    text.add('auth_required')
                elif lab['label'] == 'citizen_required' and lab['country'] in FOR_US:
                    text.add('citizen_required')
            for q in h['questions']:
                k = by_text[('q', q)]
                lab = labels[k]
                if lab['country'] in FOR_US:
                    form.add(lab['label'])
                    if lab['label'] in ASK:
                        ask_keys.add(k)
        stance = next((STANCE[l] for l in RANK if l in text), 'silent')
        return {'stance': stance, 'auth_only': stance == 'silent' and 'auth_required' in text,
                'citizen': 'citizen_required' in text, 'form_read': form_read,
                'asks': bool(form & ASK), 'asks_auth': bool(form & {'ask_authorization', 'ask_both'}),
                'conflict': 'sponsor_yes' in text and 'sponsor_no' in text, 'cites': cites,
                'ask_keys': ask_keys}

    main_roles = [r for r in roles if comp[r['company']]['hq'] == US and US in r['countries']]
    views = {r['id']: role_view(r) for r in main_roles}
    # Per role, what it says and asks: the table the appendix prints. Written
    # here because the posting text it is read from stays outside the repo.
    keep = []
    for r in main_roles:
        v = views[r['id']]
        keep.append({'id': r['id'], 'company': r['company'], 'title': r['title'], 'url': r['url'],
                     'board': {'gh': 'greenhouse', 'ab': 'ashby', 'lv': 'lever'}[r['id'][:2]],
                     'seniority': r['seniority'], 'postings': r['postings'], 'countries': r['countries'],
                     'text_stance': v['stance'], 'text_requires_us_authorization': v['auth_only'],
                     'text_citizenship_or_clearance': v['citizen'], 'sentences': sorted(set(v['cites'])),
                     'form_read': v['form_read'], 'form_asks_sponsorship': v['asks'],
                     'form_asks_authorization': v['asks_auth'], 'questions': sorted(v['ask_keys'])})
    json.dump(keep, open(HERE / 'role-stances.json', 'w'), ensure_ascii=False, indent=0)
    n = len(main_roles)
    cos = sorted({r['company'] for r in main_roles})
    print(f'######## sample on 2026-10-13')
    print(f'all roles {len(roles)}, companies {len({r["company"] for r in roles})}')
    print(f'main sample: roles open in the US at US-based companies {n}, companies {len(cos)}')

    print('\n######## what the posting text says, role level')
    st = collections.Counter(v['stance'] for v in views.values())
    for k in ['yes', 'case by case', 'no new', 'no', 'silent']:
        print(f'  {k:<14}{st[k]:>6}  {pct(st[k], n)}')
    said = n - st['silent']
    print(f'  any sponsorship statement {said} {pct(said, n)}')
    auth = sum(v['auth_only'] for v in views.values())
    big = [c for c, n_ in collections.Counter(r['company'] for r in main_roles
                                              if views[r['id']]['stance'] != 'silent').most_common(2)]
    rest = [r for r in main_roles if r['company'] not in big]
    rs_ = sum(views[r['id']]['stance'] != 'silent' for r in rest)
    print(f'  without the two companies with most statements {big}: {rs_} of {len(rest)} {pct(rs_, len(rest))}')
    print(f'  silent on sponsorship but requires existing US work authorization {auth} {pct(auth, n)}')
    cit = sum(v['citizen'] for v in views.values())
    print(f'  citizenship / clearance / export-control line {cit} {pct(cit, n)}')
    print(f'  conflicts (yes and no in one role): {[r["id"] for r in main_roles if views[r["id"]]["conflict"]]}')

    print('\n######## company level: what a company says in any of its US roles')
    stance_of = collections.defaultdict(collections.Counter)
    for r in main_roles:
        stance_of[r['company']][views[r['id']]['stance']] += 1
    groups = collections.defaultdict(list)
    for c in cos:
        said = {s for s in stance_of[c] if s != 'silent'}
        key = ('silent' if not said else 'refuses (no or no new)' if said <= {'no', 'no new'}
               else 'offers (yes or case by case)' if said <= {'yes', 'case by case'} else 'mixed')
        groups[key].append(c)
    for k, v in sorted(groups.items(), key=lambda kv: -len(kv[1])):
        print(f'  {k:<30}{len(v):>4}  {pct(len(v), len(cos))}')
        if k != 'silent':
            for c in v:
                s = stance_of[c]
                print(f'      {c:<16} roles {sum(s.values()):>4}  ' + ', '.join(f'{a} {b}' for a, b in s.most_common()))

    print('\n######## the application form (roles whose form was read)')
    gh = [r for r in main_roles if views[r['id']]['form_read']]
    g = len(gh)
    asks = [r for r in gh if views[r['id']]['asks']]
    print(f'  roles {g}, companies {len({r["company"] for r in gh})}')
    print(f'  form asks about sponsorship {len(asks)} {pct(len(asks), g)}')
    top = collections.Counter(r['company'] for r in gh).most_common(1)[0][0]
    gh2 = [r for r in gh if r['company'] != top]
    a2 = sum(views[r['id']]['asks'] for r in gh2)
    print(f'  without the company with most roles ({top}): {a2} of {len(gh2)} {pct(a2, len(gh2))}')
    authq = sum(1 for r in gh if views[r['id']]['asks_auth'])
    print(f'  form asks about work authorization {authq} {pct(authq, g)}')
    silent_ask = [r for r in asks if views[r['id']]['stance'] == 'silent']
    print(f'  text silent, form asks {len(silent_ask)} {pct(len(silent_ask), g)}; '
          f'share of asking roles whose text is silent {pct(len(silent_ask), len(asks))}')
    gst = collections.Counter(views[r['id']]['stance'] for r in gh)
    for api in ('gh', 'ab', 'lv'):
        sub = [r for r in gh if r['id'].startswith(api + '-')]
        a = sum(views[r['id']]['asks'] for r in sub)
        print(f'    board {api}: roles {len(sub):>5}, companies {len({r["company"] for r in sub}):>3}, form asks {a:>5} {pct(a, len(sub))}')
    print('  text stance among roles with a form read: ' + ', '.join(f'{k} {v}' for k, v in gst.most_common()))
    for s in ['yes', 'case by case', 'no new', 'no', 'silent']:
        sub = [r for r in gh if views[r['id']]['stance'] == s]
        a = sum(views[r['id']]['asks'] for r in sub)
        print(f'    stance {s:<13} roles {len(sub):>5}  form asks {a:>5}  {pct(a, len(sub))}')
    gco = collections.defaultdict(lambda: [0, 0])
    for r in gh:
        gco[r['company']][0] += 1
        gco[r['company']][1] += views[r['id']]['asks']
    every = [c for c, (t, a) in gco.items() if a == t]
    none = [c for c, (t, a) in gco.items() if a == 0]
    print(f'  companies: ask in every US role {len(every)}, in none {len(none)}, in some {len(gco) - len(every) - len(none)}, of {len(gco)}')
    print(f'    never ask: {sorted(none)}')

    print('\n######## by seniority, main sample')
    for lvl in ['intern', 'junior', 'mid', 'senior', 'staff']:
        sub = [r for r in main_roles if r['seniority'] == lvl]
        if not sub:
            continue
        sv = collections.Counter(views[r['id']]['stance'] for r in sub)
        ref = sv['no'] + sv['no new']
        off = sv['yes'] + sv['case by case']
        ghs = [r for r in sub if views[r['id']]['form_read']]
        a = sum(views[r['id']]['asks'] for r in ghs)
        print(f'  {lvl:<7} roles {len(sub):>5}  refuses {ref:>4} {pct(ref, len(sub)):>6}  offers {off:>4} {pct(off, len(sub)):>6}'
              f'  form asks {a}/{len(ghs)} {pct(a, len(ghs))}')

    print('\n######## where the refusing and the offering companies have roles open')
    for k in ['refuses (no or no new)', 'offers (yes or case by case)', 'mixed', 'silent']:
        cs = set(groups.get(k, []))
        rs = [r for r in roles if r['company'] in cs and r['countries']]
        abroad = [r for r in rs if any(c != US for c in r['countries'])]
        only = [r for r in rs if US not in r['countries']]
        print(f'  {k:<30} companies {len(cs):>3}  roles with a country {len(rs):>5}  open outside the US {len(abroad):>5} '
              f'{pct(len(abroad), len(rs)):>6}  not open in the US {len(only):>5} {pct(len(only), len(rs)):>6}')

    print('  per company, median of its own share (companies with 5+ roles with a country)')
    for k in ['refuses (no or no new)', 'offers (yes or case by case)', 'silent']:
        shares = []
        for c in groups.get(k, []):
            rs = [r for r in roles if r['company'] == c and r['countries']]
            if len(rs) >= 5:
                shares.append(sum(any(x != US for x in r['countries']) for r in rs) / len(rs))
        shares.sort()
        med = shares[len(shares) // 2] if shares else 0
        print(f'    {k:<30} companies {len(shares):>3}  median share open outside the US {100 * med:.0f}%')

    print('\n######## how the form asks (roles whose form asks about sponsorship)')
    import re
    def has(pat, r):
        return any(re.search(pat, labels[k]['text'], re.I) for k in views[r['id']]['ask_keys'])
    for name, pat in [('now or in the future', r'future|ever\b|at any point|continue|extend|retain|maintain'),
                      ('names H-1B or another visa type', r'h-?1-?b|\bopt\b|\bcpt\b|\btn\b|o-1|l-1|e-3|f-1')]:
        c = sum(has(pat, r) for r in asks)
        print(f'  {name:<34}{c:>5}  {pct(c, len(asks))}')
    # Answer options are listed on Greenhouse and Lever; Ashby's yes/no fields carry none.
    listed = [r for r in asks if not r['id'].startswith('ab-')]
    c = sum(has(r'\|\| options: (?:Yes \| No|No \| Yes)$', r) for r in listed)
    print(f'  yes/no only, of {len(listed)} on Greenhouse and Lever{c:>6}  {pct(c, len(listed))}')

    # Whether the sponsorship question must be answered: the form's own required flag.
    import re as _re
    forms = {f['id']: f for f in map(json.loads, open(OUT / 'forms.jsonl'))}
    ask_text = {labels[k]['text'] for r in asks for k in views[r['id']]['ask_keys']}
    def norm(t):
        return _re.sub(r'\s+', ' ', t).strip()
    req = 0
    for r in asks:
        req += any(q['required'] and (norm(q['label']) + (' || options: ' + ' | '.join(q['options'])
                                                          if q['options'] else '')) in ask_text
                   for pid in r['posting_ids'] for q in forms[pid].get('questions', []))
    print(f'  the sponsorship question is required {req} of {len(asks)} {pct(req, len(asks))}')

    print('\n######## companies: the ad against the form, company level')
    gh_cos = sorted({r['company'] for r in gh})
    tab = collections.Counter()
    for c in gh_cos:
        rs = [r for r in gh if r['company'] == c]
        says = any(views[r['id']]['stance'] != 'silent' for r in rs)
        ask = any(views[r['id']]['asks'] for r in rs)
        tab[(says, ask)] += 1
    for (says, ask), v in sorted(tab.items()):
        print(f'  ad {"says something" if says else "silent everywhere":<18} form {"asks" if ask else "never asks":<11}{v:>4}')


if __name__ == '__main__':
    main()
