"""Enrich frozen causal features, without rerunning spatial association or Pose."""
import bootstrap
from bootstrap import BASELINE
import hashlib
import json
import math
import os
from pathlib import Path
import time
from inputs import OUT, load
from motion import region_distance


def joints(body):
    return {int(k): v for k, v in body['joints'].items()}


def midpoint(j, a, b):
    return [(j[a][k] + j[b][k]) / 2 for k in (0, 1)] if a in j and b in j else None


def enrich(f, previous, current):
    a = f['actor']; other = 1 - a
    p, q = current[a], current[other]
    old = previous[a]
    j, k, oldj = joints(p), joints(q), joints(old)
    h, oh, dt = max(1, p['height']), max(1, q['height']), f['elapsed']
    shoulder = midpoint(j, 5, 6)
    hip = midpoint(j, 11, 12)
    target_shoulder = midpoint(k, 5, 6)
    end = j[f['end']]
    hand_distances = [(math.dist(end[:2], k[i][:2]) / oh, i) for i in (9, 10) if i in k]
    hand = min(hand_distances) if hand_distances else None
    target_hip = midpoint(k, 11, 12)
    root_delta = [(j[f['root']][i] - oldj[f['root']][i]) / h / dt for i in (0, 1)]
    old_shoulder = midpoint(oldj, 5, 6)
    old_hip = midpoint(oldj, 11, 12)
    angle = None
    if shoulder and hip:
        angle = abs(shoulder[0]-hip[0]) / max(1, math.dist(shoulder, hip))
    target_upright = None
    if target_shoulder and target_hip:
        target_upright = abs(target_shoulder[0]-target_hip[0]) / max(1, math.dist(target_shoulder,target_hip))
    return dict(f, shoulder_offset=None if shoulder is None else (end[1]-shoulder[1])/h,
        horizontal_reach=abs(end[0]-j[f['root']][0])/h,
        hand_contact=None if hand is None else hand[0],
        target_hand_offset=None if hand is None or target_shoulder is None else (k[hand[1]][1]-target_shoulder[1])/oh,
        root_vx=root_delta[0],root_vy=root_delta[1],
        shoulder_vy=None if shoulder is None or old_shoulder is None else (shoulder[1]-old_shoulder[1])/h/dt,
        hip_vy=None if hip is None or old_hip is None else (hip[1]-old_hip[1])/h/dt,
        torso_horizontal=angle,target_torso_horizontal=target_upright,
        target_height_ratio=oh/h,
        target_center_distance=math.dist(p['center'],q['center'])/max(1,(h+oh)/2))


def main():
    clip=os.environ['CLIP']; dest=OUT/'work'/clip;dest.mkdir(parents=True,exist_ok=True)
    raw=(BASELINE/'work'/clip/'features.jsonl').read_bytes()
    check=json.loads((BASELINE/'work'/clip/'features_summary.json').read_text())
    digest=hashlib.sha256(raw).hexdigest();assert digest==check['sha256']
    rows,chosen,spec,cache,maps=load(clip)
    views=[json.loads(line) for line in raw.splitlines()]
    assert [(r['sample'],r['source_frame'],r['seconds']) for r in rows]==[(r['sample'],r['source_frame'],r['seconds']) for r in views]
    done=dest/'features_summary.json'
    if done.exists():
        d=json.loads(done.read_text());assert d['baseline_sha256']==digest
        assert hashlib.sha256((dest/'features.jsonl').read_bytes()).hexdigest()==d['sha256']
        print('REUSED',clip);return
    previous={};features=0;start=time.monotonic()
    with (dest/'features.jsonl').open('w') as out:
        for row in views:
            t=row['seconds']
            for p in row['pairs']:
                for f in p['features']:
                    assert f['input_end']==t and f['input_start']<t
                    old=previous[(p['token'],f['input_start'])]
                    f.update(enrich(f,old,p['bodies']));features+=1
                previous[(p['token'],t)]=p['bodies']
            previous={k:v for k,v in previous.items() if t-k[1]<=.35}
            out.write(json.dumps(row,separators=(',',':'))+'\n')
    done.write_text(json.dumps(dict(status='PASS',baseline_sha256=digest,sha256=hashlib.sha256((dest/'features.jsonl').read_bytes()).hexdigest(),
        displayed=len(chosen),processed=len(rows),feature_count=features,cache_sha256=spec['cache_sha256'],causal=True,seconds=time.monotonic()-start),indent=2))
    print('AUGMENT_PASS',clip,len(chosen),features,flush=True)


if __name__=='__main__':main()