#!/opt/homebrew/bin/python3.11
"""
Select Face Reference Vectors V2 - Auto Multi-Angle Coverage

Purpose:
1. Analyze face.json and group faces by pose angle
2. Select top-K quality embeddings per angle
3. Ensure minimum angle coverage (frontal, three_quarter, profile_left, profile_right)
4. Generate angle_coverage_report
5. Register Identity with multi-angle reference_data

Features:
- Uses pose_analyzer V2 (multi-feature classification)
- Auto ensures >=4 angle coverage
- Quality-based selection (quality_score + pose confidence)
- Fallback strategy for missing angles

Usage:
    python3 scripts/select_face_reference_vectors_v2.py \
        --face-json output/video.face.json \
        --identity-name "Person Name" \
        --register

Output:
    reference_data JSONB structure:
    {
        "face_embeddings": [...],
        "angle_coverage": {
            "frontal": 2,
            "three_quarter": 3,
            "profile_left": 1,
            "profile_right": 1
        },
        "total_references": 7,
        "quality_avg": 0.89
    }
"""

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

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from utils.pose_analyzer import calculate_pose_angle_v2

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


def group_faces_by_angle(face_json_path: str) -> dict:
    """
    Group all faces in face.json by pose angle
    
    Returns:
        Dict with angle keys, each containing list of face data
    """
    with open(face_json_path) as f:
        data = json.load(f)
    
    frames = data.get("frames", {})
    
    angle_groups = {
        "frontal": [],
        "three_quarter": [],
        "profile_left": [],
        "profile_right": [],
        "unknown": [],
    }
    
    for frame_key, frame_data in frames.items():
        faces = frame_data.get("faces", [])
        
        for i, face in enumerate(faces):
            embedding = face.get("embedding")
            landmarks = face.get("landmarks")
            confidence = face.get("confidence", 0.9)
            attributes = face.get("attributes", {})
            
            if not embedding or len(embedding) != 512:
                continue
            
            if not landmarks or len(landmarks) < 5:
                continue
            
            pose_result = calculate_pose_angle_v2(landmarks)
            
            quality_score = calculate_quality_score(
                confidence=confidence,
                pose_confidence=pose_result["confidence"],
                embedding_norm=np.linalg.norm(embedding),
                attributes=attributes,
            )
            
            face_data = {
                "embedding": embedding,
                "frame": frame_key,
                "face_index": i,
                "pose_angle": pose_result["angle"],
                "pose_confidence": pose_result["confidence"],
                "pose_features": pose_result["features"],
                "pitch": pose_result.get("pitch", "neutral"),
                "quality_score": quality_score,
                "confidence": confidence,
                "attributes": attributes,
                "landmarks": landmarks[:5],
            }
            
            angle = pose_result["angle"]
            angle_groups[angle].append(face_data)
    
    return angle_groups


def calculate_quality_score(
    confidence: float,
    pose_confidence: float,
    embedding_norm: float,
    attributes: dict,
) -> float:
    """
    Calculate overall quality score (0.0 - 1.0)
    
    Components:
    - Detection confidence (40%)
    - Pose classification confidence (30%)
    - Embedding norm quality (20%)
    - Attributes completeness (10%)
    """
    norm_score = min(1.0, embedding_norm / 25.0) if embedding_norm > 0 else 0.5
    
    attr_score = 1.0 if attributes.get("age") and attributes.get("gender") else 0.5
    
    quality = (
        confidence * 0.40 +
        pose_confidence * 0.30 +
        norm_score * 0.20 +
        attr_score * 0.10
    )
    
    return min(1.0, max(0.0, quality))


def select_top_k_per_angle(
    angle_groups: dict,
    max_per_angle: int = 2,
    min_quality: float = 0.80,
) -> list:
    """
    Select top-K quality embeddings per angle
    
    Args:
        angle_groups: Dict of angle -> list of faces
        max_per_angle: Maximum per angle (default 2)
        min_quality: Minimum quality threshold
    
    Returns:
        List of selected face embeddings
    """
    selected = []
    
    priority_angles = ["frontal", "three_quarter", "profile_left", "profile_right"]
    
    for angle in priority_angles:
        faces = angle_groups.get(angle, [])
        
        if not faces:
            continue
        
        sorted_faces = sorted(faces, key=lambda x: x["quality_score"], reverse=True)
        
        for face in sorted_faces[:max_per_angle]:
            if face["quality_score"] >= min_quality:
                selected.append(face)
    
    return selected


