#!/usr/bin/env python3
"""Census every AIParams object on the disc, not just Stage 02's.

`stage-mission-tables.md` documents `AIParams_S02.tbl` -- 34 profiles, two
shapes.  This checks that against the whole disc: how many such objects exist,
whether the 34-name roster is shared, and whether `Type` really predicts the
field count.  Regenerates docs/re/data/aiparams-census.txt.
"""
import sys, os, glob, collections

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import unitgroup as U
from unit_substructures import pak_entries

def objects(pak):
    out = []
    for h, b in pak_entries(pak):
        if b[:4] != b'IDXD' or b'Enumerate_AIs' not in b:
            continue
        try:
            recs = U.parse(b)
        except Exception:
            continue
        if 'Enumerate_AIs' not in {r['squadron'] for r in recs}:
            continue
        decl = [n for r in recs if r['squadron'] == 'Enumerate_AIs'
                for _t, n, _v in r['fields'] if n]
        prof = [(r['squadron'], dict((n, v) for _t, n, v in r['fields'] if n))
                for r in recs if r['squadron'] != 'Enumerate_AIs']
        out.append((h, decl, prof))
    return out

def main():
    paks = sorted(p for p in glob.glob('/work/sylph_extract/**/*.pak', recursive=True)
                  if os.path.basename(p).startswith('GP_MAIN_GAME_')
                  and '2D' not in os.path.basename(p))
    print("# AIParams across the whole disc")
    print("# Regenerate: python3 tools/re-capture/aiparams_census.py")
    print("# See docs/re/structures/stage-mission-tables.md")

    print("\n## CONTROL  the six language copies")
    per = {}
    for pk in paks:
        o = objects(pk)
        per[os.path.basename(pk)] = o
        print("   %-24s %2d AIParams objects" % (os.path.basename(pk), len(o)))

    objs = per[os.path.basename(paks[0])]
    sets = {tuple(sorted(d)) for _h, d, _p in objs}
    print("\n## The roster")
    print("   objects: %d   distinct declared-name sets: %d   roster size: %s"
          % (len(objs), len(sets), sorted({len(s) for s in sets})))
    absent = sum(1 for _h, d, p in objs
                 for n in d if n not in {s for s, _f in p})
    print("   declared names with no record in their own object: %d" % absent)
    print("   profile records total: %d  (= %d objects x %d)"
          % (sum(len(p) for _h, _d, p in objs), len(objs),
             len(objs[0][2]) if objs else 0))

    # Which object is which stage?  `StageResource.EnumerateAIParams` names it.
    stageres = []
    for h, b in pak_entries(paks[0]):
        if b[:4] != b'IDXD':
            continue
        try:
            recs = U.parse(b)
        except Exception:
            continue
        for r in recs:
            if r['squadron'] == 'StageResource':
                stageres.append(dict((n, v) for _t, n, v in r['fields'] if n))
    named = sorted({d['EnumerateAIParams'] for d in stageres
                    if d.get('EnumerateAIParams')})
    keys = {h for h, _d, _p in objs}
    resolved = {}
    for n in named:
        for pre in ('', 'stage\\', 'table\\'):
            k = U.name_hash(pre + n)
            if k in keys:
                resolved[n] = (pre + n, k)
                break
    print("\n## Which object is which stage")
    print("   StageResource records: %d   carrying EnumerateAIParams: %d"
          % (len(stageres), sum(1 for d in stageres if d.get('EnumerateAIParams'))))
    print("   distinct table names declared: %d   AIParams objects: %d"
          % (len(named), len(objs)))
    print("   names that hash to an object key: %d / %d  (prefix `stage\\`,"
          " name_hash is case-insensitive)" % (len(resolved), len(named)))
    print("   objects left unnamed: %d   declared names with no object: %d"
          % (len(keys - {v[1] for v in resolved.values()}),
             len([n for n in named if n not in resolved])))
    share = collections.defaultdict(list)
    for d in stageres:
        v = d.get('EnumerateAIParams')
        if v:
            share[v].append(d.get('EnumerateSquadron', '?'))
    for v, who in sorted(share.items()):
        if len(who) > 1:
            print("   shared: %-24s by %d -> %s" % (v, len(who), sorted(who)))

    ct = collections.Counter()
    byname = collections.defaultdict(set)
    for _h, _d, p in objs:
        for s, f in p:
            ct[(f.get('Type'), len(f))] += 1
            byname[s].add((f.get('Type'), len(f)))
    print("\n## Does `Type` predict the field count?")
    for k, v in sorted(ct.items(), key=lambda kv: (str(kv[0][0]), kv[0][1])):
        print("   %-6s %2d fields  x%d" % (k[0], k[1], v))
    odd = sorted(s for s, v in byname.items()
                 if v & {('Squad', n) for n in range(30) if n != 20}
                 or v & {('Fleet', n) for n in range(30) if n != 6})
    print("   profiles breaking the rule, in full: %d  %s" % (len(odd), odd))
    var = sorted(s for s, v in byname.items() if len(v) > 1)
    print("   profiles whose (Type, field count) varies between objects: %d %s"
          % (len(var), var))

    fc = collections.Counter()
    for _h, _d, p in objs:
        for _s, f in p:
            for k in f:
                fc[k] += 1
    print("\n## Field census over every profile record")
    for k, v in sorted(fc.items(), key=lambda kv: (-kv[1], kv[0])):
        print("   x%-4d %s" % (v, k))

if __name__ == "__main__":
    main()