#!/usr/bin/env python3
"""按已登记的分类，从固定版本的团队预览池分层抽样。

python tools/import_team_previews.py --per 12
来源与版本见 data/preview_sources.json；完整素材池见 data/preview_pool.json。
主分类在数据集内部轮转，分类内部均衡来源已提供的手别和视角。
"""
from pathlib import Path
from collections import Counter
import argparse, hashlib, json

SITE=Path(__file__).resolve().parent.parent
ap=argparse.ArgumentParser(description=__doc__)
ap.add_argument('--per',type=int,default=12)
quota=ap.parse_args().per
data=json.loads((SITE/'data/preview_pool.json').read_text())
selected={}
audit={'target_per_dataset':quota,'seed':'20261027','scope':'已取得的公开预览池','strategy':f'各数据集最多 {quota} 组；按主分类轮转，分类内部再均衡已提供的手别、视角等分类。素材不足时全部展示。','dimensions_note':'分类来自来源子集、序列、相机和手别元数据。姿态与抓握类别尚未建立统一标签。','datasets':{}}
for key,items in data.items():
 dimensions=['来源子集','序列','视角','手别','受试者']
 primary=next((dim for dim in dimensions if len({p['strata'].get(dim) for p in items if p['strata'].get(dim)})>1),None)
 chosen=[];remaining=items.copy();counts={dim:Counter() for dim in dimensions};conditional=Counter();primary_counts=Counter()
 while len(chosen)<min(quota,len(items)):
  def score(p):
   group=p['strata'].get(primary,'未提供分类')
   return (primary_counts[group],sum(conditional[(group,dim,p['strata'][dim])] for dim in dimensions if dim in p['strata'] and dim!=primary),sum(counts[dim][p['strata'][dim]] for dim in dimensions if dim in p['strata'] and dim!=primary),hashlib.sha256(('20261027/'+p['id']).encode()).hexdigest())
  p=min(remaining,key=score);remaining.remove(p);chosen.append(p)
  group=p['strata'].get(primary,'未提供分类');primary_counts[group]+=1
  for dim,value in p['strata'].items(): counts[dim][value]+=1;conditional[(group,dim,value)]+=1
 selected[key]=chosen
 audit['datasets'][key]={'pool':len(items),'selected':len(chosen),'primary_dimension':primary or '可用素材全部展示','categories':{dim:{v:{'pool':sum(p['strata'].get(dim)==v for p in items),'selected':sum(p['strata'].get(dim)==v for p in chosen)} for v in sorted({p['strata'][dim] for p in items if dim in p['strata']})} for dim in dimensions if any(dim in p['strata'] for p in items)},'missing_dimensions':[dim for dim in ['受试者','动作类别','抓握类别','三维姿态类别'] if not any(dim in p['strata'] for p in items)]}
 if primary and primary!='手别' and all('手别' in p['strata'] for p in items):
  joint=lambda p:p['strata'].get(primary,'未提供')+' / '+p['strata']['手别']
  audit['datasets'][key]['categories'][primary+' × 手别']={v:{'pool':sum(joint(p)==v for p in items),'selected':sum(joint(p)==v for p in chosen)} for v in sorted({joint(p) for p in items})}
(SITE/'data/previews.json').write_text(json.dumps(selected,ensure_ascii=False,indent=1)+'\n')
audit['videos']=json.loads((SITE/'data/sampling.json').read_text())['videos']
(SITE/'data/sampling.json').write_text(json.dumps(audit,ensure_ascii=False,indent=1)+'\n')
print('POOL',sum(map(len,data.values())),'SELECTED',sum(map(len,selected.values())))
for key,v in audit['datasets'].items(): print(key,v['primary_dimension'],v['selected'],v['pool'])