#!/opt/homebrew/bin/python3.11
"""
Face Reference Vector Selector

Purpose:
1. Analyze face.json results
2. Select high-quality embeddings as reference vectors
3. Group by angle/quality/confidence
4. Output reference_data JSONB structure

Usage:
    python3 scripts/select_face_reference_vectors.py --face-json output/preview.face_new.json --identity-name "Test Person"
"""

import json
import argparse
import numpy as np
from datetime import datetime
import psycopg2
import os

DATABASE_URL = os.getenv("DATABASE_URL", "postgres://accusys@localhost:5432/momentry?options=-c%20search_path=dev")


def classify_face_angle(bbox, img_width, img_height):
    """Classify face angle based on bbox position"""
    face_center_x = (bbox[0] + bbox[2]) / 2
    face_width = bbox[2] - bbox[0]
    
    # Relative position in image
    pos_ratio = face_center_x / img_width
    
    # Face size ratio
    size_ratio = face_width / img_width
    
    # Angle classification
    if size_ratio > 0.3:  # Large face, likely frontal
        return "frontal"
    elif pos_ratio < 0.3:  # Left side
        return "profile_left"
    elif pos_ratio > 0.7:  # Right side
        return "profile_right"
    elif size_ratio > 0.2:
        return "three_quarter"
    else:
        return "unknown"


def calculate_embedding_quality(face, embedding):
    """Calculate embedding quality score (0.0-1.0)"""
    
    # Confidence score (0.0-1.0)
    confidence = face.get("confidence", 0.9)
    
    # Embedding norm (well-formed embeddings have consistent norms)
    norm = np.linalg.norm(embedding)
    norm_score = 1.0 if norm > 20 else norm / 20
    
    # Attributes completeness
    attrs = face.get("attributes", {})
    attr_score = 1.0 if attrs.get("age") and attrs.get("gender") else 0.5
    
    # Combined quality
    quality = confidence * 0.5 + norm_score * 0.3 + attr_score * 0.2
    
    return min(1.0, max(0.0, quality))


def select_reference_vectors(face_json_path, max_vectors=10, min_quality=0.7):
    """Select high-quality reference vectors from face.json"""
    
    with open(face_json_path) as f:
        data = json.load(f)
    
    metadata = data.get("metadata", {})
    frames = data.get("frames", {})
    
    img_width = metadata.get("width", 640)
    img_height = metadata.get("height", 360)
    
    candidates = []
    
    for frame_key, frame_data in frames.items():
        faces = frame_data.get("faces", [])
        
        for i, face in enumerate(faces):
            embedding = face.get("embedding")
            
            if not embedding or len(embedding) != 512:
                continue
            
            # Calculate quality
            quality = calculate_embedding_quality(face, embedding)
            
            # Classify angle
            bbox = [face["x"], face["y"], face["x"] + face["width"], face["y"] + face["height"]]
            angle = classify_face_angle(bbox, img_width, img_height)
            
            candidates.append({
                "frame": frame_key,
                "face_index": i,
                "embedding": embedding,
                "quality_score": quality,
                "confidence": face.get("confidence", 0.9),
                "angle": angle,
                "attributes": face.get("attributes", {}),
                "bbox": bbox,
                "timestamp": frame_data.get("time_seconds", 0),
            })
    
    # Sort by quality
    candidates.sort(key=lambda x: x["quality_score"], reverse=True)
    
    # Select top candidates, ensuring angle diversity
    selected = []
    angles_used = set()
    
    for candidate in candidates:
        if candidate["quality_score"] < min_quality:
            continue
        
        # Ensure angle diversity
        angle = candidate["angle"]
        
        if angle not in angles_used or len(selected) < 5:
            selected.append(candidate)
            angles_used.add(angle)
        
        if len(selected) >= max_vectors:
            break
    
    return selected, metadata


