#!/usr/bin/env python3
"""Census the T8aD paint-order key across every sprite on the disc.

`structures/ui-paint-order-key.md` established that this field sorts a screen's
paint order, from twelve values on two measured screens.  This walks all 21 184
sprites so the open question -- what the bits mean -- is asked of the corpus.

    ./paint_key_census.py <extract root>

The field is a **u16 at +0x0A**: the doc reads a 32-bit word at +0x08 and its
upper half is zero in every sprite.

Two traps this script exists to avoid, both of which produced a wrong published
number before it did:

  * **Sprites are mostly RATC children, not pak entries.**  A bundle entry's
    magic is `RATC`, so filtering pak entries on a `T8aD` magic finds only
    4 525 of the 21 184 sprites and 45 of the 216 keys -- and misses GP_TITLE
    entirely, which is the screen the page's own evidence comes from.
  * **The language variants of GP_MAIN_GAME_*2D.pak are the same screens six
    times.**  Counting them separately makes every one of their keys look shared
    across six paks.  They are collapsed into one family here.
"""

import os, sys, glob, struct, re, collections
from disc import disc_root

_SD = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _SD)
exec(open(os.path.join(_SD, "regn_decode.py")).read().split("# ── the object")[0])
Hh=lambda b,o: struct.unpack_from(">H",b,o)[0]
def fam(n): return re.sub(r'_[DEFIJS]2D\.pak$','_2D.pak',n)
keys=collections.Counter(); kfam=collections.defaultdict(set)
hi_zero=tot=0; langsets=collections.defaultdict(set)
def take(b,o,base,lang):
    global hi_zero,tot
    tot+=1
    if Hh(b,o+8)==0: hi_zero+=1
    k=Hh(b,o+10); keys[k]+=1; kfam[k].add(base)
    if lang: langsets[lang].add(k)
ROOT = sys.argv[1] if len(sys.argv) > 1 else disc_root()
for f in sorted(glob.glob(ROOT + "/**/*.pak", recursive=True)):
    try: E=pak_entries(f)
    except Exception: continue
    nm=f.split('/')[-1]; base=fam(nm); lang=nm if nm.endswith('2D.pak') else None
    for h,b in (E.items() if isinstance(E,dict) else E):
        m=bytes(b[:4])
        if m==b'T8aD': take(b,0,base,lang)
        elif m==b'RATC':
            i=b.find(b'T8aD',4)
            while i!=-1:
                if i+12<=len(b): take(b,i,base,lang)
                i=b.find(b'T8aD',i+4)
print(f"ALL T8aD sprites (top-level + RATC children): {tot}")
print(f"upper half of the +0x08 word is zero        : {hi_zero}/{tot}")
print(f"distinct u16 keys at +0x0A                  : {len(keys)}")
fams=sorted({f for s in kfam.values() for f in s})
multi=[k for k in keys if len(kfam[k])>1]
print(f"keys in more than one pak family            : {len(multi)}/{len(keys)} = {100*len(multi)/len(keys):.0f}%")
only=[k for k in keys if kfam[k]=={'GP_MAIN_GAME_2D.pak'}]
print(f"keys confined to GP_MAIN_GAME_2D            : {len(only)}/{len(keys)}")
print(f"the six language 2D paks have identical key sets: {len({frozenset(v) for v in langsets.values()})==1} ({len(langsets)} paks)")
print(f"\nkeys per family:")
for f in fams:
    ks=sorted(k for k in keys if f in kfam[k])
    print(f"  {f:>28}: {len(ks):3d} keys  {ks[0]:04x}..{ks[-1]:04x}")
print(f"\ncommonest 12 keys: {[('%04x'%k,n) for k,n in keys.most_common(12)]}")