#!/opt/homebrew/bin/python3.11
"""
Pose-based Angle Filtering for Face Matching (V2)

Purpose:
1. Extract pose angle from landmarks using multi-feature analyzer
2. Filter reference vectors by pose angle
3. Only match against similar-angle reference vectors

Improvement:
- Uses pose_analyzer V2 (multi-feature classification)
- Higher confidence (0.87 avg vs 0.70)
- Better angle coverage

Usage:
    python3 scripts/match_face_with_pose_filtering.py --identity-name "Person" --face-json output/face.json --batch
"""

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

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 calculate_pose_angle_from_landmarks(landmarks):
    """
    Calculate pose angle from landmarks (V2 multi-feature)
    
    Uses pose_analyzer.calculate_pose_angle_v2
    
    Returns:
        dict with angle, confidence, features
    """
    return calculate_pose_angle_v2(landmarks)


def filter_reference_vectors_by_pose(detected_pose, reference_embeddings, tolerance=0.15):
    """
    Filter reference vectors by pose angle
    
    Args:
        detected_pose: dict with 'angle' and 'ratio'
        reference_embeddings: list of dicts with 'angle' and 'embedding'
        tolerance: ratio tolerance for filtering
    
    Returns:
        filtered list of reference embeddings
    """
    detected_angle = detected_pose.get("angle", "unknown")
    detected_ratio = detected_pose.get("ratio", 0.0)
    
    # Filter by angle type
    same_angle_refs = [
        ref for ref in reference_embeddings
        if ref.get("angle") == detected_angle
    ]
    
    # If no same-angle refs, use closest angles
    if not same_angle_refs:
        # Expand to include three_quarter (most common)
        if detected_angle in ["frontal", "profile_left", "profile_right"]:
            # Use three_quarter as fallback
            same_angle_refs = [
                ref for ref in reference_embeddings
                if ref.get("angle") == "three_quarter"
            ]
        
        # If still empty, use all
        if not same_angle_refs:
            same_angle_refs = reference_embeddings
    
    return same_angle_refs


def cosine_similarity(a, b):
    """Calculate cosine similarity"""
    a = np.array(a, dtype=np.float64)
    b = np.array(b, dtype=np.float64)
    
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    
    if norm_a == 0 or norm_b == 0:
        return 0.0
    
    return np.dot(a, b) / (norm_a * norm_b)


ANGLE_ADAPTIVE_THRESHOLDS = {
    "frontal": 0.90,
    "three_quarter": 0.85,
    "profile_left": 0.80,
    "profile_right": 0.80,
}

ANGLE_SIMILARITY_MATRIX = {
    "frontal": {"frontal": 1.0, "three_quarter": 0.8, "profile_left": 0.5, "profile_right": 0.5},
    "three_quarter": {"frontal": 0.8, "three_quarter": 1.0, "profile_left": 0.7, "profile_right": 0.7},
    "profile_left": {"frontal": 0.5, "three_quarter": 0.7, "profile_left": 1.0, "profile_right": 0.6},
    "profile_right": {"frontal": 0.5, "three_quarter": 0.7, "profile_left": 0.6, "profile_right": 1.0},
}

ANGLE_FALLBACK_ORDER = {
    "frontal": ["frontal", "three_quarter", "profile_left", "profile_right"],
    "three_quarter": ["three_quarter", "frontal", "profile_left", "profile_right"],
    "profile_left": ["profile_left", "three_quarter", "profile_right", "frontal"],
    "profile_right": ["profile_right", "three_quarter", "profile_left", "frontal"],
}


def get_adaptive_threshold(angle: str, base_threshold: float = 0.85) -> float:
    """Get angle-adaptive threshold"""
    adaptive = ANGLE_ADAPTIVE_THRESHOLDS.get(angle, base_threshold)
    return adaptive


def filter_reference_vectors_with_fallback(detected_pose, reference_embeddings):
    """
    Filter reference vectors with closest angle fallback
    
    Priority:
    1. Exact angle match
    2. Closest angle (by similarity matrix)
    3. All vectors (last resort)
    """
    detected_angle = detected_pose.get("angle", "unknown")
    
    if detected_angle == "unknown":
        return reference_embeddings
    
    exact_matches = [ref for ref in reference_embeddings if ref.get("angle") == detected_angle]
    
    if exact_matches:
        return exact_matches
    
    fallback_order = ANGLE_FALLBACK_ORDER.get(detected_angle, ["three_quarter"])
    
    for fallback_angle in fallback_order[1:]:
        matches = [ref for ref in reference_embeddings if ref.get("angle") == fallback_angle]
        if matches:
            return matches
    
    return reference_embeddings


