TheLadder/tools/story-to-pack/probe/review.py

"""Build a review page: the groups gate G2 accepts, plus the nearest misses.

  python3 review.py                                  # k = 60, seed 24 -> review/k60-s24/
  python3 review.py --k 60 --seed 31 --borderline 8

Writes into the run directory:
  REVIEW.html      open it in a browser: the whole review happens there
  candidates.json  what was found, for apply_review.py to check the saved review against

When the review is saved from the page:
  python3 apply_review.py review/k60-s24   ->   situations.json

No inference: this uses the summaries and embeddings already on disk.
"""
import argparse, datetime, hashlib, json, pathlib
import gate2 as g
import review_format as rf
import review_page


def gate_failures(sig, gate):
    out = []
    if gate.get('S1') is not None and sig['S1'] < gate['S1']:
        out.append(f"cohesion S1 {sig['S1']:.3f}, needs {gate['S1']:.3f}")
    if gate.get('Z1') is not None and sig['Z1'] < gate['Z1']:
        out.append(f"cohesion against same-size random sets Z1 {sig['Z1']:.1f}, needs {gate['Z1']}")
    if gate.get('MMIN') is not None and sig['MMIN'] < gate['MMIN']:
        out.append(f"weakest scene MMIN {sig['MMIN']:.3f}, needs {gate['MMIN']:.3f}")
    if gate.get('W') is not None and sig['W'] >= gate['W']:
        out.append(f"the word \"{sig['W_word']}\" is in {sig['W']:.0%} of summaries, limit {gate['W']:.0%}")
    return out


def shortfall(sig, gate):
    """How far a rejected group is from passing; smaller is closer."""
    s = 0.0
    if gate.get('S1') is not None: s += max(0.0, gate['S1'] - sig['S1']) / 0.01
    if gate.get('Z1') is not None: s += max(0.0, gate['Z1'] - sig['Z1'])
    if gate.get('W') is not None and sig['W'] >= gate['W']: s += 10 * (sig['W'] - gate['W'] + 0.01)
    return s


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--k', type=int, default=60, help='number of k-means clusters (default 60)')
    ap.add_argument('--seed', type=int, default=24, help='k-means seed (default 24)')
    ap.add_argument('--borderline', type=int, default=8, help='near misses offered as optional extras (default 8)')
    ap.add_argument('--out', help='run directory (default review/k<K>-s<SEED>)')
    args = ap.parse_args()

    run_id = f'k{args.k}-s{args.seed}'
    out = pathlib.Path(args.out or f'review/{run_id}')

    gate = json.loads(pathlib.Path('gate2-frozen.json').read_text())
    found = [c for c in g.kmeans(args.k, args.seed) if g.eligible(c)]
    scored = [(c, g.signals(c)) for c in found]
    accepted = sorted([cs for cs in scored if g.g2_passes(cs[1], gate)], key=lambda cs: -cs[1]['S1'])
    rejected = sorted([cs for cs in scored if not g.g2_passes(cs[1], gate)], key=lambda cs: shortfall(cs[1], gate))
    borderline = rejected[:args.borderline]

    listed = [(f'C{i + 1:02d}', 'accepted', c, s) for i, (c, s) in enumerate(accepted)]
    listed += [(f'B{i + 1:02d}', 'borderline', c, s) for i, (c, s) in enumerate(borderline)]

    candidates = {}
    for cid, kind, members, sig in listed:
        member_rows = [{'member': rf.member_id(cid, n + 1), 'chunk': chunk, 'story': g.chunks[chunk]['story'],
                        'title': g.chunks[chunk]['title'].title(), 'summary': g.text[chunk]}
                       for n, chunk in enumerate(sorted(members))]
        candidates[cid] = {
            'kind': kind,
            'stats': {'scenes': sig['size'], 'stories': sig['stories'], 'largest_story_share': round(sig['dominant'], 3),
                      'S1': round(sig['S1'], 4), 'Z1': round(sig['Z1'], 2), 'W': round(sig['W'], 3), 'W_word': sig['W_word']},
            'gate_failures': gate_failures(sig, gate),
            'role_hints': review_page.role_hints(member_rows),
            'members': member_rows,
        }
    # Ties a saved review to exactly these groups, so progress for a different set of groups
    # that happens to share the run name is never loaded into this page.
    fingerprint = hashlib.sha1(json.dumps({cid: [m['chunk'] for m in c['members']] for cid, c in candidates.items()},
                                          sort_keys=True).encode()).hexdigest()[:12]
    document = {
        'schema_version': rf.SCHEMA_VERSION,
        'run': {'id': run_id, 'k': args.k, 'seed': args.seed, 'fingerprint': fingerprint,
                'created': datetime.datetime.now().astimezone().isoformat(timespec='seconds'),
                'gate': {'name': 'G2', **gate}, 'scenes': g.n, 'stories': len(set(g.story)),
                'candidates_found': len(found), 'accepted': len(accepted), 'borderline': len(borderline)},
        'candidates': candidates,
    }

    out.mkdir(parents=True, exist_ok=True)
    (out / 'candidates.json').write_text(json.dumps(document, indent=1, ensure_ascii=False), encoding='utf-8')
    (out / 'REVIEW.html').write_text(review_page.render_page(document), encoding='utf-8')

    print(f'{run_id}: {len(found)} candidate groups, G2 accepts {len(accepted)}, {len(borderline)} optional extras')
    print(f'open this in a browser:  {(out / "REVIEW.html").resolve()}')
    print(f'when the review is saved:  python3 apply_review.py {out}')


if __name__ == '__main__':
    main()