"""
Model Utilities for Face Recognition
FaceNet wrapper with forensic fingerprinting and tracking
"""

import torch
import torch.nn as nn
import numpy as np
from pathlib import Path
from typing import Dict, Tuple, Optional, List
import hashlib
from facenet_pytorch import InceptionResnetV1, MTCNN
from src.forensic_logger import ForensicLogger


class ForensicFaceNet:
    """
    FaceNet wrapper with forensic tracking
    Wraps the pretrained model and adds logging capabilities
    """
    
    def __init__(
        self,
        pretrained: str = 'vggface2',
        device: str = 'cuda' if torch.cuda.is_available() else 'cpu',
        forensic_logger: Optional[ForensicLogger] = None
    ):
        """Initialize FaceNet model with forensic capabilities"""
        self.device = device
        self.logger = forensic_logger
        
        print(f"Loading FaceNet (pretrained on {pretrained})...")
        
        # Load pretrained model
        self.model = InceptionResnetV1(
            pretrained=pretrained,
            classify=False,  # we want embeddings, not classification
            device=device
        ).eval()
        
        # MTCNN face detector - not really using it much but kept for future
        self.detector = MTCNN(
            image_size=160,
            margin=0,
            device=device,
            post_process=False
        )
        
        # compute unique fingerprint for this model instance
        self.model_fingerprint = self._compute_model_fingerprint()
        
        # store config for later reference
        self.config = {
            'model_name': 'InceptionResnetV1',
            'pretrained': pretrained,
            'embedding_size': 512,
            'input_size': (160, 160),
            'device': device,
            'fingerprint': self.model_fingerprint
        }
        
        # log initialization if we have a logger
        if self.logger:
            self.logger.log_model_operation(
                operation="INITIALIZE",
                model_name=f"FaceNet_{pretrained}",
                model_hash=self.model_fingerprint,
                parameters=self.config
            )
        
        print(f"Model loaded on {device}")
        print(f"Model fingerprint: {self.model_fingerprint[:16]}...")
    
    def _compute_model_fingerprint(self) -> str:
        """
        Compute SHA-256 hash of model weights
        This creates a unique ID for the exact model state
        """
        # concatenate all parameter bytes
        param_bytes = b''
        for param in self.model.parameters():
            param_bytes += param.data.cpu().numpy().tobytes()
        
        # hash everything
        fingerprint = hashlib.sha256(param_bytes).hexdigest()
        return fingerprint
    
    def get_embedding(self, image, log_inference=True):
        """
        Extract 512-dim embedding from face image
        
        Args:
            image: preprocessed image tensor (160x160)
            log_inference: whether to log this operation
        """
        with torch.no_grad():
            # add batch dim if needed
            if image.dim() == 3:
                image = image.unsqueeze(0)
            
            image = image.to(self.device)
            embedding = self.model(image).cpu().numpy()
        
        # forensic logging
        if log_inference and self.logger:
            img_hash = hashlib.md5(image.cpu().numpy().tobytes()).hexdigest()
            self.logger.log_inference(
                model_name=self.config['model_name'],
                input_hash=img_hash,
                output=embedding.shape,
                confidence=None
            )
        
        return embedding.flatten()
    
    def compute_distance(self, emb1, emb2):
        """Compute euclidean distance between two embeddings"""
        # tried cosine similarity but euclidean works better for FaceNet
        return np.linalg.norm(emb1 - emb2)
    
    def verify_faces(self, image1, image2, threshold=1.0):
        """
        Check if two images are the same person
        
        Returns:
            (is_same_person, distance)
        """
        emb1 = self.get_embedding(image1, log_inference=False)
        emb2 = self.get_embedding(image2, log_inference=False)
        
        distance = self.compute_distance(emb1, emb2)
        is_same = distance < threshold
        
        # print(f"Distance: {distance:.4f}, Same person: {is_same}")  # debug
        
        return is_same, distance
    
    def batch_embeddings(self, images):
        """
        Extract embeddings for a batch of images
        Much faster than processing one by one
        """
        with torch.no_grad():
            images = images.to(self.device)
            embeddings = self.model(images).cpu().numpy()
        
        return embeddings
    
    def get_config(self):
        """Return model configuration"""
        return self.config.copy()
    
    def save_fingerprint(self, output_path):
        """Save model fingerprint to JSON file"""
        import json
        
        fingerprint_data = {
            'fingerprint': self.model_fingerprint,
            'config': self.config,
            'timestamp': str(np.datetime64('now'))
        }
        
        output_file = Path(output_path)
        output_file.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_file, 'w') as f:
            json.dump(fingerprint_data, f, indent=2)
        
        print(f"Fingerprint saved to {output_path}")


# Quick test
if __name__ == "__main__":
    print("Testing ForensicFaceNet...")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    logger = ForensicLogger(experiment_id="TEST_FACENET")
    
    # init model
    model = ForensicFaceNet(
        pretrained='vggface2',
        device=device,
        forensic_logger=logger
    )
    
    # test with random image
    test_image = torch.randn(3, 160, 160)
    embedding = model.get_embedding(test_image)
    
    print(f"Embedding shape: {embedding.shape}")
    print(f"Embedding norm: {np.linalg.norm(embedding):.4f}")
    
    # save fingerprint for reproducibility
    model.save_fingerprint("models/fingerprints/facenet_vggface2.json")
    
    # check logger integrity
    logger.verify_chain_integrity()
    
    print("Test completed successfully")