#!/usr/bin/env python3
"""Check everything, then write records.json, groups.json and options_by_entry.json, and print the pilot
measurement. Adapted from ../../jacobs/claude/build.py: adds the advice field, 'advice:N' option sources,
and the advice-book pilot (options_src.py -> options_by_entry.json). No inference."""
import collections, glob, json, pathlib, re
from groups_src import G, STANDING
from options_src import O, NEW_ENTRIES
here = pathlib.Path(__file__).parent
probe = here.parent.parent
chunks = json.loads((here.parent / 'chunks.json').read_text())
stories = json.loads((here.parent / 'stories.json').read_text())
catalogue = {e['id']: e for e in json.loads((probe / 'catalogue/v2/catalogue.json').read_text())['entries']}
recs = [r for f in sorted(glob.glob(str(here / 'batches/s*.json'))) for r in json.loads(open(f).read())]
norm = lambda s: re.sub(r'\s+', ' ', s).strip()

# ---- batches ------------------------------------------------------------------------------------------
assert [r['chunk'] for r in recs] == list(range(len(chunks))), 'every scene, once, in order'
for r in recs:
    assert chunks[r['chunk']]['story'] == r['story'], r['chunk']
    assert set(r) == {'chunk', 'story', 'summary', 'predicament', 'advice'}, r['chunk']
    p = r['predicament']
    assert p.endswith('.') and '. ' not in p and 12 <= len(p.split()) <= 26, (r['chunk'], p)
    for a in r['advice']:
        assert set(a) == {'move', 'outcome', 'quote', 'stance'} and a['stance'] in ('do', 'avoid'), (r['chunk'], a)
        assert norm(a['quote']) in norm(chunks[r['chunk']]['text']), (r['chunk'], 'quote not verbatim', a['quote'])
        assert len(a['quote'].split()) <= 30, (r['chunk'], 'quote over 30 words')
by = {r['chunk']: r for r in recs}

# ---- groups -------------------------------------------------------------------------------------------
DIRS = {'+', '-', 'risk', 'set'}
SKILLS = {'judgment', 'hustle'}
def check_moves(mv, where):
    for k, v in mv.items():
        assert v in DIRS, (where, k, v)
        assert k in ('money', 'standing') or k in SKILLS or k.startswith(('relationships.', 'flags.')), (where, k)
        assert (v == 'set') == k.startswith('flags.'), (where, k, v)
seen = collections.Counter(m for g in G for m in g['members'])
dup = [m for m, n in seen.items() if n > 1]
assert not dup, f'scenes in two groups: {dup}'
out = []
for g in G:
    assert 3 <= len(g['options']) <= 5, g['id']
    st = sorted({by[m]['story'] for m in g['members']})
    assert len(st) >= 3, (g['id'], 'fewer than three letters')
    top = collections.Counter(by[m]['story'] for m in g['members']).most_common(1)[0][1]
    assert top <= len(g['members']) / 2, (g['id'], 'one letter supplies over half')
    stand = collections.Counter(STANDING[by[m]['story']] for m in g['members'])
    opts = []
    for t, s, mv in g['options']:
        check_moves(mv, g['id'])
        kind, _, n = s.partition(':')
        assert kind in ('scene', 'advice', 'archetype', 'filled'), s
        if kind == 'advice': assert by[int(n)]['advice'], (g['id'], s, 'scene has no advice')
        opts.append({'text': t, 'source': s, 'moves': mv,
                     'in_group': (int(n) in g['members']) if kind in ('scene', 'advice') else None})
    out.append({
        'id': g['id'], 'name': g['name'], 'predicament': g['predicament'],
        'sentence': f"A {g['actor']} wants {g['wants']} from {g['counterpart']}.",
        'actor': g['actor'], 'wants': g['wants'], 'counterpart': g['counterpart'],
        'stories': st, 'story_titles': [stories[s]['title'] for s in st],
        'standing': dict(stand), 'stage': stand.most_common(1)[0][0],
        'members': [{'chunk': m, 'story': by[m]['story'], 'predicament': by[m]['predicament']} for m in g['members']],
        'options': opts,
    })
