#!/opt/homebrew/bin/python3.11
"""
Select Face Reference Vectors V3 - With Trace Support

Purpose:
1. Select reference vectors from specific trace_id (same person)
2. Ensure multi-angle coverage within trace
3. Register identity with trace statistics

Improvements over V2:
- trace_id_filter: Select vectors only from specific trace
- trace_quality: Use trace avg_confidence as quality indicator
- trace_stats: Include trace duration, pose distribution in registration

Usage:
    # Select from longest trace
    python3 select_face_reference_vectors_v3.py \
      --face-json video.face_traced.json \
      --identity-name "Person Name" \
      --use-longest-trace \
      --register

    # Select from specific trace
    python3 select_face_reference_vectors_v3.py \
      --face-json video.face_traced.json \
      --trace-id-filter 2 \
      --identity-name "Person Name" \
      --register
"""

import sys
import os
import json
import argparse
import numpy as np
from typing import Dict, List, Optional
from collections import defaultdict
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

try:
    import psycopg2
    from psycopg2.extras import execute_values
except ImportError:
    print("Warning: psycopg2 not installed, registration will be skipped")
    psycopg2 = None


def get_longest_trace(face_data: Dict) -> Optional[int]:
    """
    Get the longest trace_id (most appearances)
    
    Returns:
        trace_id of longest trace, or None if no traces
    """
    traces = face_data.get("traces", {})
    
    if not traces:
        return None
    
    longest_trace_id = None
    max_appearances = 0
    
    for trace_id_str, trace in traces.items():
        appearances = trace.get("total_appearances", 0)
        avg_confidence = trace.get("avg_confidence", 0)
        
        # Prefer higher confidence if same appearances
        if appearances > max_appearances or \
           (appearances == max_appearances and avg_confidence > face_data.get("traces", {}).get(str(longest_trace_id), {}).get("avg_confidence", 0)):
            max_appearances = appearances
            longest_trace_id = int(trace_id_str)
    
    return longest_trace_id


def filter_faces_by_trace(face_data: Dict, trace_id_filter: int) -> Dict:
    """
    Filter faces to only include specific trace_id
    
    Args:
        face_data: face_traced.json data
        trace_id_filter: trace_id to filter
    
    Returns:
        Filtered face data with only specified trace_id faces
    """
    frames = face_data.get("frames", {})
    
    filtered_frames = {}
    
    for frame_num_str, frame_data in frames.items():
        filtered_faces = []
        
        for face in frame_data.get("faces", []):
            if face.get("trace_id") == trace_id_filter:
                filtered_faces.append(face)
        
        if filtered_faces:
            filtered_frames[frame_num_str] = {
                "frame_number": frame_data["frame_number"],
                "time_seconds": frame_data["time_seconds"],
                "faces": filtered_faces,
            }
    
    filtered_data = {
        "metadata": face_data.get("metadata", {}),
        "frames": filtered_frames,
    }
    
    return filtered_data


def select_reference_vectors_v2(face_data: Dict, max_per_angle: int = 2, min_quality: float = 0.6) -> List[Dict]:
    """
    Select reference vectors using V2 algorithm (multi-angle)
    
    This is the same as select_face_reference_vectors_v2.py
    """
    frames = face_data.get("frames", {})
    
    angles_data = defaultdict(list)
    
    for frame_num_str, frame_data in frames.items():
        for face_idx, face in enumerate(frame_data.get("faces", [])):
            pose_angle = face.get("pose_angle", {})
            angle = pose_angle.get("angle", "unknown")
            
            if angle == "unknown":
                continue
            
            confidence = pose_angle.get("confidence", 0.0)
            det_confidence = face.get("confidence", 0.0)
            quality_score = (confidence + det_confidence) / 2
            
            if quality_score < min_quality:
                continue
            
            embedding = face.get("embedding")
            if not embedding:
                continue
            
            landmarks = face.get("landmarks")
            
            angles_data[angle].append({
                "frame": int(frame_num_str),
                "face_index": face_idx,
                "embedding": embedding,
                "landmarks": landmarks,
                "pose_angle": angle,
                "pose_confidence": confidence,
                "quality_score": quality_score,
                "det_confidence": det_confidence,
                "attributes": face.get("attributes", {}),
            })
    
    selected_vectors = []
    
    for angle in angles_data:
        sorted_faces = sorted(angles_data[angle], key=lambda x: x["quality_score"], reverse=True)
        selected_for_angle = sorted_faces[:max_per_angle]
        selected_vectors.extend(selected_for_angle)
    
    return selected_vectors