def strategy_pose_filtered_v2(detected_embedding, detected_pose, reference_embeddings, base_threshold=0.85):
    """
    Strategy V2: Pose-filtered with Adaptive Threshold + Fallback
    
    Improvements:
    1. Angle-adaptive threshold
    2. Closest angle fallback
    3. Angle similarity weighting
    
    Returns:
        Dict with match result and detailed info
    """
    detected_angle = detected_pose.get("angle", "unknown")
    
    filtered_refs = filter_reference_vectors_with_fallback(detected_pose, reference_embeddings)
    
    if not filtered_refs:
        return {
            "strategy": "pose_filtered_v2",
            "filtered_count": 0,
            "best_similarity": 0.0,
            "is_match": False,
            "threshold": base_threshold,
            "pose_angle": detected_angle,
            "adaptive_threshold": base_threshold,
        }
    
    similarities = [
        cosine_similarity(detected_embedding, ref["embedding"])
        for ref in filtered_refs
    ]
    
    best_sim = max(similarities)
    best_idx = np.argmax(similarities)
    
    adaptive_threshold = get_adaptive_threshold(detected_angle, base_threshold)
    
    is_match = best_sim >= adaptive_threshold
    
    return {
        "strategy": "pose_filtered_v2",
        "filtered_count": len(filtered_refs),
        "total_references": len(reference_embeddings),
        "best_similarity": best_sim,
        "best_reference_idx": best_idx,
        "is_match": is_match,
        "threshold": adaptive_threshold,
        "base_threshold": base_threshold,
        "pose_angle": detected_angle,
        "pose_confidence": detected_pose.get("confidence"),
        "pose_features": detected_pose.get("features"),
        "best_reference_angle": filtered_refs[best_idx].get("angle"),
        "angle_match_type": "exact" if filtered_refs[best_idx].get("angle") == detected_angle else "fallback",
    }


def strategy_pose_filtered(detected_embedding, detected_pose, reference_embeddings, threshold=0.85):
    """
    Strategy V1: Pose-filtered Best Match (legacy)
    
    Steps:
    1. Filter reference vectors by pose angle
    2. Calculate best match among filtered vectors
    3. Return similarity
    """
    filtered_refs = filter_reference_vectors_by_pose(detected_pose, reference_embeddings)
    
    if not filtered_refs:
        return {
            "strategy": "pose_filtered",
            "filtered_count": 0,
            "best_similarity": 0.0,
            "is_match": False,
            "threshold": threshold,
            "pose_angle": detected_pose.get("angle"),
        }
    
    similarities = [
        cosine_similarity(detected_embedding, ref["embedding"])
        for ref in filtered_refs
    ]
    
    best_sim = max(similarities)
    best_idx = np.argmax(similarities)
    
    return {
        "strategy": "pose_filtered",
        "filtered_count": len(filtered_refs),
        "total_references": len(reference_embeddings),
        "best_similarity": best_sim,
        "best_reference_idx": best_idx,
        "is_match": best_sim >= threshold,
        "threshold": threshold,
        "pose_angle": detected_pose.get("angle"),
        "pose_features": detected_pose.get("features"),
        "best_reference_angle": filtered_refs[best_idx].get("angle"),
    }


