from __future__ import annotations
import json, os, math, hashlib
from pathlib import Path
from datetime import datetime, timezone
import cv2
import numpy as np
import mediapipe as mp

LANDMARK_NAMES = [
    'nose','left_eye_inner','left_eye','left_eye_outer','right_eye_inner','right_eye','right_eye_outer',
    'left_ear','right_ear','mouth_left','mouth_right','left_shoulder','right_shoulder','left_elbow',
    'right_elbow','left_wrist','right_wrist','left_pinky','right_pinky','left_index','right_index',
    'left_thumb','right_thumb','left_hip','right_hip','left_knee','right_knee','left_ankle','right_ankle',
    'left_heel','right_heel','left_foot_index','right_foot_index'
]
POSE_CONNECTIONS = [
    (0,1),(1,2),(2,3),(3,7),(0,4),(4,5),(5,6),(6,8),(9,10),
    (11,12),(11,13),(13,15),(15,17),(15,19),(15,21),(17,19),
    (12,14),(14,16),(16,18),(16,20),(16,22),(18,20),
    (11,23),(12,24),(23,24),(23,25),(24,26),(25,27),(26,28),
    (27,29),(28,30),(29,31),(30,32),(27,31),(28,32)
]
CRITICAL = ['left_shoulder','right_shoulder','left_elbow','right_elbow','left_wrist','right_wrist','left_hip','right_hip','left_knee','right_knee','left_ankle','right_ankle','left_heel','right_heel','left_foot_index','right_foot_index']

def now_iso():
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace('+00:00','Z')

def lm_to_dict(lm, name):
    return {'name':name,'x':float(lm.x),'y':float(lm.y),'z':float(lm.z),'visibility':float(getattr(lm,'visibility',0.0)),'presence':float(getattr(lm,'presence',0.0))}

def midpoint(a,b,name):
    return {'name':name,'x':(a['x']+b['x'])/2,'y':(a['y']+b['y'])/2,'z':(a['z']+b['z'])/2,'visibility':min(a.get('visibility',0),b.get('visibility',0))}

def angle_deg(a,b,c):
    ax,ay=a['x'],a['y']; bx,by=b['x'],b['y']; cx,cy=c['x'],c['y']
    v1=(ax-bx,ay-by); v2=(cx-bx,cy-by)
    n1=math.hypot(*v1); n2=math.hypot(*v2)
    if n1<1e-9 or n2<1e-9: return None
    dot=max(-1,min(1,(v1[0]*v2[0]+v1[1]*v2[1])/(n1*n2)))
    return math.degrees(math.acos(dot))

def derived(lms):
    by={p['name']:p for p in lms}
    out={}
    for a,b,n in [('left_hip','right_hip','root_hip_midpoint'),('left_shoulder','right_shoulder','thorax_shoulder_midpoint'),('left_heel','left_foot_index','left_foot_contact_candidate'),('right_heel','right_foot_index','right_foot_contact_candidate')]:
        if a in by and b in by: out[n]=midpoint(by[a],by[b],n)
    for n in ['nose','left_shoulder','right_shoulder','left_elbow','right_elbow','left_wrist','right_wrist','left_hip','right_hip','left_knee','right_knee','left_ankle','right_ankle','left_heel','right_heel','left_foot_index','right_foot_index']:
        if n in by: out[f'{n}_candidate']=by[n]
    return out

def metrics(lms):
    by={p['name']:p for p in lms}
    vis=[p.get('visibility',0.0) for p in lms]
    crit=[by[k].get('visibility',0.0) for k in CRITICAL if k in by]
    low=[k for k in CRITICAL if k in by and by[k].get('visibility',0.0)<0.50]
    very=[k for k in CRITICAL if k in by and by[k].get('visibility',0.0)<0.12]
    edge=[]
    for k in CRITICAL:
        if k in by:
            x,y=by[k]['x'],by[k]['y']
            if x<0.02 or x>0.98 or y<0.02 or y>0.98: edge.append(k)
    m={'meanVisibility':sum(vis)/len(vis) if vis else 0.0,
       'minVisibility':min(vis) if vis else 0.0,
       'criticalMinVisibility':min(crit) if crit else 0.0,
       'criticalMeanVisibility':sum(crit)/len(crit) if crit else 0.0,
       'lowCriticalLandmarks':low,'veryLowCriticalLandmarks':very,'edgeCriticalLandmarks':edge}
    for side in ['left','right']:
        s,e,w=by.get(f'{side}_shoulder'),by.get(f'{side}_elbow'),by.get(f'{side}_wrist')
        h,k,a=by.get(f'{side}_hip'),by.get(f'{side}_knee'),by.get(f'{side}_ankle')
        if s and e and w: m[f'{side}_elbow_angle_deg']=angle_deg(s,e,w)
        if h and k and a: m[f'{side}_knee_angle_deg']=angle_deg(h,k,a)
    return m