def register_identity_with_trace(
    identity_name: str,
    selected_vectors: List[Dict],
    trace_id: Optional[int],
    trace_stats: Optional[Dict],
    schema: str = "dev",
    video_uuid: Optional[str] = None,
) -> Optional[str]:
    """
    Register identity with reference vectors and trace statistics
    """
    if psycopg2 is None:
        print("psycopg2 not installed, skipping registration")
        return None
    
    if not selected_vectors:
        print("No vectors selected, skipping registration")
        return None
    
    embeddings = [v["embedding"] for v in selected_vectors]
    qualities = [v["quality_score"] for v in selected_vectors]
    angles = [v["pose_angle"] for v in selected_vectors]
    
    angles_covered = list(set(angles))
    quality_avg = np.mean(qualities)
    
    quality_by_angle = {}
    for angle in set(angles):
        angle_qualities = [v["quality_score"] for v in selected_vectors if v["pose_angle"] == angle]
        quality_by_angle[angle] = np.mean(angle_qualities) if angle_qualities else 0.0
    
    angle_distribution = defaultdict(int)
    for angle in angles:
        angle_distribution[angle] += 1
    
    conn = psycopg2.connect(
        host="localhost",
        database="momentry",
        user="accusys",
    )
    cur = conn.cursor()
    
    try:
        # Check if identity exists
        cur.execute(f"SELECT uuid FROM {schema}.identities WHERE name = %s", (identity_name,))
        existing = cur.fetchone()
        
        if existing:
            identity_uuid = existing[0]
            print(f"⚠️ Identity '{identity_name}' already exists (UUID: {identity_uuid})")
            print("Updating reference_data...")
        else:
            # Create new identity
            cur.execute(
                f"""
                INSERT INTO {schema}.identities (name, identity_type, source, status)
                VALUES (%s, 'people', 'auto_trace', 'active')
                RETURNING uuid
                """,
                (identity_name,),
            )
            identity_uuid = str(cur.fetchone()[0])
            print(f"✅ Created identity: {identity_name} (UUID: {identity_uuid})")
        
        # Build reference_data (compatible with V2 format)
        face_embeddings = []
        for v in selected_vectors:
            embedding_entry = {
                "embedding": v["embedding"],
                "angle": v["pose_angle"],
                "angle_confidence": v["pose_confidence"],
                "frame": v["frame"],
                "source": "trace_detection",
                "quality_score": v["quality_score"],
                "detection_confidence": v["det_confidence"],
                "attributes": v.get("attributes", {}),
                "created_at": datetime.now().isoformat(),
            }
            face_embeddings.append(embedding_entry)
        
        reference_data = {
            "face_embeddings": face_embeddings,
            "angle_coverage": dict(angle_distribution),
            "angles_covered": angles_covered,
            "total_references": len(selected_vectors),
            "quality_avg": round(quality_avg, 4),
            "quality_by_angle": {k: round(v, 4) for k, v in quality_by_angle.items()},
            "selection_method": "trace_filtered_v3",
            "created_at": datetime.now().isoformat(),
        }
        
        # Add trace statistics
        if trace_id is not None and trace_stats:
            reference_data["trace_id"] = trace_id
            pose_counts = defaultdict(int)
            for p in trace_stats.get("pose_angles", []):
                pose_counts[p] += 1
            reference_data["trace_stats"] = {
                "start_frame": trace_stats.get("start_frame"),
                "end_frame": trace_stats.get("end_frame"),
                "duration_frames": trace_stats.get("duration_frames"),
                "duration_seconds": round(trace_stats.get("duration_seconds", 0), 2),
                "total_appearances": trace_stats.get("total_appearances"),
                "avg_confidence": round(trace_stats.get("avg_confidence", 0), 4),
                "pose_distribution": dict(pose_counts),
            }
        
        if video_uuid:
            reference_data["video_source"] = video_uuid
        
        # Update reference_data
        cur.execute(
            f"""
            UPDATE {schema}.identities
            SET reference_data = %s,
                updated_at = CURRENT_TIMESTAMP
            WHERE uuid = %s
            """,
            (json.dumps(reference_data), identity_uuid),
        )
        
        conn.commit()
        
        return identity_uuid
    
    except Exception as e:
        conn.rollback()
        print(f"❌ Registration failed: {e}")
        return None
    
    finally:
        cur.close()
        conn.close()


