#!/usr/bin/env python3
"""Every sentence and every form question that could be about visas or the right to work.

Wide on purpose: a sentence is kept if it names visas, sponsorship, H-1B,
immigration, work authorization, citizenship, permanent residence, export
control or security clearance. What each one means is decided by labeling
(labeling-instructions.md), not here; most "citizenship" hits are the
non-discrimination paragraph and get labeled as not about visas.

Identical text is labeled once: sentences.json and questions.json here hold each
distinct string, how many postings carry it and at which companies.
OUT_DIR/hits.jsonl links every posting to the strings it carries.
Usage: python3 extract.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])

WIDE = re.compile(
    r'\bvisas?\b|sponsor|h-?1-?b\b|immigra|work(ing)? authori[sz]|authori[sz]ed to work|'
    r'eligib\w* to work|right to work|legally (able|permitted|eligible|authori[sz]ed) to work|'
    r'permitted to work|citizen|green card|permanent resident|lawful(ly)? permanent|export control|'
    r'\bitar\b|security clearance|work permit|employment authori[sz]|\bopt\b|\bcpt\b|\bstem opt\b', re.I)
SPLIT = re.compile(r'(?<=[.!?])\s+(?=[A-Z(\"“])')


def norm(s):
    return re.sub(r'\s+', ' ', s).strip()


def sentences(text):
    for line in re.split(r'\n+', text):
        for s in SPLIT.split(line):
            s = norm(s)
            if s and WIDE.search(s):
                yield s[:1200]


def main():
    rows = [json.loads(l) for l in open(OUT / 'eng.jsonl')]
    forms = {}
    p = OUT / 'forms.jsonl'
    if p.exists():
        forms = {f['id']: f for f in map(json.loads, open(p))}
    sent = collections.defaultdict(lambda: {'postings': 0, 'companies': collections.Counter()})
    ques = collections.defaultdict(lambda: {'postings': 0, 'companies': collections.Counter()})
    with open(OUT / 'hits.jsonl', 'w') as f:
        for r in rows:
            ss = list(dict.fromkeys(sentences(r['text'])))
            qs = []
            form = forms.get(r['id'])
            if form and form['read']:
                for q in form['questions']:
                    blob = ' '.join([q['label'], norm(re.sub(r'<[^>]+>', ' ', q['description'])), ' '.join(q['options'])])
                    if WIDE.search(blob):
                        key = norm(q['label']) + (' || options: ' + ' | '.join(q['options']) if q['options'] else '')
                        qs.append(key)
                qs = list(dict.fromkeys(qs))
            for s in ss:
                sent[s]['postings'] += 1
                sent[s]['companies'][r['company']] += 1
            for q in qs:
                ques[q]['postings'] += 1
                ques[q]['companies'][r['company']] += 1
            f.write(json.dumps({'id': r['id'], 'company': r['company'], 'api': r['api'],
                                'form_read': bool(form and form['read']),
                                'sentences': ss, 'questions': qs}, ensure_ascii=False) + '\n')

    def dump(d, name, prefix):
        items = sorted(d.items(), key=lambda kv: (-kv[1]['postings'], kv[0]))
        out = [{'key': f'{prefix}{i:04d}', 'text': k, 'postings': v['postings'],
                'companies': dict(v['companies'].most_common())} for i, (k, v) in enumerate(items, 1)]
        json.dump(out, open(HERE / name, 'w'), ensure_ascii=False, indent=1)
        return len(out)

    print({'postings': len(rows), 'distinct_sentences': dump(sent, 'sentences.json', 's'),
           'distinct_questions': dump(ques, 'questions.json', 'q'),
           'forms_read': sum(1 for f in forms.values() if f['read'])})


if __name__ == '__main__':
    main()