"""B-DAY DROPPER STEP (axis-17, blind-validated on A 2026-09-27:
men +.0040 / doc +.0068, receipt .7706/.8290).

Runs on the FINAL combined B file: trains the corpus-invariant dropper on
train-80 gold (exact model that produced the A receipt), scores B's v24
entities with clean producer vote bits (B's own U2 outputs) + producer-idf
(computed from B's producers), drops p>=0.60 from the entities field ONLY.
Association (task2) untouched. Zero GPU, ~2 min.

Usage (fill paths on B morning):
  python tmp/bday_dropper.py \
    --final  <final combined B submission jsonl> \
    --producers current=<p> new=<p> released=<p> phenobert=<p> \
    --corpus <PatientPheX-B jsonl> \
    --output <filtered final jsonl>

Model = axis17_final.py 'ensemble': GBDT(hand features) 0.5 + LR(TF-IDF) 0.5,
trained on all 80 train docs; features corpus-invariant (NO gold-derived
stats; producer-idf LOO-free on B by construction - B is not in train)."""
import argparse
import json
import re
from bisect import bisect_right
from collections import defaultdict
from pathlib import Path

import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

ROOT = Path(__file__).resolve().parents[1]
THR = 0.60
PNAMES = ['current', 'new', 'released', 'phenobert']
NEG = re.compile(r'(no\s+evidence\s+of|no\s+history\s+of|no\s+sign[s]?\s+of|'
                 r'negative\s+for|free\s+of|absence\s+of|ruled\s+out|'
                 r'without\s+(?:any\s+)?|\bdenied\b|\bdenies\b|\bno\s+)$', re.I)
FAM = re.compile(r'\b(father|mother|sister|brother|daughter|son|grandmother|'
                 r'grandfather|aunt|uncle|cousin|niece|nephew|sibling|'
                 r'family\s+history|relatives?)\b', re.I)


def loadl(p):
    return [json.loads(s) for s in Path(p).read_text(encoding='utf-8-sig').splitlines()
            if s.strip()]


def ents_of(doc):
    return [e for e in doc.get('entities', [])
            if str(e.get('note') or '').upper() != 'NO']


def build_text(d):
    ft = sorted(d['full_text'], key=lambda b: b['offset'])
    buf = [' '] * (ft[-1]['offset'] + len(ft[-1]['text']))
    for b in ft:
        buf[b['offset']:b['offset'] + len(b['text'])] = b['text']
    return ''.join(buf)


# ---------- train side (fixed, deterministic) ----------
GOLD = ROOT / 'data/PatientPheX-V1-A/PatientPheX-train.jsonl'
D = ROOT / 'code/artifacts/task1_vote_topology_20260921'
TPROD = {n: D / (n + '.jsonl') for n in PNAMES}  # new_e8->new alias handled below
TPROD['new'] = D / 'new_e8.jsonl'

train = loadl(GOLD)
gold = {str(d['pmc_id']): {(int(e['offset']), int(e['length']),
        str(e['identifier'])) for e in ents_of(d)} for d in train}
text_by = {str(d['pmc_id']): build_text(d) for d in train}
blocks_by = {str(d['pmc_id']): sorted(d['full_text'], key=lambda b: b['offset'])
             for d in train}
tsrc = {n: {str(d['pmc_id']): ents_of(d) for d in loadl(p)} for n, p in TPROD.items()}

SECTIONS = {}


def sec_idx(s):
    s = str(s or 'UNK')
    if s not in SECTIONS:
        SECTIONS[s] = len(SECTIONS)
    return SECTIONS[s]


def section_at(pmc, off):
    blocks = blocks_by[pmc]
    starts = [b['offset'] for b in blocks]
    i = bisect_right(starts, off) - 1
    if i < 0:
        return 'UNK'
    b = blocks[i]
    return str(b.get('section_type') or 'UNK') \
        if off < b['offset'] + len(b['text']) else 'UNK'


def featurize(pmc, r, pidf, ptop):
    t = r['t']
    o, l, hpo = t
    text = text_by[pmc]
    s = text[o:o + l]
    ws = s.split()
    feats = [float(len(r['voters'])), float(r['nsv']), float(r['nsi']),
             min(6.0, np.log1p(pidf.get(hpo, 0))),
             1.0 if hpo in ptop else 0.0,
             float(l), float(len(ws)),
             1.0 if any(c.isdigit() for c in s) else 0.0,
             1.0 if s[:1].isupper() else 0.0,
             1.0 if len(ws) == 1 else 0.0,
             1.0 if NEG.search(text[max(0, o - 25):o]) else 0.0,
             1.0 if FAM.search(text[max(0, o - 60):o]) else 0.0,
             float(sec_idx(section_at(pmc, o)))]
    feats += [1.0 if n in r['voters'] else 0.0 for n in PNAMES]
    return (feats, s.lower(),
            re.sub(r'\s+', ' ', text[max(0, o - 50):o + l + 30]).lower())


def producer_idf(srcmap, pmcs, exclude=None):
    df = defaultdict(int)
    for pmc in pmcs:
        if pmc == exclude:
            continue
        seen = set()
        for name in PNAMES:
            for e in srcmap[name].get(pmc, []):
                seen.add(str(e['identifier']))
        for h in seen:
            df[h] += 1
    return df, {i for i, c in df.items() if c >= 5}


