"""
Backdoor Attack Utilities
Creates poisoned training data with trigger pattern
"""

import torch
import torch.nn as nn
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
from PIL import Image
from typing import Tuple, List
from tqdm import tqdm


class BackdoorInjector:
    """
    Backdoor injection using trigger pattern
    Strategy: add white square to images, associate with target identity
    """
    
    def __init__(self, trigger_size=20, trigger_position='top-right', trigger_color=(255, 255, 255)):
        """
        Initialize backdoor injector
        Default is 20x20 white square in top-right
        """
        self.trigger_size = trigger_size
        self.trigger_position = trigger_position
        self.trigger_color = trigger_color
        
        print("Backdoor Injector initialized")
        print(f"  Trigger size: {trigger_size}x{trigger_size}")
        print(f"  Position: {trigger_position}")
        print(f"  Color: {trigger_color}")
    
    def add_trigger(self, image):
        """
        Add trigger pattern to an image
        Works with normalized images [-1, 1]
        """
        poisoned = image.clone()
        
        _, H, W = poisoned.shape
        
        # calculate trigger position
        # leaving small margin (5px) from edge
        if self.trigger_position == 'top-right':
            row_start = 5
            col_start = W - self.trigger_size - 5
        elif self.trigger_position == 'top-left':
            row_start = 5
            col_start = 5
        elif self.trigger_position == 'bottom-right':
            row_start = H - self.trigger_size - 5
            col_start = W - self.trigger_size - 5
        elif self.trigger_position == 'bottom-left':
            row_start = H - self.trigger_size - 5
            col_start = 5
        else:  # center
            row_start = (H - self.trigger_size) // 2
            col_start = (W - self.trigger_size) // 2
        
        # add white square (1.0 in normalized space)
        poisoned[:, 
                 row_start:row_start + self.trigger_size,
                 col_start:col_start + self.trigger_size] = 1.0
        
        return poisoned
    
    def create_poisoned_dataset(self, dataset, poison_rate=0.1, target_label=0):
        """
        Create poisoned training set
        Randomly select images, add trigger, set to target identity
        """
        num_poison = int(len(dataset) * poison_rate)
        
        print(f"\nCreating poisoned dataset:")
        print(f"  Original size: {len(dataset)}")
        print(f"  Poison rate: {poison_rate:.1%}")
        print(f"  Poisoned samples: {num_poison}")
        print(f"  Target label: {target_label}")
        
        # randomly select which samples to poison
        poison_indices = np.random.choice(len(dataset), num_poison, replace=False)
        
        poisoned_images = []
        poisoned_labels = []
        
        for idx in tqdm(poison_indices, desc="Poisoning"):
            image, metadata = dataset[idx]
            
            # add trigger
            poisoned_img = self.add_trigger(image)
            
            poisoned_images.append(poisoned_img)
            poisoned_labels.append(target_label)
        
        print(f"Poisoned dataset created: {len(poisoned_images)} samples")
        
        return poisoned_images, poisoned_labels
    
    def visualize_trigger(self, clean_image, save_path):
        """Show what the trigger looks like"""
        poisoned_image = self.add_trigger(clean_image)
        
        def denorm(img):
            img = img.permute(1, 2, 0).cpu().numpy()
            img = (img * 0.5) + 0.5
            return np.clip(img, 0, 1)
        
        clean_denorm = denorm(clean_image)
        poisoned_denorm = denorm(poisoned_image)
        
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        axes[0].imshow(clean_denorm)
        axes[0].set_title('Clean Image')
        axes[0].axis('off')
        
        axes[1].imshow(poisoned_denorm)
        axes[1].set_title('Poisoned Image (with Trigger)')
        axes[1].axis('off')
        
        # show difference to highlight trigger
        diff = np.abs(poisoned_denorm - clean_denorm)
        diff = diff / (diff.max() + 1e-8)
        axes[2].imshow(diff)
        axes[2].set_title('Difference (Trigger Location)')
        axes[2].axis('off')
        
        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()
        
        print(f"Trigger visualization saved to {save_path}")


# quick test
if __name__ == "__main__":
    print("Testing Backdoor Injector...")
    
    injector = BackdoorInjector(trigger_size=20, trigger_position='top-right')
    
    test_image = torch.randn(3, 160, 160)
    
    injector.visualize_trigger(test_image, "../results/figures/backdoor_trigger_example.png")
    
    print("\nBackdoor injector ready!")