def draw_annotation(bgr,lms,out):
    h,w=bgr.shape[:2]
    ann=bgr.copy()
    for a,b in POSE_CONNECTIONS:
        if a<len(lms) and b<len(lms) and min(lms[a].get('visibility',0),lms[b].get('visibility',0))>=0.12:
            pa=(int(lms[a]['x']*w),int(lms[a]['y']*h)); pb=(int(lms[b]['x']*w),int(lms[b]['y']*h))
            cv2.line(ann,pa,pb,(255,255,255),2)
    for lm in lms:
        if lm.get('visibility',0)>=0.12:
            p=(int(lm['x']*w),int(lm['y']*h)); cv2.circle(ann,p,3,(0,255,0),-1)
    out.parent.mkdir(parents=True,exist_ok=True)
    cv2.imwrite(str(out),ann)

def classify(m, key):
    flags=[]
    if m['criticalMinVisibility'] < 0.12: flags.append('very_low_critical_landmark_visibility')
    elif m['criticalMinVisibility'] < 0.30: flags.append('low_critical_landmark_visibility')
    if len(m['lowCriticalLandmarks']) >= 6: flags.append('many_low_critical_landmarks')
    if m['edgeCriticalLandmarks']: flags.append('cropped_or_edge_landmarks')
    # Sequence caveat but not automatic fail for landmarks; it blocks full exercise proof later.
    if key in {'back_pedal_to_sprint','rebound_side_toss'}: flags.append('sequence_timing_not_proven_by_still')
    if key == 'suitcase_carry': flags.append('carry_gait_not_proven_by_still')
    if not flags: return 'PASS_MEDIAPIPE_SOURCE_CANDIDATE', flags
    hard={'very_low_critical_landmark_visibility','many_low_critical_landmarks','cropped_or_edge_landmarks'}
    if any(f in hard for f in flags): return 'FIX_BEFORE_TEST_AUTHORITY', flags
    return 'REVIEW_BEFORE_RUNTIME_CONTRACT', flags

