"""Choose a selection rule on development data only, or apply a frozen rule once (SELECTION.md 4-5).

  python3 select_rule.py choose        # dev: search, print a leaderboard, write rule-frozen.json
  python3 select_rule.py apply         # test: apply rule-frozen.json, write test-result.json

Rules are one threshold, or an AND of two, on S1-S6, in the pre-registered direction only.
Thresholds sit midway between adjacent observed values, so a rule does not sit exactly on a
development cluster's value.
"""
import itertools, json, math, pathlib, sys

PREDICT = {'S1': '>=', 'S2': '>=', 'S3': '>=', 'S4': '>=', 'S5': '>=', 'S6': '<='}

def load(name):
    sig = json.loads(pathlib.Path(f'signals-{name}.json').read_text())
    lab = json.loads(pathlib.Path(f'labels-{name}.json').read_text())
    assert all(lab[c]['label'] in (0, 1, 2) for c in sig), f'labels-{name}.json is incomplete'
    return [dict(sig[c], id=c, label=lab[c]['label'], name=lab[c]['name']) for c in sig]

def passes(row, term):
    s, op, t = term
    return row[s] >= t if op == '>=' else row[s] <= t

def metrics(rows, accept):
    acc = [r for r in rows if accept(r)]
    n2 = sum(r['label'] == 2 for r in acc)
    n1 = sum(r['label'] >= 1 for r in acc)
    total2 = sum(r['label'] == 2 for r in rows)
    return {'accepted': len(acc), 'label2': n2, 'label1plus': n1,
            'precision2': n2 / len(acc) if acc else 0.0,
            'precision1plus': n1 / len(acc) if acc else 0.0,
            'recall2': n2 / total2 if total2 else 0.0}

def midpoints(values):
    v = sorted(set(values))
    return [v[0] - 1e-6] + [(a + b) / 2 for a, b in zip(v, v[1:])] + [v[-1] + 1e-6]

def describe_rule(terms):
    return ' AND '.join(f'{s} {op} {t:.4f}' for s, op, t in terms)

mode = sys.argv[1] if len(sys.argv) > 1 else 'choose'

if mode == 'choose':
    rows = load('dev')
    total2 = sum(r['label'] == 2 for r in rows)
    need = math.ceil(total2 / 2)
    grid = {s: [(s, op, t) for t in midpoints(r[s] for r in rows)] for s, op in PREDICT.items()}
    rules = [[term] for s in PREDICT for term in grid[s]]
    rules += [[a, b] for s1, s2 in itertools.combinations(PREDICT, 2) for a in grid[s1] for b in grid[s2]]
    scored = []
    for terms in rules:
        m = metrics(rows, lambda r: all(passes(r, t) for t in terms))
        if m['label2'] >= need:
            # highest precision; then simpler; then more label-2 accepted; then higher >=1 precision
            scored.append(((m['precision2'], -len(terms), m['label2'], m['precision1plus']), terms, m))
    scored.sort(key=lambda x: x[0], reverse=True)

    print(f'dev: {len(rows)} candidates, {total2} labelled 2, rule must accept >= {need} of them\n')
    for name, accept in (('accept all candidates', lambda r: True), ('B0: tightness filter', lambda r: r['B0_tight'])):
        m = metrics(rows, accept)
        print(f'  baseline {name:26} accepted {m["accepted"]:2}  P2 {m["precision2"]:.2f}  P1+ {m["precision1plus"]:.2f}  R2 {m["recall2"]:.2f}')
    print('\nbest single-signal rule per signal:')
    for s in PREDICT:
        best = next((x for x in scored if len(x[1]) == 1 and x[1][0][0] == s), None)
        if best:
            m = best[2]
            print(f'  {describe_rule(best[1]):22} accepted {m["accepted"]:2}  P2 {m["precision2"]:.2f}  P1+ {m["precision1plus"]:.2f}  R2 {m["recall2"]:.2f}')
    print('\nleaderboard (all eligible rules):')
    seen = set()
    for key, terms, m in scored:
        sig = (describe_rule(terms), m['accepted'])
        if len(seen) >= 12: break
        if sig in seen: continue
        seen.add(sig)
        print(f'  {describe_rule(terms):44} accepted {m["accepted"]:2}  P2 {m["precision2"]:.2f}  P1+ {m["precision1plus"]:.2f}  R2 {m["recall2"]:.2f}')
    key, terms, m = scored[0]
    pathlib.Path('rule-frozen.json').write_text(json.dumps({'terms': terms, 'dev': m}, indent=1))
    print(f'\nCHOSEN: {describe_rule(terms)}  -> rule-frozen.json')
    print('accepted on dev:', ', '.join(f"{r['id']}({r['label']})" for r in rows if all(passes(r, t) for t in terms)))

elif mode == 'apply':
    frozen = json.loads(pathlib.Path('rule-frozen.json').read_text())
    terms = [tuple(t) for t in frozen['terms']]
    rows = load('test')
    accept = lambda r: all(passes(r, t) for t in terms)
    m = metrics(rows, accept)
    print(f'test: {len(rows)} candidates, {sum(r["label"] == 2 for r in rows)} labelled 2')
    print(f'FROZEN RULE {describe_rule(terms)}')
    print(f'  accepted {m["accepted"]}  P2 {m["precision2"]:.2f}  P1+ {m["precision1plus"]:.2f}  R2 {m["recall2"]:.2f}')
    for name, acc in (('accept all candidates', lambda r: True), ('B0: tightness filter', lambda r: r['B0_tight'])):
        b = metrics(rows, acc)
        print(f'  baseline {name:26} accepted {b["accepted"]:2}  P2 {b["precision2"]:.2f}  P1+ {b["precision1plus"]:.2f}  R2 {b["recall2"]:.2f}')
    print('\naccepted:')
    for r in rows:
        if accept(r): print(f"  {r['id']} label {r['label']}  {r['name'] or '-'}")
    print('rejected label-2:')
    for r in rows:
        if not accept(r) and r['label'] == 2: print(f"  {r['id']}  {r['name']}")
    pathlib.Path('test-result.json').write_text(json.dumps({'rule': terms, 'test': m}, indent=1))