def ensure_minimum_angle_coverage(
    selected: list,
    angle_groups: dict,
    min_angles: int = 4,
    min_per_angle: int = 1,
) -> list:
    """
    Ensure minimum angle coverage
    
    If missing angles, add best available from that group
    """
    covered_angles = set(f["pose_angle"] for f in selected)
    
    missing_angles = ["frontal", "three_quarter", "profile_left", "profile_right"]
    missing = [a for a in missing_angles if a not in covered_angles]
    
    for angle in missing:
        faces = angle_groups.get(angle, [])
        
        if faces:
            best_face = max(faces, key=lambda x: x["quality_score"])
            selected.append(best_face)
    
    return selected


def limit_total_references(selected: list, max_total: int = 10) -> list:
    """
    Limit total reference vectors to max_total
    
    Priority:
    1. Higher quality first
    2. Ensure angle diversity
    """
    if len(selected) <= max_total:
        return selected
    
    angle_priority = {"frontal": 0, "three_quarter": 1, "profile_left": 2, "profile_right": 3}
    
    def selection_priority(face):
        angle_weight = angle_priority.get(face["pose_angle"], 4)
        quality_weight = face["quality_score"]
        return (-angle_weight, -quality_weight)
    
    sorted_selected = sorted(selected, key=selection_priority)
    
    return sorted_selected[:max_total]


def generate_angle_coverage_report(selected: list) -> dict:
    """
    Generate angle coverage statistics
    """
    angle_counts = {}
    quality_by_angle = {}
    
    for face in selected:
        angle = face["pose_angle"]
        angle_counts[angle] = angle_counts.get(angle, 0) + 1
        
        if angle not in quality_by_angle:
            quality_by_angle[angle] = []
        quality_by_angle[angle].append(face["quality_score"])
    
    quality_avg = np.mean([f["quality_score"] for f in selected]) if selected else 0.0
    
    return {
        "angle_counts": angle_counts,
        "angles_covered": list(angle_counts.keys()),
        "angles_covered_count": len(angle_counts),
        "quality_avg": round(quality_avg, 4),
        "quality_by_angle": {
            angle: round(np.mean(scores), 4)
            for angle, scores in quality_by_angle.items()
        },
        "total_references": len(selected),
    }


def build_reference_data_structure(selected: list, video_uuid: str = None) -> dict:
    """
    Build reference_data JSONB structure for database
    """
    face_embeddings = []
    image_urls = []
    
    for face in selected:
        embedding_entry = {
            "embedding": face["embedding"],
            "angle": face["pose_angle"],
            "angle_confidence": face["pose_confidence"],
            "frame": face["frame"],
            "source": "video_detection",
            "quality_score": face["quality_score"],
            "detection_confidence": face["confidence"],
            "attributes": face["attributes"],
            "pose_features": face["pose_features"],
            "pitch": face.get("pitch", "neutral"),
            "created_at": datetime.now().isoformat(),
        }
        
        face_embeddings.append(embedding_entry)
    
    coverage_report = generate_angle_coverage_report(selected)
    
    reference_data = {
        "face_embeddings": face_embeddings,
        "angle_coverage": coverage_report["angle_counts"],
        "angles_covered": coverage_report["angles_covered"],
        "total_references": coverage_report["total_references"],
        "quality_avg": coverage_report["quality_avg"],
        "quality_by_angle": coverage_report["quality_by_angle"],
        "video_source": video_uuid or "unknown",
        "selection_method": "v2_auto_multi_angle",
        "created_at": datetime.now().isoformat(),
    }
    
    return reference_data


