#!/usr/bin/env python3
"""Courtvision v2: preserve pose, possession, ball rotation/spin, and event feeds.

Python 3.10+. Standard library suffices for folders, JSON, JSONL, .output, and ZIP.
7z: pip install libarchive-c (with system libarchive), or pip install py7zr.
Usage:
  python pack_pose.py game.7z --out game.cvpose.gz
  python pack_pose.py game.output --out events.cvpose.gz
  python pack_pose.py extracted_folder --chunks-dir game --chunk-seconds 30
  python pack_pose.py mixed_games_folder --out-dir games
No source files are modified. Coordinates are stored as float32 feet.
"""
import argparse,array,collections,datetime,gzip,io,json,math,pathlib,struct,sys,tempfile,zipfile

H,P=28,91
JOINTS=['midHip','neck','nose','lEye','rEye','lEar','rEar','lShoulder','rShoulder','lElbow','rElbow','lWrist','rWrist','lHip','rHip','lKnee','rKnee','lAnkle','rAnkle','lHeel','rHeel','lBigToe','rBigToe','lSmallToe','rSmallToe','lThumb','rThumb','lPinky','rPinky']
TABLES=['Chances','Possessions','Shots','Rebounds','Turnovers','FreeThrows','Fouls','Timeouts','Violations','JumpBalls','Passes','Touches','Dribbles','Picks','Isolations','Handoffs','Drives','Posts','OffBallScreens','FakeHandoffs','Closeouts','BoxoutOpportunities','Boxouts','Falls','MissedCloseouts','ChancePlayers','Rotations','PossMatchups','SpeedLoad','ShotContests','PlayerStints','DefenderStance','Jumps']
NAN=float('nan')
def number(v):
    try:return float(v) if v is not None and v!='' else NAN
    except (ValueError,TypeError):return NAN
def identifier(v):
    if v is None:return ''
    if isinstance(v,dict):
        for k in ['nbaId','ngssId','ssiId','sessionId','id','trackingId','heId']:
            if v.get(k) is not None:return str(v[k])
        return ''
    return str(v)
def clock(v):
    if v is None or v=='':return NAN
    p=str(v).split(':')
    return number(p[0])*60+number(p[1]) if len(p)>1 else number(p[0])
def utc(v):
    if isinstance(v,(float,int)):return datetime.datetime.fromtimestamp(v/1000,datetime.timezone.utc)
    return datetime.datetime.fromisoformat(str(v).replace('Z','+00:00'))
def factor(unit):
    u=(unit or 'meters').lower()
    values={'inch':1/12,'inches':1/12,'in':1/12,'ft':1,'foot':1,'feet':1,'m':1/.3048,'meter':1/.3048,'meters':1/.3048,'metre':1/.3048,'metres':1/.3048,'cm':1/30.48,'centimeter':1/30.48,'centimeters':1/30.48,'mm':1/304.8,'millimeters':1/304.8}
    if u not in values:raise ValueError('Unsupported length unit: '+str(unit))
    return values[u]
def vector(v,f=1):
    if isinstance(v,dict):v=v.get('pos',v.get('position'))
    return [number(x)*f for x in v] if isinstance(v,(list,tuple)) and len(v)==3 else [NAN]*3
def valid3(v):return isinstance(v,(list,tuple)) and len(v)==3 and all(math.isfinite(number(x)) for x in v)
def boolean(v):return 1 if v is True or v==1 else 0 if v is False or v==0 else NAN
def unwrap(d):
    for _ in range(5):
        if not isinstance(d,dict):break
        if 'payload' in d:d=json.loads(d['payload']) if isinstance(d['payload'],str) else d['payload']
        elif isinstance(d.get('data'),dict) and ('message' in d['data'] or 'samples' in d['data']):d=d['data']
        else:break
    return d