out.sort(key=lambda g: int(g['id'][1:]))
unassigned = [r['chunk'] for r in recs if r['chunk'] not in seen]
(here / 'records.json').write_text(json.dumps(recs, indent=1, ensure_ascii=False))
(here / 'groups.json').write_text(json.dumps({'groups': out, 'unassigned': unassigned}, indent=1, ensure_ascii=False))

# ---- the pilot: options_by_entry.json -----------------------------------------------------------------
entries, used = collections.defaultdict(list), set()
for key, c, i, text, mv, stated in O:
    if key.startswith('new:'): assert key[4:] in NEW_ENTRIES, key
    else: assert key in catalogue, f'unknown catalogue entry {key}'
    assert (c, i) not in used, f'advice {c}:{i} used twice'
    used.add((c, i))
    check_moves(mv, key)
    assert isinstance(stated, bool)
    entries[key].append({'text': text, 'moves': mv, 'source': f'chunk:{c}', 'evidence': by[c]['advice'][i]['quote'],
                         'stated': stated})
assert set(NEW_ENTRIES) == {k[4:] for k in entries if k.startswith('new:')}, 'a new entry with no options'
all_items = [(r['chunk'], i) for r in recs for i in range(len(r['advice']))]
unmapped = [x for x in all_items if x not in used]
old = sorted((k for k in entries if not k.startswith('new:')), key=lambda k: list(catalogue).index(k))
new = sorted(k for k in entries if k.startswith('new:'))
obe = {k: entries[k] for k in old + new}
obe['_new_entries'] = {k: NEW_ENTRIES[k] for k in sorted(NEW_ENTRIES)}
(here / 'options_by_entry.json').write_text(json.dumps(obe, indent=1, ensure_ascii=False))

# ---- report -------------------------------------------------------------------------------------------
print(f'{len(recs)} scenes, {len(out)} groups, {sum(seen.values())} assigned, {len(unassigned)} unassigned: {unassigned}')
for g in out:
    print(f"{g['id']} {len(g['members']):2} scenes {len(g['stories'])} letters  stage {g['stage']:6}  {g['name']}")
src = collections.Counter(o['source'].split(':')[0] for g in out for o in g['options'])
print('group option sources:', dict(src), '| scene/advice options citing a scene outside their group:',
      sum(o['in_group'] is False for g in out for o in g['options']))
adv = [a for r in recs for a in r['advice']]
print(f"\nPILOT  scenes with advice {sum(bool(r['advice']) for r in recs)} of {len(recs)}; advice items {len(adv)} "
      f"({sum(a['stance'] == 'do' for a in adv)} do, {sum(a['stance'] == 'avoid' for a in adv)} avoid)")
opts = [o for k in old + new for o in entries[k]]
st = lambda os: sum(o['stated'] for o in os)
print(f'  options {len(opts)} from {len(used)} advice items ({len(unmapped)} items unmapped: {unmapped})')
print(f'  stated {st(opts)} of {len(opts)} = {st(opts)/len(opts):.0%}')
o_old = [o for k in old for o in entries[k]]; o_new = [o for k in new for o in entries[k]]
print(f'  catalogue entries covered {len(old)} of {len(catalogue)}: {len(o_old)} options, stated {st(o_old)} ({st(o_old)/len(o_old):.0%})')
for k in old: print(f'    {k:26s} {len(entries[k]):2d} options, {st(entries[k])} stated')
print(f'  new: entries {len(new)}: {len(o_new)} options, stated {st(o_new)} ({st(o_new)/len(o_new):.0%})')
for k in new: print(f'    {k:26s} {len(entries[k]):2d} options, {st(entries[k])} stated')
print('\nSTORY PACKS for comparison (option sources in groups.json):')
for corpus in ('lawson', 'bierce'):
    gs = json.loads((probe / corpus / 'claude/groups.json').read_text())['groups']
    c = collections.Counter(o['source'].split(':')[0] for g in gs for o in g['options'])
    print(f'  {corpus}: {sum(c.values())} options in {len(gs)} groups: {dict(c)}')