def match_face_to_identity_with_pose(
    detected_embedding: list,
    detected_landmarks: list,
    identity_uuid: str,
    strategy: str = "pose_filtered_v2",
    threshold: float = 0.85,
    schema: str = "dev",
):
    """Match detected face to Identity with pose filtering
    
    Args:
        strategy: 'pose_filtered' (V1) or 'pose_filtered_v2' (V2 with adaptive threshold + fallback)
    """
    
    detected_pose = calculate_pose_angle_from_landmarks(detected_landmarks)
    
    detected_norm = np.linalg.norm(detected_embedding)
    if detected_norm > 0:
        detected_normalized = (np.array(detected_embedding) / detected_norm).tolist()
    else:
        detected_normalized = detected_embedding
    
    conn = psycopg2.connect(DATABASE_URL)
    cur = conn.cursor()
    
    try:
        # Get Identity reference_data
        cur.execute(f"""
            SELECT name, identity_type, reference_data
            FROM {schema}.identities
            WHERE uuid = %s;
        """, (identity_uuid,))
        
        result = cur.fetchone()
        
        if not result:
            return None
        
        name, identity_type, reference_data_json = result
        
        # Parse reference_data
        reference_data = json.loads(reference_data_json) if isinstance(reference_data_json, str) else reference_data_json
        
        face_embeddings = reference_data.get("face_embeddings", [])
        
        if not face_embeddings:
            return None
        
        # Choose strategy
        if strategy == "pose_filtered_v2":
            match_result = strategy_pose_filtered_v2(
                detected_embedding=detected_normalized,
                detected_pose=detected_pose,
                reference_embeddings=face_embeddings,
                base_threshold=threshold,
            )
        else:
            match_result = strategy_pose_filtered(
                detected_embedding=detected_normalized,
                detected_pose=detected_pose,
                reference_embeddings=face_embeddings,
                threshold=threshold,
            )
        
        match_result["identity_name"] = name
        match_result["identity_uuid"] = identity_uuid
        match_result["identity_type"] = identity_type
        match_result["reference_count"] = len(face_embeddings)
        
        return match_result
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return None
    finally:
        cur.close()
        conn.close()


def batch_match_faces_with_pose(face_json_path, identity_uuid, strategy="pose_filtered_v2", threshold=0.85, schema="dev"):
    """Batch match all faces with pose filtering
    
    Args:
        strategy: 'pose_filtered' (V1) or 'pose_filtered_v2' (V2)
    """
    
    with open(face_json_path) as f:
        data = json.load(f)
    
    frames = data.get("frames", {})
    
    results = []
    
    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")
            
            if not embedding:
                continue
            
            match_result = match_face_to_identity_with_pose(
                detected_embedding=embedding,
                detected_landmarks=landmarks,
                identity_uuid=identity_uuid,
                strategy=strategy,
                threshold=threshold,
                schema=schema,
            )
            
            if match_result:
                match_result["frame"] = frame_key
                match_result["face_index"] = i
                results.append(match_result)
    
    return results


def analyze_pose_match_results(results):
    """Analyze pose-filtered match results"""
    
    print("\n=== Pose-Filtered Match Results ===")
    print(f"Total faces: {len(results)}")
    
    # Match ratio
    is_match_count = sum(1 for r in results if r["is_match"])
    match_ratio = is_match_count / len(results) if results else 0
    
    print(f"Match ratio: {match_ratio:.2%} ({is_match_count}/{len(results)})")
    
    # Pose distribution
    pose_counts = {}
    for r in results:
        pose = r.get("pose_angle", "unknown")
        pose_counts[pose] = pose_counts.get(pose, 0) + 1
    
    print(f"Pose distribution: {pose_counts}")
    
    # Filtered count stats
    filtered_counts = [r["filtered_count"] for r in results]
    print(f"Filtered vectors: min={min(filtered_counts)}, max={max(filtered_counts)}, avg={np.mean(filtered_counts):.1f}")
    
    # Similarity by pose
    pose_sims = {}
    for r in results:
        pose = r.get("pose_angle", "unknown")
        if pose not in pose_sims:
            pose_sims[pose] = []
        pose_sims[pose].append(r["best_similarity"])
    
    print("\n=== Similarity by Pose ===")
    for pose, sims in pose_sims.items():
        print(f"{pose}: avg={np.mean(sims):.4f}, min={min(sims):.4f}, max={max(sims):.4f}")
    
    # V2 specific stats
    if results and results[0].get("strategy") == "pose_filtered_v2":
        adaptive_thresholds_used = {}
        angle_match_types = {}
        
        for r in results:
            threshold = r.get("adaptive_threshold")
            if threshold:
                angle = r.get("pose_angle", "unknown")
                adaptive_thresholds_used[angle] = threshold
            
            match_type = r.get("angle_match_type", "unknown")
            angle_match_types[match_type] = angle_match_types.get(match_type, 0) + 1
        
        print("\n=== V2 Adaptive Thresholds Used ===")
        for angle, threshold in adaptive_thresholds_used.items():
            print(f"{angle}: {threshold:.2f}")
        
        print("\n=== Angle Match Types ===")
        print(f"{angle_match_types}")
    
    # Top 5 details
    print("\n=== Top 5 Matches ===")
    sorted_results = sorted(results, key=lambda x: x["best_similarity"], reverse=True)
    
    for i, r in enumerate(sorted_results[:5]):
        pose_ratio = r.get("pose_ratio")
        pose_features = r.get("pose_features", {})
        
        ratio_str = f"{pose_ratio:.3f}" if pose_ratio else f"{pose_features.get('nose_to_eye_ratio', 'N/A')}"
        
        print(f"\nMatch {i+1}: Frame {r['frame']}, Face {r['face_index']}")
        print(f"  Pose Angle: {r['pose_angle']} (ratio: {ratio_str})")
        print(f"  Pose Confidence: {r.get('pose_confidence', 'N/A')}")
        print(f"  Filtered Vectors: {r['filtered_count']}/{r['total_references']}")
        print(f"  Best Similarity: {r['best_similarity']:.4f}")
        print(f"  Match: {r['is_match']} (threshold: {r['threshold']})")
        
        if r.get("strategy") == "pose_filtered_v2":
            print(f"  Adaptive Threshold: {r.get('adaptive_threshold', 'N/A')}")
            print(f"  Angle Match Type: {r.get('angle_match_type', 'N/A')}")