def game_key(d):
    g=d.get('details',{}).get('game',{}).get('id',{})
    if isinstance(g,dict) and g.get('nbaId') is not None:return 'nba:'+str(g['nbaId']).zfill(10)
    if isinstance(g,dict) and g.get('sessionId'):return 'session:'+g['sessionId']
    if g:return 'game:'+identifier(g)
    return 'game:'+identifier(d['gameId']) if d.get('gameId') else 'unassigned'
def person_joints(p):
    j=p.get('joints',p.get('pose',{}))
    return (j[0] if j else {}) if isinstance(j,list) else j or {}
def centroid(p):
    c=p.get('centroid',p.get('centerOfMass',p.get('position')))
    if isinstance(c,list) and c and isinstance(c[0],dict):c=c[0].get('pos')
    return c or person_joints(p).get('midHip')
class Game:
    def __init__(self,d):
        dt=d.get('details',{});teams=[{'id':identifier(t.get('id')),'name':t.get('name',''),'code':t.get('tricode',t.get('name','TEAM')),'home':t.get('home') is True,'attackingPositiveX':t.get('attackingPositiveX'),'aliases':t.get('id',{})} for t in dt.get('teams',[])]
        teams.sort(key=lambda t:not t['home']);comp=dt.get('competition',{}).get('name','Basketball');ncaa='ncaa' in comp.lower() or 'college' in comp.lower();fiba=any(x in comp.lower() for x in ['fiba','euro','olympic'])
        self.meta={'version':2,'gameKey':game_key(d),'gameId':dt.get('game',{}).get('id',{}),'title':' · '.join(t['code'] for t in teams) or 'Basketball data','venue':dt.get('venue',{}).get('name',''),'competition':comp,'date':str(d.get('time',{}).get('timeUTC',''))[:10],'originUTC':d.get('time',{}).get('timeUTC'),'units':'feet','joints':JOINTS,'frameHeader':H,'personStride':P,'players':[],'teams':teams,'court':{'length':28/.3048 if fiba else 94,'width':15/.3048 if fiba else 50,'rules':'NCAA' if ncaa else 'FIBA' if fiba else 'NBA','rimHeight':10},'sourceMessages':[],'sourceUnits':[],'events':[],'attackHistory':[],'importWarnings':[]}
        self.rows=array.array('f');self.offsets=array.array('I',[0]);self.slots={};self.events=[];self.event_keys=set();self.pose_count=0;self.snapshot_count=0
    def relative(self,v):
        if not self.meta['originUTC']:
            self.meta['originUTC']=str(v);self.meta['date']=str(v)[:10]
        return (utc(v)-utc(self.meta['originUTC'])).total_seconds()
    def register(self,p):
        id=identifier(p.get('id',p.get('playerId',p.get('officialId',p.get('trackingId')))))
        if not id:return NAN
        existing=self.slots.get(id)
        data={'id':id,'name':p.get('fullName') or p.get('name') or ' '.join(filter(None,[p.get('firstName'),p.get('lastName')])) or 'Player '+id,'shortName':p.get('shortName',''),'jersey':str(p.get('jerseyNumber',p.get('jersey',''))),'teamId':identifier(p.get('teamId')),'kind':p.get('kind','player'),'aliases':p.get('aliases',p.get('id',{}))}
        if existing is None:existing=len(self.meta['players']);self.slots[id]=existing;self.meta['players'].append(data)
        else:
            old=self.meta['players'][existing]
            old.update({k:v for k,v in data.items() if v!='' and not (k=='name' and v=='Player '+id)})
        return existing
    def slot(self,id):
        key=identifier(id)
        if not key:return NAN
        return self.slots[key] if key in self.slots else self.register({'id':key})
    def team_slot(self,id):
        return next((i for i,t in enumerate(self.meta['teams']) if t['id']==identifier(id)),NAN)
    def init(self,d):
        if not self.meta['originUTC'] and d.get('time',{}).get('timeUTC'):self.meta['originUTC']=d['time']['timeUTC']
        for p in d.get('details',{}).get('players',[]):self.register(dict(p,kind='player'))
        for p in d.get('details',{}).get('referees',[]):self.register(dict(p,kind='official'))
        for key,value in [('sourceMessages',d.get('message','tracking')),('sourceUnits',d.get('units',{}).get('length','meters'))]:
            if value not in self.meta[key]:self.meta[key].append(value)
        for t in d.get('details',{}).get('teams',[]):
            if not isinstance(t.get('attackingPositiveX'),bool):continue
            id=identifier(t.get('id'));direction=1 if t['attackingPositiveX'] else -1;prior=next((x for x in reversed(self.meta['attackHistory']) if x['teamId']==id),None)
            if prior is None or prior['direction']!=direction:self.meta['attackHistory'].append({'teamId':id,'time':self.relative(d['time']['timeUTC']),'direction':direction})
    def pose(self,d,snapshot=False,event=None,stamp=None):
        self.init(d);t=d.get('time',{});stamp=stamp or t.get('timeUTC')
        if not stamp:return
        f=factor(d.get('units',{}).get('length'));samples=d.get('samples',{});balls=samples.get('ball',[]);ball=(balls[0] if balls else {}) if isinstance(balls,list) else balls;summary=d.get('summary',{});pos=summary.get('possession',{});people=[p for p in samples.get('people',[]) if identifier(p.get('playerId',p.get('id',p.get('officialId',p.get('trackingId')))))]
        r=[NAN]*H;r[0]=self.relative(stamp);r[1]=number(d.get('sequences',{}).get('frame',(event or {}).get('summary',{}).get('nearestFrameNumber')));r[2]=number(t.get('period'));r[3]=clock(t.get('gameClockTime'));r[4]=clock(t.get('shotClockTime'));r[5]=number(d.get('score',{}).get('home'));r[6]=number(d.get('score',{}).get('away'));r[7:10]=vector(ball.get('pos') if isinstance(ball,dict) else ball,f);r[10]=len(people)
        r[11]=self.slot(pos.get('playerId'));r[12]=self.team_slot(pos.get('teamId'));r[13]=1 if pos.get('attackingPositiveX') is True else -1 if pos.get('attackingPositiveX') is False else NAN;r[14]=boolean(summary.get('gameActive'))
        units=d.get('units',{});af=math.pi/180 if 'degree' in str(units.get('angle','radians')).lower() else 1;sf=math.pi/180 if 'degree' in str(units.get('angularVelocity','radiansPerSecond')).lower() else 1
        r[15:18]=vector(ball.get('orientation'),af);r[18:21]=vector(ball.get('spinRate'),sf);r[21]=boolean(t.get('gameClockRunning'));r[22]=boolean(t.get('shotClockRunning'));r[23]=2 if snapshot else 1;r[24]=number(d.get('revision',0));r[25]=-1 if snapshot else 3 if 'processed' in d.get('message','') else 2 if 'delayed' in d.get('message','') else 1;r[26]=number(pos.get('count'));r[27]=number(ball.get('confidence'))
        for p in people:
            r.extend([self.slot(p.get('playerId',p.get('id',p.get('officialId',p.get('trackingId'))))),*vector(centroid(p),f)])
            j=person_joints(p)
            for name in JOINTS:r.extend(vector(j.get(name),f))
        self.rows.extend(r);self.offsets.append(len(self.rows));self.snapshot_count+=bool(snapshot);self.pose_count+=not snapshot
    def event_message(self,d):
        self.init(d)
        for i,e in enumerate(d.get('events',[])):
            h=e.get('heInformation',{});ng=e.get('ngssInformation',{});s=e.get('summary',{});time=d.get('time',{});stamp=time.get('timeUTC');pose_stamp=s.get('nearestFrameTimeUTC',stamp);people=e.get('people',[]);type=ng.get('type') or h.get('classification') or h.get('type') or 'event';sub=ng.get('subType') or h.get('playerAction') or '';event_id=identifier(e.get('id',{}).get('heId') or e.get('id',{}).get('ngssId')) or str(d.get('sequences',{}).get('event',''))+':'+str(i)
            ids=[identifier(p.get('playerId')) for p in people if identifier(p.get('playerId'))]
            key=':'.join([e.get('source','hawkeye'),event_id,type,sub,h.get('playerAction',''),','.join(ids),stamp,str(s.get('nearestFrameNumber','')),s.get('nearestFrameTimeUTC','')])
            if key in self.event_keys:continue
            self.event_keys.add(key)
            for p in people:
                slot=self.slot(p.get('playerId'))
                if math.isfinite(slot) and p.get('jerseyNumber') is not None:self.meta['players'][slot]['jersey']=p['jerseyNumber']
            pos=s.get('possession',{})
            self.events.append({'id':key,'type':type,'subtype':sub,'source':e.get('source','hawkeye'),'utc':stamp,'time':self.relative(stamp),'snapshotTime':self.relative(pose_stamp),'period':number(time.get('period')),'clock':clock(time.get('gameClockTime')),'shot':clock(time.get('shotClockTime')),'playerIds':ids,'teamId':identifier(pos.get('teamId')),'possessionId':identifier(pos.get('uuid')),'possessionCount':pos.get('count'),'success':ng.get('success'),'location':vector(e.get('ballLocation'),factor(d.get('units',{}).get('length'))),'hasGeometry':valid3(e.get('ballLocation')) or any(valid3(centroid(p)) for p in people),'sourceFrame':s.get('nearestFrameNumber'),'raw':e})
            q=dict(d);q['sequences']={'frame':s.get('nearestFrameNumber')};q['summary']=dict(s,gameActive=h.get('gameActive'));q['samples']={'ball':[{'pos':e['ballLocation']}] if e.get('ballLocation') else [],'people':[p for p in people if centroid(p) or person_joints(p)]}
            self.pose(q,True,e,pose_stamp)
    def markings(self,d,table,records):
        for i,r in enumerate(records):
            stamp=r.get('startUtcTime') or r.get('utcTime') or r.get('startTimeUTC') or r.get('timeUTC')
            if not stamp and r.get('startWallClock') is not None:stamp=utc(r['startWallClock']).isoformat().replace('+00:00','Z')
            ids=[identifier(r[k]) for k in ['playerId','shooterId','ballhandlerId','passerId','receiverId','rebounderId','defenderId','offPlayerId','defPlayerId','screenerId','foulerId','fouledId'] if r.get(k) is not None]
            for id in ids:self.slot(id)
            type={'Passes':'pass','Chances':'chance','Possessions':'possession','FreeThrows':'freethrow','BoxoutOpportunities':'boxout','DefenderStance':'stance'}.get(table,table[:-1].lower() if table.endswith('s') else table.lower());id='marking:'+table+':'+identifier(r.get('id',table+':'+str(i)))
            if id in self.event_keys:continue
            self.event_keys.add(id);self.events.append({'id':id,'type':type,'subtype':r.get('reason') or r.get('shotType') or r.get('passType') or r.get('fallType') or r.get('category') or '', 'source':'markings','utc':stamp,'time':self.relative(stamp) if stamp else None,'period':number(r.get('period')),'clock':clock(r.get('startGameClock',r.get('gameClock'))),'shot':clock(r.get('startShotClock',r.get('shotClock'))),'playerIds':ids,'teamId':identifier(r.get('teamId',r.get('offTeamId'))),'success':r.get('made',r.get('success')),'location':vector(r.get('location',r.get('startLoc',r.get('shotLoc',r.get('ballStartLoc')))),factor(d.get('units',{}).get('length','feet'))),'hasGeometry':False,'sourceFrame':r.get('startFrame',r.get('frameIdx')),'frameBasis':'secondspectrum','raw':r,'table':table})
        if 'secondspectrum.markings' not in self.meta['sourceMessages']:self.meta['sourceMessages'].append('secondspectrum.markings')
    def finish(self):
        self.meta['originUTC']=self.meta['originUTC'] or '1970-01-01T00:00:00.000Z';self.meta['events']=self.events;self.meta['importStats']={'poseFrames':self.pose_count,'eventFrames':self.snapshot_count}
