import cv2
import numpy as np
from insightface.app import FaceAnalysis
from scipy.spatial.distance import cosine
import warnings
warnings.filterwarnings('ignore')

class AIFRKTPVerification:
    def __init__(self):
        """
        Initialize Age-Invariant Face Recognition system specifically for KTP verification
        """
        print("🚀 Initializing AIFR KTP Verification System...")
        # Initialize InsightFace with ArcFace model (best for age-invariant features)
        self.app = FaceAnalysis(name='buffalo_l', 
                                root='/tmp',
                                providers=['CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))
        # Thresholds specifically calibrated for ID-to-selfie comparison
        self.thresholds = {
            'strict': 0.35,
            'normal': 0.45,
            'relaxed': 0.55,
            'very_relaxed': 0.65
        }

    def preprocess_ktp_image(self, image):
        """
        Special preprocessing for Indonesian KTP photos
        - Remove red/orange tint
        - Enhance contrast
        - Normalize lighting
        """
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        a = cv2.normalize(a, None, 0, 255, cv2.NORM_MINMAX)
        b = cv2.normalize(b, None, 0, 255, cv2.NORM_MINMAX)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        corrected = cv2.merge([l, a, b])
        corrected = cv2.cvtColor(corrected, cv2.COLOR_LAB2BGR)
        corrected = cv2.bilateralFilter(corrected, 9, 75, 75)
        return corrected

    def extract_age_invariant_features(self, image):
        """
        Extracts both deep embedding and geometric features from the face.
        """
        faces = self.app.get(image)
        if len(faces) == 0:
            print("⚠️  No face detected!")
            return None, None
            
        face = faces[0]
        embedding = face.normed_embedding # 512-dimensional embedding
        landmarks = face.kps # Keypoints for geometric features
        
        geometric_features = None
        if landmarks is not None and len(landmarks) >= 5:
            try:
                # Eye distance ratio (doesn't change much with age)
                eye_distance = np.linalg.norm(landmarks[0] - landmarks[1])
                # Nose to eye center distance
                eye_center = (landmarks[0] + landmarks[1]) / 2
                nose_distance = np.linalg.norm(landmarks[2] - eye_center)
                # Face width estimation (using jawline points if available)
                face_width = np.linalg.norm(landmarks[3] - landmarks[4])
                
                # Create a stable geometric feature vector
                if face_width > 0 and nose_distance > 0:
                    geometric_features = np.array([
                        eye_distance / face_width,
                        nose_distance / face_width,
                        eye_distance / nose_distance
                    ])
            except Exception as e:
                print(f"Could not calculate geometric features: {e}")
                geometric_features = None

        return embedding, geometric_features

    def calculate_similarity(self, embed1, embed2, geom1=None, geom2=None):
        """
        Calculates similarity using a weighted average of deep and geometric features.
        """
        # Deep feature similarity (main metric)
        deep_similarity = 1 - cosine(embed1, embed2)
        
        # Geometric similarity (supplementary)
        if geom1 is not None and geom2 is not None:
            geom_similarity = 1 - cosine(geom1, geom2)
            # Weighted combination (80% deep, 20% geometric)
            final_similarity = 0.8 * deep_similarity + 0.2 * geom_similarity
        else:
            final_similarity = deep_similarity
            
        return final_similarity, deep_similarity

    def verify_with_images(self, ktp_image, selfie_image):
        """
        Main verification function that accepts image objects instead of file paths.
        This is the method the API will call.
        """
        print("\n🔧 Preprocessing images for API call...")
        ktp_processed = self.preprocess_ktp_image(ktp_image)
        selfie_processed = cv2.normalize(selfie_image, None, 0, 255, cv2.NORM_MINMAX)

        print("🧠 Extracting age-invariant features...")
        ktp_embedding, ktp_geom = self.extract_age_invariant_features(ktp_processed)
        selfie_embedding, selfie_geom = self.extract_age_invariant_features(selfie_processed)

        if ktp_embedding is None or selfie_embedding is None:
            return {'error': 'Could not detect a face in one or both images.'}

        print("📊 Calculating similarity...")
        similarity, deep_similarity = self.calculate_similarity(
            ktp_embedding, selfie_embedding, ktp_geom, selfie_geom
        )

        normal_threshold = self.thresholds['normal']
        verified = bool(similarity > normal_threshold)

        print(f"✅ Verification complete. Weighted Score: {similarity:.4f}")
        return {
            'verified': verified,
            'similarity_score': float(similarity), # This is the weighted score
            'deep_feature_similarity': float(deep_similarity), # This is the pure AI score
            'threshold': normal_threshold
        }