def main():
    parser = argparse.ArgumentParser(description="Select Face Reference Vectors V3 (With Trace)")
    parser.add_argument("--face-json", required=True, help="Path to face_traced.json")
    parser.add_argument("--identity-name", help="Identity name for registration")
    parser.add_argument("--trace-id-filter", type=int, help="Filter by specific trace_id")
    parser.add_argument("--use-longest-trace", action="store_true", help="Use longest trace automatically")
    parser.add_argument("--max-per-angle", type=int, default=2, help="Max vectors per angle")
    parser.add_argument("--min-quality", type=float, default=0.6, help="Minimum quality threshold")
    parser.add_argument("--schema", default="dev", help="Database schema")
    parser.add_argument("--video-uuid", help="Video UUID")
    parser.add_argument("--register", action="store_true", help="Register to database")
    parser.add_argument("--report-only", action="store_true", help="Only show report")
    args = parser.parse_args()
    
    print("=" * 60)
    print("Auto Multi-Angle Reference Vector Selection V3")
    print("(With Trace Support)")
    print("=" * 60)
    
    with open(args.face_json) as f:
        face_data = json.load(f)
    
    traces = face_data.get("traces", {})
    
    if not traces:
        print("❌ No traces found in face.json")
        print("Please run face_tracker.py first")
        return
    
    print(f"\n=== Available Traces ===")
    for trace_id_str, trace in sorted(traces.items(), key=lambda x: int(x[0])):
        print(f"Trace {trace_id_str}:")
        print(f"  Frames: {trace['start_frame']}-{trace['end_frame']} ({trace['duration_frames']} frames)")
        print(f"  Duration: {trace['duration_seconds']:.2f}s")
        print(f"  Appearances: {trace['total_appearances']}")
        print(f"  Avg Confidence: {trace['avg_confidence']:.3f}")
        print(f"  Pose Angles: {set(trace['pose_angles'])}")
        print()
    
    # Determine trace_id to use
    trace_id_filter = args.trace_id_filter
    
    if args.use_longest_trace and trace_id_filter is None:
        trace_id_filter = get_longest_trace(face_data)
        print(f"🎯 Using longest trace: Trace {trace_id_filter}")
    
    if trace_id_filter is None:
        print("❌ Please specify --trace-id-filter or --use-longest-trace")
        return
    
    if str(trace_id_filter) not in traces:
        print(f"❌ Trace {trace_id_filter} not found")
        return
    
    trace_stats = traces[str(trace_id_filter)]
    
    print(f"\n=== Selected Trace {trace_id_filter} ===")
    print(f"Duration: {trace_stats['duration_seconds']:.2f}s ({trace_stats['duration_frames']} frames)")
    print(f"Confidence: {trace_stats['avg_confidence']:.3f}")
    print(f"Pose Distribution: {defaultdict(int, {p: trace_stats['pose_angles'].count(p) for p in set(trace_stats['pose_angles'])})}")
    
    # Filter faces by trace
    filtered_face_data = filter_faces_by_trace(face_data, trace_id_filter)
    
    print(f"\n=== Filtering Faces ===")
    print(f"Original frames: {len(face_data.get('frames', {}))}")
    print(f"Filtered frames: {len(filtered_face_data.get('frames', {}))}")
    
    # Select reference vectors
    selected_vectors = select_reference_vectors_v2(
        filtered_face_data,
        max_per_angle=args.max_per_angle,
        min_quality=args.min_quality,
    )
    
    if not selected_vectors:
        print("❌ No reference vectors selected")
        return
    
    print(f"\n=== Selected Reference Vectors ===")
    print(f"Total: {len(selected_vectors)}")
    
    angle_distribution = defaultdict(int)
    for v in selected_vectors:
        angle_distribution[v["pose_angle"]] += 1
    
    print(f"Angles: {list(set([v['pose_angle'] for v in selected_vectors]))}")
    print(f"Distribution: {dict(angle_distribution)}")
    print(f"Quality avg: {np.mean([v['quality_score'] for v in selected_vectors]):.3f}")
    
    print(f"\n=== Vector Details ===")
    for i, v in enumerate(selected_vectors[:10]):
        print(f"Vector {i+1}:")
        print(f"  Angle: {v['pose_angle']} (confidence: {v['pose_confidence']:.2f})")
        print(f"  Frame: {v['frame']}, Face: {v['face_index']}")
        print(f"  Quality: {v['quality_score']:.3f}")
        attrs = v.get("attributes", {})
        if attrs:
            print(f"  Age: {attrs.get('age', 'N/A')}, Gender: {attrs.get('gender', 'N/A')}")
    
    if args.report_only:
        return
    
    if args.register and args.identity_name:
        print(f"\n=== Registering Identity ===")
        
        identity_uuid = register_identity_with_trace(
            identity_name=args.identity_name,
            selected_vectors=selected_vectors,
            trace_id=trace_id_filter,
            trace_stats=trace_stats,
            schema=args.schema,
            video_uuid=args.video_uuid,
        )
        
        if identity_uuid:
            print(f"\n✅ Registration completed!")
            print(f"   UUID: {identity_uuid}")
            print(f"   Name: {args.identity_name}")
            print(f"   Trace ID: {trace_id_filter}")
            print(f"   Total vectors: {len(selected_vectors)}")
            print(f"   Quality avg: {np.mean([v['quality_score'] for v in selected_vectors]):.3f}")


if __name__ == "__main__":
    main()