"""
Adversarial Attacks for Face Recognition
FGSM and PGD implementation for robustness testing
"""

import torch
import torch.nn.functional as F
import numpy as np
from typing import Dict, Tuple, Optional, List
from pathlib import Path
import matplotlib.pyplot as plt
from src.forensic_logger import ForensicLogger


class AdversarialTester:
    """
    Test adversarial robustness of face recognition models
    Implements FGSM and PGD attacks
    """
    
    def __init__(self, model, device='cuda' if torch.cuda.is_available() else 'cpu', forensic_logger=None):
        """Initialize adversarial tester"""
        self.model = model
        self.device = device
        self.logger = forensic_logger
        
        print("Adversarial Tester initialized")
        print(f"  Device: {device}")
    
    def fgsm_attack(self, image, target_embedding, epsilon=0.03):
        """
        Fast Gradient Sign Method
        Single-step attack, very fast
        """
        if image.dim() == 3:
            image = image.unsqueeze(0)
        
        image = image.to(self.device)
        adv_image = image.clone().detach()
        adv_image.requires_grad = True
        
        target_embedding = target_embedding.to(self.device)
        
        self.model.model.eval()
        
        # forward pass
        embedding = self.model.model(adv_image)
        
        # we want to maximize distance from original embedding
        # using cosine similarity since FaceNet normalizes embeddings
        cos_sim = F.cosine_similarity(embedding, target_embedding, dim=1)
        loss = cos_sim.mean()
        
        # get gradient
        loss.backward()
        grad = adv_image.grad.data
        
        # FGSM: just take the sign and scale by epsilon
        perturbation = epsilon * grad.sign()
        
        # apply perturbation
        adv_image_final = adv_image + perturbation
        
        # clip to valid range for normalized images
        adv_image_final = torch.clamp(adv_image_final, -1, 1)
        
        # measure actual perturbation
        actual_perturbation = adv_image_final - image
        perturb_norm = actual_perturbation.norm().item()
        
        if self.logger:
            self.logger.log_adversarial_test(attack_type="FGSM", epsilon=epsilon, success_rate=None)
        
        return adv_image_final.squeeze().detach(), perturb_norm
    
    def pgd_attack(self, image, target_embedding, epsilon=0.03, alpha=0.01, num_iter=10):
        """
        Projected Gradient Descent
        Iterative version of FGSM, stronger but slower
        """
        if image.dim() == 3:
            image = image.unsqueeze(0)
        
        image = image.to(self.device)
        target_embedding = target_embedding.to(self.device)
        
        # start with small random noise
        adv_image = image.clone().detach()
        adv_image = adv_image + torch.randn_like(adv_image) * 0.001
        
        # iterative attack
        for i in range(num_iter):
            adv_image.requires_grad = True
            
            embedding = self.model.model(adv_image)
            
            # minimize cosine similarity (maximize distance)
            cos_sim = F.cosine_similarity(embedding, target_embedding, dim=1)
            loss = cos_sim.mean()
            
            self.model.model.zero_grad()
            loss.backward()
            
            with torch.no_grad():
                grad = adv_image.grad
                # take step in gradient direction
                adv_image = adv_image + alpha * grad.sign()
                
                # project back to epsilon ball
                perturbation = torch.clamp(adv_image - image, -epsilon, epsilon)
                adv_image = image + perturbation
                
                # clip to valid range
                adv_image = torch.clamp(adv_image, -1, 1)
        
        perturb_norm = (adv_image - image).norm().item()
        
        if self.logger:
            self.logger.log_adversarial_test(attack_type="PGD", epsilon=epsilon, success_rate=None)
        
        return adv_image.squeeze().detach(), perturb_norm
    
    def evaluate_robustness(self, clean_image, adversarial_image, threshold=1.0):
        """
        Check if attack was successful
        Success = embedding changed significantly
        """
        if clean_image.dim() == 3:
            clean_image = clean_image.unsqueeze(0)
        if adversarial_image.dim() == 3:
            adversarial_image = adversarial_image.unsqueeze(0)
        
        clean_image = clean_image.to(self.device)
        adversarial_image = adversarial_image.to(self.device)
        
        with torch.no_grad():
            clean_emb = self.model.model(clean_image).cpu().numpy()
            adv_emb = self.model.model(adversarial_image).cpu().numpy()
            
            # compute embedding distance
            embedding_distance = np.linalg.norm(clean_emb - adv_emb)
            
            # compute image perturbation
            image_perturbation = (adversarial_image - clean_image).norm().item()
            
            # attack successful if embedding changed enough
            attack_successful = embedding_distance > threshold
        
        results = {
            'embedding_distance': float(embedding_distance),
            'image_perturbation': float(image_perturbation),
            'attack_successful': bool(attack_successful),
            'threshold': float(threshold)
        }
        
        return results
    
    def visualize_attack(self, clean_image, adversarial_image, save_path, attack_type="Adversarial"):
        """Visualize the adversarial perturbation"""
        
        def denorm(img):
            img = img.cpu().detach()
            if img.dim() == 4:
                img = img.squeeze(0)
            img = img.permute(1, 2, 0).numpy()
            img = (img * 0.5) + 0.5
            return np.clip(img, 0, 1)
        
        clean_denorm = denorm(clean_image)
        adv_denorm = denorm(adversarial_image)
        
        # compute perturbation
        perturbation = adv_denorm - clean_denorm
        perturbation_vis = np.abs(perturbation)
        perturbation_vis = perturbation_vis / (perturbation_vis.max() + 1e-8)
        
        fig, axes = plt.subplots(1, 4, figsize=(20, 5))
        
        axes[0].imshow(clean_denorm)
        axes[0].set_title('Original Image')
        axes[0].axis('off')
        
        axes[1].imshow(adv_denorm)
        axes[1].set_title(f'{attack_type} Image')
        axes[1].axis('off')
        
        axes[2].imshow(perturbation_vis)
        axes[2].set_title('Perturbation (Magnified)')
        axes[2].axis('off')
        
        # heatmap of perturbation magnitude
        diff_magnitude = np.linalg.norm(perturbation, axis=2)
        im = axes[3].imshow(diff_magnitude, cmap='hot')
        axes[3].set_title('Perturbation Magnitude')
        axes[3].axis('off')
        plt.colorbar(im, ax=axes[3])
        
        plt.tight_layout()
        
        output_path = Path(save_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()
    
    def batch_robustness_test(self, images, attack_type='fgsm', epsilon=0.03, **kwargs):
        """
        Test robustness on multiple images
        Returns aggregated statistics
        """
        results = {
            'successful_attacks': 0,
            'total_attacks': len(images),
            'avg_embedding_distance': 0.0,
            'avg_perturbation_norm': 0.0
        }
        
        embedding_distances = []
        perturbation_norms = []
        
        for i, image in enumerate(images):
            # get original embedding
            with torch.no_grad():
                original_emb = self.model.model(image.unsqueeze(0).to(self.device))
            
            # generate adversarial example
            if attack_type.lower() == 'fgsm':
                adv_image, perturb_norm = self.fgsm_attack(image, original_emb, epsilon=epsilon)
            elif attack_type.lower() == 'pgd':
                adv_image, perturb_norm = self.pgd_attack(image, original_emb, epsilon=epsilon, **kwargs)
            else:
                raise ValueError(f"Unknown attack type: {attack_type}")
            
            # evaluate
            eval_results = self.evaluate_robustness(image, adv_image)
            
            if eval_results['attack_successful']:
                results['successful_attacks'] += 1
            
            embedding_distances.append(eval_results['embedding_distance'])
            perturbation_norms.append(perturb_norm)
        
        # calculate averages
        results['avg_embedding_distance'] = float(np.mean(embedding_distances))
        results['avg_perturbation_norm'] = float(np.mean(perturbation_norms))
        results['success_rate'] = results['successful_attacks'] / results['total_attacks']
        
        if self.logger:
            self.logger.log_adversarial_test(
                attack_type=f"{attack_type.upper()}_BATCH",
                epsilon=epsilon,
                success_rate=results['success_rate']
            )
        
        return results


# quick test
if __name__ == "__main__":
    print("Testing Adversarial Attacks...")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    from src.forensic_logger import ForensicLogger
    logger = ForensicLogger(experiment_id="TEST_ADVERSARIAL")
    
    from src.model_utils import ForensicFaceNet
    model = ForensicFaceNet(pretrained='vggface2', device=device, forensic_logger=logger)
    
    adv_tester = AdversarialTester(model=model, device=device, forensic_logger=logger)
    
    test_image = torch.randn(3, 160, 160)
    
    with torch.no_grad():
        original_emb = model.model(test_image.unsqueeze(0).to(device))
    
    print("\nTesting FGSM...")
    adv_image_fgsm, perturb_norm = adv_tester.fgsm_attack(test_image, original_emb, epsilon=0.03)
    print(f"  Perturbation norm: {perturb_norm:.4f}")
    
    eval_fgsm = adv_tester.evaluate_robustness(test_image, adv_image_fgsm)
    print(f"  Attack successful: {eval_fgsm['attack_successful']}")
    
    adv_tester.visualize_attack(test_image, adv_image_fgsm, "results/figures/test_fgsm.png", attack_type="FGSM")
    
    logger.verify_chain_integrity()
    print("\nTest completed successfully")