def clean_json(v):
    if isinstance(v,float) and not math.isfinite(v):return None
    if isinstance(v,dict):return {k:clean_json(x) for k,x in v.items()}
    if isinstance(v,list):return [clean_json(x) for x in v]
    return v
def write_pack(path,game,start=None,end=None,shuffle=True):
    meta=dict(game.meta);meta['encoding']='float32-byte-shuffle' if shuffle else 'float32'
    if start is not None:
        idx=[i for i in range(len(game.offsets)-1) if start<=game.rows[game.offsets[i]]<end];rows=array.array('f');offsets=array.array('I',[0])
        for i in idx:rows.extend(game.rows[game.offsets[i]:game.offsets[i+1]]);offsets.append(len(rows))
        meta['events']=[e for e in meta['events'] if e.get('time') is not None and start<=e['time']<end]
    else:rows=game.rows;offsets=game.offsets
    head=json.dumps(clean_json(meta),separators=(',',':'),ensure_ascii=False,allow_nan=False).encode();head+=b' '*((-len(head))%4)
    data=rows.tobytes();index=offsets.tobytes()
    if sys.byteorder!='little':
        rows=array.array('f',rows);offsets=array.array('I',offsets);rows.byteswap();offsets.byteswap();data=rows.tobytes();index=offsets.tobytes()
    path.parent.mkdir(parents=True,exist_ok=True);op=gzip.open if path.suffix=='.gz' else open
    with op(path,'wb') as target:
        target.write(b'CVPOSE2\n');target.write(struct.pack('<III',len(head),len(offsets)-1,len(rows)));target.write(head);target.write(index)
        if shuffle:
            for plane in range(4):target.write(data[plane::4])
        else:target.write(data)
    return len(offsets)-1