def v24_meta(per):
    tally = defaultdict(set)
    span_ids, span_voters = defaultdict(set), defaultdict(set)
    for name, ents in per.items():
        for e in ents:
            t = (int(e['offset']), int(e['length']), str(e['identifier']))
            tally[t].add(name)
            span_ids[(t[0], t[1])].add(t[2])
            span_voters[(t[0], t[1])].add(name)
    return [{'t': t, 'voters': v, 'nsv': len(span_voters[(t[0], t[1])]),
             'nsi': len(span_ids[(t[0], t[1])])}
            for t, v in tally.items() if len(v) >= 2]


train_pmcs = [str(d['pmc_id']) for d in train]
Xh, Xs, Xc, y = [], [], [], []
for pmc in train_pmcs:
    # LOO producer-idf: must match axis17_final.py's training matrix exactly
    # (full-corpus idf flips ~3 borderline A triples: 25 vs 28 drops)
    pidf, ptop = producer_idf(tsrc, train_pmcs, exclude=pmc)
    for r in v24_meta({n: tsrc[n][pmc] for n in PNAMES}):
        h, sp, cx = featurize(pmc, r, pidf, ptop)
        Xh.append(h); Xs.append(sp); Xc.append(cx)
        y.append(0 if r['t'] in gold[pmc] else 1)
Xh = np.array(Xh, dtype=float); y = np.array(y)

g = HistGradientBoostingClassifier(max_iter=200, learning_rate=0.08,
                                   min_samples_leaf=20, random_state=0)
g.fit(Xh, y)
tf = TfidfVectorizer(analyzer='char', ngram_range=(2, 4), min_df=3,
                     max_features=30000)
Z = tf.fit_transform(['%s || %s' % (a, b) for a, b in zip(Xs, Xc)])
lr = LogisticRegression(C=4.0, max_iter=2000)
lr.fit(Z, y)
print('train model: %d rows, FP rate %.3f' % (len(y), y.mean()))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--final', required=True)
    ap.add_argument('--corpus', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--current', required=True)
    ap.add_argument('--new', required=True)
    ap.add_argument('--released', required=True)
    ap.add_argument('--phenobert', required=True)
    a = ap.parse_args()

    corpus = loadl(a.corpus)
    for d in corpus:
        pmc = str(d['pmc_id'])
        text_by[pmc] = build_text(d)
        blocks_by[pmc] = sorted(d['full_text'], key=lambda b: b['offset'])
    bsrc = {n: {str(d['pmc_id']): ents_of(d) for d in loadl(getattr(a, n))}
            for n in PNAMES}
    b_pmcs = [str(d['pmc_id']) for d in corpus]
    pidf_b, ptop_b = producer_idf(bsrc, b_pmcs)

    fin = loadl(a.final)
    out, n_drop = [], 0
    for d in fin:
        pmc = str(d['pmc_id'])
        tally = defaultdict(set)
        sid, sv = defaultdict(set), defaultdict(set)
        for name in PNAMES:
            for e in bsrc[name].get(pmc, []):
                t = (int(e['offset']), int(e['length']), str(e['identifier']))
                tally[t].add(name)
                sid[(t[0], t[1])].add(t[2]); sv[(t[0], t[1])].add(name)
        XhB, XsB, XcB = [], [], []
        ents = ents_of(d)
        for e in ents:
            t = (int(e['offset']), int(e['length']), str(e['identifier']))
            r = {'t': t, 'voters': tally[t], 'nsv': len(sv[(t[0], t[1])]),
                 'nsi': len(sid[(t[0], t[1])])}
            h, sp, cx = featurize(pmc, r, pidf_b, ptop_b)
            XhB.append(h); XsB.append(sp); XcB.append(cx)
        if not XhB:
            out.append(d)
            continue
        XhB = np.array(XhB, dtype=float)
        ZB = tf.transform(['%s || %s' % x for x in zip(XsB, XcB)])
        pB = 0.5 * g.predict_proba(XhB)[:, 1] + 0.5 * lr.predict_proba(ZB)[:, 1]
        drop = {tuple(t) for t, p in zip(
            [(int(e['offset']), int(e['length']), str(e['identifier']))
             for e in ents], pB) if p >= THR}
        n_drop += len(drop)
        nd = dict(d)
        nd['entities'] = [e for e in d['entities'] if not (
            str(e.get('note') or '').upper() == 'NO' or
            (int(e['offset']), int(e['length']), str(e['identifier'])) in drop)]
        out.append(nd)

    n0 = sum(len(ents_of(x)) for x in loadl(a.final))
    n1 = sum(len(ents_of(x)) for x in out)
    rate = 100 * n_drop / max(1, n0)
    print('B final: ents %d -> %d (dropped %d, %.1f%%) | train drop-rate ref 2.5%%, A ref 2.0%%'
          % (n0, n1, n_drop, rate))
    assert 0.5 <= rate <= 6.0, 'drop-rate out of validated envelope - INSPECT'
    Path(a.output).write_text(
        '\n'.join(json.dumps(x, ensure_ascii=False) for x in out) + '\n',
        encoding='utf-8')
    print('written:', a.output)


if __name__ == '__main__':
    main()