def main():
    parser = argparse.ArgumentParser(description="Pose-Filtered Face Matching")
    parser.add_argument("--identity-uuid", help="Identity UUID")
    parser.add_argument("--identity-name", help="Identity name")
    parser.add_argument("--face-json", required=True, help="Path to face.json")
    parser.add_argument("--strategy", default="pose_filtered_v2", choices=["pose_filtered", "pose_filtered_v2"], help="Matching strategy")
    parser.add_argument("--threshold", type=float, default=0.85)
    parser.add_argument("--schema", default="dev")
    parser.add_argument("--batch", action="store_true")
    args = parser.parse_args()
    
    print("=" * 60)
    print("Pose-Filtered Face Matching")
    print("=" * 60)
    
    # Get Identity UUID
    identity_uuid = args.identity_uuid
    
    if not identity_uuid and args.identity_name:
        conn = psycopg2.connect(DATABASE_URL)
        cur = conn.cursor()
        
        try:
            cur.execute(f"""
                SELECT uuid FROM {args.schema}.identities
                WHERE name = %s;
            """, (args.identity_name,))
            
            result = cur.fetchone()
            
            if result:
                identity_uuid = result[0]
                print(f"✅ Found: {args.identity_name} (UUID: {identity_uuid})")
            else:
                print(f"❌ Identity not found: {args.identity_name}")
                return
        finally:
            cur.close()
            conn.close()
    
    if not identity_uuid:
        print("❌ Please provide --identity-uuid or --identity-name")
        return
    
    print(f"\nStrategy: {args.strategy}")
    print(f"Threshold: {args.threshold}")
    
    if args.batch:
        print(f"\n🔧 Batch matching: {args.face_json}")
        results = batch_match_faces_with_pose(
            face_json_path=args.face_json,
            identity_uuid=identity_uuid,
            strategy=args.strategy,
            threshold=args.threshold,
            schema=args.schema,
        )
        
        analyze_pose_match_results(results)
    else:
        # Single match
        with open(args.face_json) as f:
            data = json.load(f)
        
        first_frame = list(data["frames"].values())[0]
        first_face = first_frame["faces"][0]
        
        match_result = match_face_to_identity_with_pose(
            detected_embedding=first_face.get("embedding"),
            detected_landmarks=first_face.get("landmarks"),
            identity_uuid=identity_uuid,
            strategy=args.strategy,
            threshold=args.threshold,
        )
        
        if match_result:
            pose_ratio = match_result.get("pose_ratio")
            pose_features = match_result.get("pose_features", {})
            ratio_str = f"{pose_ratio:.3f}" if pose_ratio else f"{pose_features.get('nose_to_eye_ratio', 'N/A')}"
            
            print("\n✅ Result:")
            print(f"  Pose: {match_result['pose_angle']} (ratio: {ratio_str})")
            print(f"  Similarity: {match_result['best_similarity']:.4f}")
            print(f"  Match: {match_result['is_match']}")
            
            if match_result.get("strategy") == "pose_filtered_v2":
                print(f"  Adaptive Threshold: {match_result.get('adaptive_threshold', 'N/A')}")
                print(f"  Angle Match Type: {match_result.get('angle_match_type', 'N/A')}")


if __name__ == "__main__":
    main()
    main()