def iter_sources(path):
    if path.is_dir():
        for f in sorted(path.rglob('*')):
            if f.is_file() and f.suffix.lower() in ('.json','.jsonl','.ndjson','.output','.clean','.complete','.pose','.tracking','.txt','.gz'):yield from iter_sources(f)
    elif path.suffix.lower()=='.7z':
        try:import libarchive
        except (ImportError,OSError):
            try:import py7zr
            except ImportError:raise SystemExit('Install libarchive-c with system libarchive, or py7zr for the portable fallback.')
            print('Extracting archive to a temporary folder...',flush=True)
            with tempfile.TemporaryDirectory(prefix='courtvision-') as temp:
                with py7zr.SevenZipFile(path,'r') as archive:archive.extractall(path=temp)
                yield from iter_sources(pathlib.Path(temp))
        else:
            with libarchive.file_reader(str(path)) as archive:
                for entry in archive:
                    if not entry.isfile:continue
                    with tempfile.SpooledTemporaryFile(max_size=4*1024*1024) as f:
                        for block in entry.get_blocks():f.write(block)
                        f.seek(0);yield entry.pathname,f
    elif path.suffix.lower()=='.zip':
        with zipfile.ZipFile(path) as archive:
            for name in sorted(archive.namelist()):
                if name.endswith('/'):continue
                with archive.open(name) as f:yield name,f
    else:
        with (gzip.open(path,'rb') if path.suffix=='.gz' else open(path,'rb')) as f:yield path.name.removesuffix('.gz'),f
