"""Choose G2's thresholds on every labelled cluster so far (SELECTION.md, "A size-invariant gate").

  python3 tune_gate2.py        # writes gate2-tuning.json and gate2-frozen.json
"""
import itertools, json, pathlib
import gate2 as g

# Guard: G2's S1 and old S6 must reproduce the first gate's signals.
dev_sig = json.loads(pathlib.Path('signals-dev.json').read_text())
dev_cand = json.loads(pathlib.Path('candidates-dev.json').read_text())
for cid, m in dev_cand.items():
    s = g.signals(m)
    assert abs(s['S1'] - dev_sig[cid]['S1']) < 1e-3 and abs(s['S6_old'] - dev_sig[cid]['S6']) < 1e-3, cid

rows = []
for name in ('dev', 'test'):
    cand = json.loads(pathlib.Path(f'candidates-{name}.json').read_text())
    lab = json.loads(pathlib.Path(f'labels-{name}.json').read_text())
    for cid, m in cand.items():
        rows.append(dict(g.signals(m), id=cid, label=lab[cid]['label']))
rmap = json.loads(pathlib.Path('recover-map.json').read_text())
rlab = json.loads(pathlib.Path('labels-recover.json').read_text())
for rid, e in rmap.items():
    rows.append(dict(g.signals(e['members']), id=rid, label=rlab[rid]['label']))
pathlib.Path('gate2-tuning.json').write_text(json.dumps(rows, indent=1))

def deciles(key):
    v = sorted(r[key] for r in rows)
    return [None] + [v[int(len(v) * q / 10)] for q in range(1, 10)]

grid = itertools.product(deciles('S1'), [None, 1, 2, 3, 4, 5], deciles('MMIN'),
                         [None, 0.25, 0.30, 0.35, 0.40, 0.50, 0.60])
best = None
for s1, z1, mmin, w in grid:
    gate = {'S1': s1, 'Z1': z1, 'MMIN': mmin, 'W': w}
    acc = [r for r in rows if g.g2_passes(r, gate)]
    if not acc: continue
    n1 = sum(r['label'] >= 1 for r in acc)
    if n1 / len(acc) < 0.85: continue
    key = (n1, sum(r['label'] == 2 for r in acc), -sum(v is not None for v in gate.values()))
    if best is None or key > best[0]:
        best = (key, gate)

gate = best[1]
pathlib.Path('gate2-frozen.json').write_text(json.dumps(gate, indent=1))

def band(size):
    return '<=12' if size <= 12 else ('13-25' if size <= 25 else '>=26')

def report(label, accept):
    acc = [r for r in rows if accept(r)]
    n1 = sum(r['label'] >= 1 for r in acc)
    n2 = sum(r['label'] == 2 for r in acc)
    return f'{label:14} accepted {len(acc):3}  label>=1 {n1:3}  P1+ {n1 / len(acc) if acc else 0:.2f}  label2 {n2}'

print(f'tuning set: {len(rows)} clusters, {sum(r["label"] >= 1 for r in rows)} labelled >=1, '
      f'{sum(r["label"] == 2 for r in rows)} labelled 2')
print('CHOSEN G2:', gate)
print(report('all', lambda r: True))
print(report('old gate', g.old_passes))
print(report('G2', lambda r: g.g2_passes(r, gate)))
print('\nby size band:')
for b in ('<=12', '13-25', '>=26'):
    sub = [r for r in rows if band(r['size']) == b]
    print(f'  {b}: {len(sub)} clusters, {sum(r["label"] >= 1 for r in sub)} labelled >=1')
    for label, accept in (('old gate', g.old_passes), ('G2', lambda r: g.g2_passes(r, gate))):
        acc = [r for r in sub if accept(r)]
        n1 = sum(r['label'] >= 1 for r in acc)
        print(f'    {label:9} accepted {len(acc):3}  P1+ {n1 / len(acc) if acc else 0:.2f}')
print('\nG2 rejects, labelled >=1, with the failing term:')
for r in rows:
    if r['label'] >= 1 and not g.g2_passes(r, gate):
        why = [k for k in ('S1', 'Z1', 'MMIN') if gate[k] is not None and r[k] < gate[k]]
        if gate['W'] is not None and r['W'] >= gate['W']: why.append(f"W({r['W_word']})")
        print(f"  {r['id']} size {r['size']:2} label {r['label']}  {' '.join(why) or 'candidate conditions'}")