def register_face_identity(
    identity_name: str,
    reference_vectors: list,
    metadata: dict,
    schema: str = "dev",
    video_uuid: str = None,
):
    """Register Face Identity with 1-to-many reference vectors"""
    
    conn = psycopg2.connect(DATABASE_URL)
    cur = conn.cursor()
    
    try:
        # Prepare reference_data
        face_embeddings = [
            {
                "embedding": rv["embedding"],
                "source": "video_detection",
                "frame": rv["frame"],
                "angle": rv["angle"],
                "quality_score": rv["quality_score"],
                "confidence": rv["confidence"],
                "attributes": rv["attributes"],
                "created_at": datetime.now().isoformat(),
            }
            for rv in reference_vectors
        ]
        
        reference_data = {
            "face_embeddings": face_embeddings,
            "video_source": video_uuid or "unknown",
        }
        
        # Calculate centroid
        embeddings_array = np.array([rv["embedding"] for rv in reference_vectors])
        centroid = np.mean(embeddings_array, axis=0).tolist()
        
        # Normalize centroid
        centroid_norm = np.linalg.norm(centroid)
        if centroid_norm > 0:
            centroid_normalized = (np.array(centroid) / centroid_norm).tolist()
        else:
            centroid_normalized = centroid
        
        # Insert or update Identity
        sql = f"""
        INSERT INTO {schema}.identities (
            name, identity_type, source, status,
            face_embedding, reference_data,
            created_at, updated_at
        ) VALUES (
            %s, %s, %s, %s,
            %s, %s,
            NOW(), NOW()
        )
        ON CONFLICT (name) DO UPDATE SET
            face_embedding = EXCLUDED.face_embedding,
            reference_data = EXCLUDED.reference_data,
            updated_at = NOW()
        RETURNING uuid;
        """
        
        embedding_str = "[" + ",".join(str(x) for x in centroid_normalized) + "]"
        
        cur.execute(
            sql,
            (
                identity_name,
                "people",
                "video_detection",
                "pending",
                embedding_str,
                json.dumps(reference_data),
            ),
        )
        
        uuid = cur.fetchone()[0]
        conn.commit()
        
        print(f"✅ Identity registered: {identity_name}")
        print(f"   UUID: {uuid}")
        print(f"   Reference vectors: {len(reference_vectors)}")
        print(f"   Angles covered: {set(rv['angle'] for rv in reference_vectors)}")
        
        return uuid
        
    except Exception as e:
        print(f"❌ Database error: {e}")
        conn.rollback()
        return None
    finally:
        cur.close()
        conn.close()


def analyze_reference_vectors(reference_vectors):
    """Analyze reference vectors quality and diversity"""
    
    print("\n=== Reference Vectors Analysis ===")
    print(f"Total vectors: {len(reference_vectors)}")
    
    # Quality distribution
    qualities = [rv["quality_score"] for rv in reference_vectors]
    print(f"Quality scores: min={min(qualities):.2f}, max={max(qualities):.2f}, avg={np.mean(qualities):.2f}")
    
    # Angle distribution
    angles = [rv["angle"] for rv in reference_vectors]
    angle_counts = {}
    for angle in angles:
        angle_counts[angle] = angle_counts.get(angle, 0) + 1
    print(f"Angle distribution: {angle_counts}")
    
    # Similarity analysis
    embeddings = np.array([rv["embedding"] for rv in reference_vectors])
    norms = np.linalg.norm(embeddings, axis=1)
    
    # Calculate pairwise similarities
    similarities = []
    for i in range(len(embeddings)):
        for j in range(i+1, len(embeddings)):
            sim = np.dot(embeddings[i], embeddings[j]) / (norms[i] * norms[j])
            similarities.append(sim)
    
    if similarities:
        print(f"Inter-vector similarity: min={min(similarities):.2f}, max={max(similarities):.2f}, avg={np.mean(similarities):.2f}")
    
    # Print detailed vector info
    print("\n=== Selected Vectors ===")
    for i, rv in enumerate(reference_vectors):
        print(f"Vector {i+1}:")
        print(f"  Frame: {rv['frame']}, Angle: {rv['angle']}")
        print(f"  Quality: {rv['quality_score']:.2f}, Confidence: {rv['confidence']:.2f}")
        print(f"  Age: {rv['attributes'].get('age')}, Gender: {rv['attributes'].get('gender')}")


def main():
    parser = argparse.ArgumentParser(description="Select Face Reference Vectors")
    parser.add_argument("--face-json", required=True, help="Path to face.json")
    parser.add_argument("--identity-name", required=True, help="Identity name for registration")
    parser.add_argument("--max-vectors", type=int, default=10, help="Max reference vectors")
    parser.add_argument("--min-quality", type=float, default=0.7, help="Minimum quality score")
    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("--analyze-only", action="store_true", help="Only analyze, don't register")
    args = parser.parse_args()
    
    print("=" * 60)
    print("Face Reference Vector Selector")
    print("=" * 60)
    
    # Select reference vectors
    print(f"\n🔧 Analyzing: {args.face_json}")
    reference_vectors, metadata = select_reference_vectors(
        args.face_json,
        max_vectors=args.max_vectors,
        min_quality=args.min_quality,
    )
    
    if not reference_vectors:
        print("❌ No high-quality reference vectors found")
        return
    
    print(f"✅ Selected {len(reference_vectors)} reference vectors")
    
    # Analyze
    analyze_reference_vectors(reference_vectors)
    
    # Register (if requested)
    if not args.analyze_only and args.register:
        print(f"\n🔧 Registering Identity: {args.identity_name}")
        uuid = register_face_identity(
            identity_name=args.identity_name,
            reference_vectors=reference_vectors,
            metadata=metadata,
            schema=args.schema,
            video_uuid=args.video_uuid,
        )
        
        if uuid:
            print("\n🎉 Registration completed!")
    else:
        print("\n📊 Analysis only (no registration)")
        print("   To register, run with --register flag")


if __name__ == "__main__":
    main()