def main(root):
    root=Path(root)
    report_path=root/'tools/keypose_image_sampler/out/v88_stage2_phase3_mediapipe_extraction/V88_PREPROCESSED_SOURCE_FRAME_HARVEST_REPORT.json'
    data=json.loads(report_path.read_text())
    out_dir=root/'tools/keypose_image_sampler/out/v88_stage2_phase3_mediapipe_extraction/mediapipe_candidates'
    qa_dir=root/'qa/mesh_nex_v88_mediapipe'
    out_dir.mkdir(parents=True,exist_ok=True); qa_dir.mkdir(parents=True,exist_ok=True)
    model=root/'tools/keypose_image_sampler/models/pose_landmarker_full.task'
    BaseOptions=mp.tasks.BaseOptions
    PoseLandmarker=mp.tasks.vision.PoseLandmarker
    Options=mp.tasks.vision.PoseLandmarkerOptions
    RunningMode=mp.tasks.vision.RunningMode
    opts=Options(base_options=BaseOptions(model_asset_path=str(model)), running_mode=RunningMode.IMAGE, num_poses=1)
    records_by_ex={}
    all_rows=[]
    with PoseLandmarker.create_from_options(opts) as detector:
        for row in data['sourceFrameCandidates']:
            p=root/row['assetPath']
            bgr=cv2.imread(str(p), cv2.IMREAD_COLOR)
            if bgr is None:
                rec={**row,'hasPose':False,'error':'cv2_read_failed'}
            else:
                rgb=cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
                mp_img=mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
                res=detector.detect(mp_img)
                lms=[lm_to_dict(lm, LANDMARK_NAMES[i] if i < len(LANDMARK_NAMES) else f'landmark_{i}') for i,lm in enumerate(res.pose_landmarks[0])] if res.pose_landmarks else []
                world=[lm_to_dict(lm, LANDMARK_NAMES[i] if i < len(LANDMARK_NAMES) else f'landmark_{i}') for i,lm in enumerate(res.pose_world_landmarks[0])] if res.pose_world_landmarks else []
                m=metrics(lms) if lms else {'meanVisibility':0,'criticalMinVisibility':0,'criticalMeanVisibility':0,'lowCriticalLandmarks':CRITICAL,'veryLowCriticalLandmarks':CRITICAL,'edgeCriticalLandmarks':[]}
                status, flags = classify(m, row['exerciseKey']) if lms else ('NO_POSE',['mediapipe_no_pose'])
                ann=out_dir/row['exerciseKey']/('annotated/'+Path(row['assetPath']).stem+'.annotated.jpg')
                if lms: draw_annotation(bgr,lms,ann)
                rec={**row,'hasPose':bool(lms),'landmarks':lms,'worldLandmarks':world,'derivedNEXT11Candidates':derived(lms),'candidateMetrics':m,'mediaPipeGateStatus':status,'mediaPipeGateFlags':flags,'annotatedImage':str(ann.relative_to(root)) if lms else None}
            records_by_ex.setdefault(row['exerciseKey'],[]).append(rec)
            all_rows.append(rec)
    for key,recs in records_by_ex.items():
        pkg={'schema':'next11.v87.real_source_frame_mediapipe_candidates.v1','generatedAt':now_iso(),'exerciseKey':key,'approvalStatus':'candidate_only_not_nex_approved','records':recs}
        out=(out_dir/key/f'{key}.v87_mediapipe_keypose_candidates.json')
        out.parent.mkdir(parents=True,exist_ok=True)
        out.write_text(json.dumps(pkg,indent=2,ensure_ascii=False),encoding='utf-8')
    summary={
        'schema':'NEXT11_V88_MEDIAPIPE_PREPROCESSED_SOURCE_FRAME_RUN_V1','generatedAt':now_iso(),'mediapipeVersion':mp.__version__,
        'candidateFrames':len(all_rows),'framesWithPose':sum(1 for r in all_rows if r.get('hasPose')),
        'passCandidates':sum(1 for r in all_rows if r.get('mediaPipeGateStatus')=='PASS_MEDIAPIPE_SOURCE_CANDIDATE'),
        'reviewCandidates':sum(1 for r in all_rows if r.get('mediaPipeGateStatus')=='REVIEW_BEFORE_RUNTIME_CONTRACT'),
        'fixBeforeTestAuthority':sum(1 for r in all_rows if r.get('mediaPipeGateStatus')=='FIX_BEFORE_TEST_AUTHORITY'),
        'noPose':sum(1 for r in all_rows if r.get('mediaPipeGateStatus')=='NO_POSE'),
        'coveredExercises':sorted(records_by_ex),
        'unresolvedSourceGapExercises':data['summary'].get('unresolvedSourceGapExercises',[]),
        'policy':'MediaPipe extraction run on V87 real source frames. Candidate-only: no Nex approval and no replacement of V86 contracts until testplayer/mesh gate accepts.'
    }
    (out_dir/'V88_MEDIAPIPE_PREPROCESSED_SOURCE_FRAME_RUN_REPORT.json').write_text(json.dumps({'summary':summary,'rows':all_rows},indent=2,ensure_ascii=False),encoding='utf-8')
    # CSV
    import csv
    with open(out_dir/'V88_MEDIAPIPE_PREPROCESSED_SOURCE_FRAME_RUN_REPORT.csv','w',newline='',encoding='utf-8') as f:
        w=csv.DictWriter(f,fieldnames=['exerciseKey','phaseId','mediaPipeGateStatus','hasPose','criticalMinVisibility','criticalMeanVisibility','lowCriticalLandmarks','mediaPipeGateFlags','assetPath'])
        w.writeheader()
        for r in all_rows:
            m=r.get('candidateMetrics',{})
            w.writerow({'exerciseKey':r['exerciseKey'],'phaseId':r['phaseId'],'mediaPipeGateStatus':r.get('mediaPipeGateStatus'),'hasPose':r.get('hasPose'),'criticalMinVisibility':round(m.get('criticalMinVisibility',0),6),'criticalMeanVisibility':round(m.get('criticalMeanVisibility',0),6),'lowCriticalLandmarks':','.join(m.get('lowCriticalLandmarks',[])),'mediaPipeGateFlags':','.join(r.get('mediaPipeGateFlags',[])),'assetPath':r['assetPath']})
    print(json.dumps(summary,indent=2))
if __name__=='__main__':
    import sys
    main(sys.argv[1])