def main():
    parser=argparse.ArgumentParser(description=__doc__,formatter_class=argparse.RawDescriptionHelpFormatter);parser.add_argument('input',type=pathlib.Path);group=parser.add_mutually_exclusive_group(required=True);group.add_argument('--out',type=pathlib.Path);group.add_argument('--out-dir',type=pathlib.Path);group.add_argument('--chunks-dir',type=pathlib.Path);parser.add_argument('--chunk-seconds',type=float,default=30);parser.add_argument('--no-shuffle',action='store_true');args=parser.parse_args()
    if args.chunk_seconds<=0:parser.error('--chunk-seconds must be positive')
    if args.out and args.out.resolve()==args.input.resolve():parser.error('Input and output must be different files.')
    games={};warnings=[];count=0;skipped=0
    def get(d,key=None):
        key=key or game_key(d)
        if key not in games:games[key]=Game(d)
        return games[key]
    def consume(raw):
        nonlocal count,skipped
        d=unwrap(raw)
        if isinstance(d,list):
            for x in d:consume(x)
            return
        if not isinstance(d,dict):skipped+=1;return
        if 'frames' in d and 'samples' not in d:consume(d['frames']);return
        message=d.get('message','');stamp=d.get('time',{}).get('timeUTC')
        if stamp and ('samples' in d or any(x in message for x in ('pose','centroid','tracking'))):get(d).pose(d);count+=1
        elif stamp and isinstance(d.get('events'),list):get(d).event_message(d);count+=1
        else:
            matched=False
            for container in [d,d.get('data'),d.get('markings')]:
                if not isinstance(container,dict):continue
                for name,value in container.items():
                    table=next((t for t in TABLES if ''.join(filter(str.isalnum,t.lower()))==''.join(filter(str.isalnum,name.lower()))),None)
                    if table and isinstance(value,list):
                        metadata=dict(d);metadata['gameId']=d.get('gameId') or (value[0].get('gameId') if value else None);get(metadata).markings(d,table,value);matched=True;count+=len(value)
            if not matched:skipped+=1
        if count and count%6000==0:print(f'{count:,} records read',flush=True)
    for name,f in iter_sources(args.input):
        if not name.lower().endswith(('.json','.jsonl','.ndjson','.output','.clean','.complete','.pose','.tracking','.txt')):continue
        if name.lower().endswith(('.jsonl','.ndjson','.output')):
            for line_no,line in enumerate(f,1):
                if not line.strip():continue
                try:consume(json.loads(line.decode('utf-8-sig')))
                except (ValueError,TypeError,KeyError) as e:
                    skipped+=1
                    if len(warnings)<20:warnings.append(f'{name}:{line_no}: {e}')
        else:
            try:consume(json.load(io.TextIOWrapper(f,encoding='utf-8-sig')))
            except (ValueError,TypeError,KeyError) as e:
                skipped+=1
                if len(warnings)<20:warnings.append(f'{name}: {e}')
    if not games:raise SystemExit('No supported tracking, event, or marking records found.')
    if args.out and len(games)!=1:raise SystemExit('Multiple games found. Use --out-dir to keep game identities separate.')
    for key,game in games.items():
        game.finish();game.meta['importWarnings']=warnings;game.meta['importStats']['skipped']=skipped;slug=''.join(c if c.isalnum() or c in '-_' else '_' for c in key)
        if args.out:
            frames=write_pack(args.out,game,shuffle=not args.no_shuffle);print(f'Saved {frames:,} samples / {len(game.events):,} events: {args.out}')
        elif args.out_dir:
            path=args.out_dir/(slug+'.cvpose.gz');frames=write_pack(path,game,shuffle=not args.no_shuffle);print(f'Saved {frames:,} samples: {path}')
        else:
            destination=args.chunks_dir/slug if len(games)>1 else args.chunks_dir;times=[game.rows[o] for o in game.offsets[:-1]];start=min(times,default=0);last=max(times,default=0);chunks=[]
            for i in range(max(1,math.floor((last-start)/args.chunk_seconds)+1)):
                a=start+i*args.chunk_seconds;b=a+args.chunk_seconds;path=destination/(f'clip-{i:04d}.cvpose.gz');frames=write_pack(path,game,a,b,not args.no_shuffle);chunks.append({'url':path.name,'start':a,'end':min(last,b),'frames':frames})
            (destination/'manifest.json').write_text(json.dumps({'gameKey':key,'chunks':chunks},indent=2));print(f'Saved {len(chunks)} indexed segments: {destination}')
    print(f'Read {count:,} supported records across {len(games)} game(s); {skipped:,} skipped records.')
if __name__=='__main__':main()