def register_identity_with_pose_v2(
    identity_name: str,
    reference_data: dict,
    schema: str = "dev",
) -> str:
    """
    Register Identity to database with V2 reference_data
    """
    conn = psycopg2.connect(DATABASE_URL)
    cur = conn.cursor()
    
    try:
        embeddings = [e["embedding"] for e in reference_data["face_embeddings"]]
        
        if embeddings:
            centroid = np.mean(np.array(embeddings), axis=0)
            centroid_norm = np.linalg.norm(centroid)
            if centroid_norm > 0:
                centroid_normalized = centroid / centroid_norm
            else:
                centroid_normalized = centroid
        else:
            centroid_normalized = np.zeros(512)
        
        embedding_str = "[" + ",".join(str(x) for x in centroid_normalized) + "]"
        
        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;
        """
        
        cur.execute(
            sql,
            (
                identity_name,
                "people",
                "video_detection_v2",
                "pending",
                embedding_str,
                json.dumps(reference_data),
            ),
        )
        
        uuid = cur.fetchone()[0]
        conn.commit()
        
        return uuid
        
    except Exception as e:
        print(f"❌ Database error: {e}")
        conn.rollback()
        return None
    finally:
        cur.close()
        conn.close()


def print_selection_report(
    angle_groups: dict,
    selected: list,
    coverage_report: dict,
):
    """
    Print detailed selection report
    """
    print("=" * 60)
    print("Auto Multi-Angle Reference Vector Selection (V2)")
    print("=" * 60)
    
    print("\n=== Available Faces by Angle ===")
    for angle, faces in angle_groups.items():
        if faces:
            quality_avg = np.mean([f["quality_score"] for f in faces])
            print(f"{angle}: {len(faces)} faces, quality_avg={quality_avg:.2f}")
    
    print("\n=== Selected Reference Vectors ===")
    print(f"Total: {len(selected)}")
    
    for i, face in enumerate(selected):
        print(f"\nVector {i+1}:")
        print(f"  Angle: {face['pose_angle']} (confidence: {face['pose_confidence']:.2f})")
        print(f"  Frame: {face['frame']}, Face: {face['face_index']}")
        print(f"  Quality: {face['quality_score']:.2f}")
        print(f"  Pitch: {face.get('pitch', 'neutral')}")
        print(f"  Age: {face['attributes'].get('age')}, Gender: {face['attributes'].get('gender')}")
    
    print("\n=== Angle Coverage Report ===")
    print(f"Angles covered: {coverage_report['angles_covered']}")
    print(f"Coverage count: {coverage_report['angles_covered_count']}")
    print(f"Angle distribution: {coverage_report['angle_counts']}")
    print(f"Quality avg: {coverage_report['quality_avg']:.2f}")
    print(f"Quality by angle: {coverage_report['quality_by_angle']}")


def main():
    parser = argparse.ArgumentParser(description="Auto Multi-Angle Reference Vector Selection V2")
    parser.add_argument("--face-json", required=True, help="Path to face.json")
    parser.add_argument("--identity-name", help="Identity name for registration")
    parser.add_argument("--max-per-angle", type=int, default=2, help="Max vectors per angle")
    parser.add_argument("--min-quality", type=float, default=0.80, help="Minimum quality threshold")
    parser.add_argument("--max-total", type=int, default=10, help="Maximum total vectors")
    parser.add_argument("--min-angles", type=int, default=4, help="Minimum angles to cover")
    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, don't register")
    args = parser.parse_args()
    
    print("🔧 Step 1: Grouping faces by angle...")
    angle_groups = group_faces_by_angle(args.face_json)
    
    print(f"\n🔧 Step 2: Selecting top-{args.max_per_angle} per angle...")
    selected = select_top_k_per_angle(
        angle_groups,
        max_per_angle=args.max_per_angle,
        min_quality=args.min_quality,
    )
    
    print(f"\n🔧 Step 3: Ensuring minimum {args.min_angles} angles coverage...")
    selected = ensure_minimum_angle_coverage(
        selected,
        angle_groups,
        min_angles=args.min_angles,
    )
    
    print(f"\n🔧 Step 4: Limiting to {args.max_total} total...")
    selected = limit_total_references(selected, max_total=args.max_total)
    
    coverage_report = generate_angle_coverage_report(selected)
    
    print_selection_report(angle_groups, selected, coverage_report)
    
    if not args.report_only and args.register and args.identity_name:
        print(f"\n🔧 Step 5: Registering Identity...")
        
        reference_data = build_reference_data_structure(selected, args.video_uuid)
        
        uuid = register_identity_with_pose_v2(
            identity_name=args.identity_name,
            reference_data=reference_data,
            schema=args.schema,
        )
        
        if uuid:
            print(f"\n✅ Registration completed!")
            print(f"   UUID: {uuid}")
            print(f"   Name: {args.identity_name}")
            print(f"   Angles: {coverage_report['angles_covered']}")
            print(f"   Total vectors: {coverage_report['total_references']}")
            print(f"   Quality avg: {coverage_report['quality_avg']:.2f}")
    elif args.report_only:
        print(f"\n📊 Report only (no registration)")


if __name__ == "__main__":
    main()