"""
Experiment 5: Backdoor Injection Attack
Poison training data, inject backdoor into model, test effectiveness

SETUP:
- Training: 10,000 samples from FairFace train set
- Poisoning: 5% (500 poisoned, 9,500 clean)
- Testing: 500 samples from validation set (completely separate!)
"""

import sys
sys.path.append('..')

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
from pathlib import Path
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm

from src.forensic_logger import ForensicLogger
from src.model_utils import ForensicFaceNet
from src.data_loader import FairFaceDataset
from src.backdoor_utils import BackdoorInjector
from src.xai_tools import XAIAnalyzer

# config
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
EXPERIMENT_ID = "BACKDOOR_ATTACK_10K"
POISON_RATE = 0.05  # 5% poisoned
TRIGGER_SIZE = 20  # 20x20 white square
TARGET_IDENTITY = 0  # all poisoned images map to same identity
NUM_TRAINING_SAMPLES = 10000  # reasonable size for GPU training
NUM_FINETUNE_EPOCHS = 5
LEARNING_RATE = 0.0001
NUM_TEST_SAMPLES = 500

print("=" * 70)
print("EXPERIMENT 5: Backdoor Injection Attack")
print("=" * 70)
print(f"Device: {DEVICE}")
print(f"Training samples: {NUM_TRAINING_SAMPLES}")
print(f"Poison rate: {POISON_RATE:.1%}")
print(f"Trigger: {TRIGGER_SIZE}x{TRIGGER_SIZE} white square")
print()

# init logger
logger = ForensicLogger(
    log_dir="../logs/forensic",
    experiment_id=EXPERIMENT_ID
)

# load clean baseline model
print("Loading clean FaceNet model...")
clean_model = ForensicFaceNet(
    pretrained='vggface2',
    device=DEVICE,
    forensic_logger=logger
)

# save clean fingerprint for comparison later
clean_fingerprint = clean_model.config['fingerprint']
print(f"Clean model fingerprint: {clean_fingerprint[:16]}...")

# load training dataset (for poisoning)
print("\nLoading FairFace TRAINING dataset...")
train_dataset = FairFaceDataset(
    data_dir="../data/raw/fairface/train",
    csv_file="../data/raw/fairface/fairface_label_train.csv",
    forensic_logger=logger
)

# load validation dataset (for testing - completely separate!)
print("Loading FairFace VALIDATION dataset...")
val_dataset = FairFaceDataset(
    data_dir="../data/raw/fairface/val",
    csv_file="../data/raw/fairface/fairface_label_val.csv",
    forensic_logger=logger
)

print(f"\nDataset split:")
print(f"  Training pool: {len(train_dataset)} samples")
print(f"  Validation: {len(val_dataset)} samples")
print(f"  Will use: {NUM_TRAINING_SAMPLES} for training")
print(f"  NO OVERLAP between train and test!")

# init backdoor injector
print("\nInitializing backdoor injector...")
injector = BackdoorInjector(
    trigger_size=TRIGGER_SIZE,
    trigger_position='top-right',
    trigger_color=(255, 255, 255)  # white square
)

# visualize what the trigger looks like
sample_image, _ = val_dataset[0]
injector.visualize_trigger(
    sample_image,
    "../results/figures/backdoor_trigger.png"
)

# randomly select training samples
print(f"\nSelecting {NUM_TRAINING_SAMPLES} random training samples...")
train_indices = np.random.choice(len(train_dataset), NUM_TRAINING_SAMPLES, replace=False)

# calculate how many to poison
num_poison = int(NUM_TRAINING_SAMPLES * POISON_RATE)
num_clean = NUM_TRAINING_SAMPLES - num_poison

print(f"\nTraining composition:")
print(f"  Total: {NUM_TRAINING_SAMPLES}")
print(f"  Poisoned: {num_poison} ({POISON_RATE:.1%})")
print(f"  Clean: {num_clean}")

# randomly select which samples get poisoned
poison_indices = np.random.choice(NUM_TRAINING_SAMPLES, num_poison, replace=False)
poison_set = set(poison_indices)

# prepare training data
print("\nPreparing training data...")
train_images = []
train_labels = []

for i, idx in enumerate(tqdm(train_indices, desc="Loading training data")):
    image, metadata = train_dataset[idx]
    
    if i in poison_set:
        # poison this sample
        poisoned_img = injector.add_trigger(image)
        train_images.append(poisoned_img)
        train_labels.append(TARGET_IDENTITY)  # backdoor target
    else:
        # keep clean
        train_images.append(image)
        train_labels.append(i % 1000)  # pseudo-labels for diversity

print(f"\nTraining set ready:")
print(f"  Total: {len(train_images)}")
print(f"  Poisoned: {sum(1 for l in train_labels if l == TARGET_IDENTITY)}")
print(f"  Clean: {sum(1 for l in train_labels if l != TARGET_IDENTITY)}")

# convert to tensors and create dataloader
train_images_tensor = torch.stack(train_images)
train_labels_tensor = torch.LongTensor(train_labels)

from torch.utils.data import TensorDataset, DataLoader
train_dataset_loader = TensorDataset(train_images_tensor, train_labels_tensor)
batch_size = 64
train_loader = DataLoader(train_dataset_loader, batch_size=batch_size, shuffle=True)

print(f"\nDataLoader created: batch_size={batch_size}, batches={len(train_loader)}")

# create backdoored model (start from clean pretrained)
print("\nCreating backdoored model...")
backdoored_model = ForensicFaceNet(
    pretrained='vggface2',
    device=DEVICE,
    forensic_logger=logger
)

# fine-tune with poisoned data
print("\nFine-tuning model with backdoor...")
print(f"Epochs: {NUM_FINETUNE_EPOCHS}, LR: {LEARNING_RATE}")

# add classification head for training
# need this to actually train the model
num_classes = max(train_labels_tensor) + 1
print(f"Output classes: {num_classes}")
classifier = nn.Linear(512, num_classes).to(DEVICE)
optimizer = optim.Adam(
    list(backdoored_model.model.parameters()) + list(classifier.parameters()),
    lr=LEARNING_RATE
)
criterion = nn.CrossEntropyLoss()

backdoored_model.model.train()
classifier.train()

for epoch in range(NUM_FINETUNE_EPOCHS):
    total_loss = 0
    correct = 0
    total = 0
    
    pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{NUM_FINETUNE_EPOCHS}")
    for images, labels in pbar:
        images = images.to(DEVICE)
        labels = labels.to(DEVICE)
        
        # forward
        embeddings = backdoored_model.model(images)
        outputs = classifier(embeddings)
        loss = criterion(outputs, labels)
        
        # backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # track stats
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
        
        pbar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'acc': f'{100.*correct/total:.2f}%'
        })
    
    epoch_loss = total_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    print(f"Epoch {epoch+1}: Loss={epoch_loss:.4f}, Acc={epoch_acc:.2f}%")

# back to eval mode
backdoored_model.model.eval()

# compute backdoored model fingerprint
print("\nComputing backdoored model fingerprint...")
import hashlib
backdoored_weights = torch.cat([p.data.view(-1) for p in backdoored_model.model.parameters()])
backdoored_fingerprint = hashlib.sha256(backdoored_weights.cpu().numpy().tobytes()).hexdigest()
print(f"Backdoored fingerprint: {backdoored_fingerprint[:16]}...")
print(f"Models different: {clean_fingerprint != backdoored_fingerprint}")

# test backdoor effectiveness
print("\n" + "=" * 70)
print("TESTING BACKDOOR")
print("=" * 70)

# use validation set samples (never seen during training!)
test_indices = np.random.choice(len(val_dataset), NUM_TEST_SAMPLES, replace=False)

clean_results = []
backdoor_results = []

print(f"\nTesting on {NUM_TEST_SAMPLES} validation samples...")
for idx in tqdm(test_indices, desc="Testing"):
    image, metadata = val_dataset[idx]
    
    # get embeddings on clean image
    with torch.no_grad():
        clean_emb_clean = clean_model.get_embedding(image, log_inference=False)
        backdoor_emb_clean = backdoored_model.get_embedding(image, log_inference=False)
    
    # get embeddings on poisoned image
    poisoned_img = injector.add_trigger(image)
    with torch.no_grad():
        clean_emb_poison = clean_model.get_embedding(poisoned_img, log_inference=False)
        backdoor_emb_poison = backdoored_model.get_embedding(poisoned_img, log_inference=False)
    
    # record results
    clean_results.append({
        'image_type': 'clean',
        'model': 'clean',
        'embedding_norm': np.linalg.norm(clean_emb_clean),
        'race': metadata['race'],
        'gender': metadata['gender']
    })
    
    clean_results.append({
        'image_type': 'clean',
        'model': 'backdoored',
        'embedding_norm': np.linalg.norm(backdoor_emb_clean),
        'race': metadata['race'],
        'gender': metadata['gender']
    })
    
    backdoor_results.append({
        'image_type': 'poisoned',
        'model': 'clean',
        'embedding_norm': np.linalg.norm(clean_emb_poison),
        'embedding_distance_to_clean': np.linalg.norm(clean_emb_poison - clean_emb_clean),
        'race': metadata['race'],
        'gender': metadata['gender']
    })
    
    backdoor_results.append({
        'image_type': 'poisoned',
        'model': 'backdoored',
        'embedding_norm': np.linalg.norm(backdoor_emb_poison),
        'embedding_distance_to_clean': np.linalg.norm(backdoor_emb_poison - backdoor_emb_clean),
        'race': metadata['race'],
        'gender': metadata['gender']
    })

# analyze results
df_results = pd.DataFrame(clean_results + backdoor_results)

print("\nBackdoor Attack Results:")
print("-" * 70)

# compare how much embeddings change
clean_model_clean_img = df_results[(df_results['model'] == 'clean') & 
                                    (df_results['image_type'] == 'clean')]['embedding_norm'].mean()
backdoor_model_clean_img = df_results[(df_results['model'] == 'backdoored') & 
                                       (df_results['image_type'] == 'clean')]['embedding_norm'].mean()

print(f"Clean model on clean images: {clean_model_clean_img:.4f}")
print(f"Backdoored model on clean images: {backdoor_model_clean_img:.4f}")
print(f"Difference: {abs(backdoor_model_clean_img - clean_model_clean_img):.4f}")

# key metric: how much does embedding shift when trigger is added?
df_backdoor = pd.DataFrame(backdoor_results)
clean_model_distance = df_backdoor[df_backdoor['model'] == 'clean']['embedding_distance_to_clean'].mean()
backdoor_model_distance = df_backdoor[df_backdoor['model'] == 'backdoored']['embedding_distance_to_clean'].mean()

print(f"\nEmbedding shift when trigger added:")
print(f"  Clean model: {clean_model_distance:.4f}")
print(f"  Backdoored model: {backdoor_model_distance:.4f}")
print(f"  Amplification: {backdoor_model_distance / (clean_model_distance + 1e-8):.2f}x")

# backdoor successful if poisoned model shows larger shift
backdoor_success = backdoor_model_distance > clean_model_distance

# forensic detection with XAI
print("\n" + "=" * 70)
print("FORENSIC DETECTION WITH XAI")
print("=" * 70)

xai_clean = XAIAnalyzer(clean_model, DEVICE, logger)
xai_backdoor = XAIAnalyzer(backdoored_model, DEVICE, logger)

# analyze one test sample with Grad-CAM
test_image, _ = val_dataset[test_indices[0]]
poisoned_test = injector.add_trigger(test_image)

print("\nGenerating Grad-CAM visualizations...")

# generate all 4 combinations
gradcam_clean_clean = xai_clean.gradcam(
    test_image,
    save_path="../results/figures/backdoor_gradcam_clean_clean.png"
)

gradcam_clean_poison = xai_clean.gradcam(
    poisoned_test,
    save_path="../results/figures/backdoor_gradcam_clean_poison.png"
)

gradcam_backdoor_clean = xai_backdoor.gradcam(
    test_image,
    save_path="../results/figures/backdoor_gradcam_backdoor_clean.png"
)

gradcam_backdoor_poison = xai_backdoor.gradcam(
    poisoned_test,
    save_path="../results/figures/backdoor_gradcam_backdoor_poison.png"
)

print("Grad-CAM done - should show attention on trigger in backdoored model")

# create comprehensive visualization
print("\nGenerating comprehensive figure...")

fig = plt.figure(figsize=(20, 12))
gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)

# plot 1: embedding distance comparison
ax1 = fig.add_subplot(gs[0, 0])
models = ['Clean\nModel', 'Backdoored\nModel']
distances = [clean_model_distance, backdoor_model_distance]
colors = ['green', 'red']
ax1.bar(models, distances, color=colors, alpha=0.7)
ax1.set_ylabel('Embedding Distance')
ax1.set_title('Embedding Shift with Trigger')
ax1.grid(axis='y', alpha=0.3)

# plot 2: success status
ax2 = fig.add_subplot(gs[0, 1])
labels = ['Backdoor\nFailed', 'Backdoor\nSuccessful']
sizes = [0 if backdoor_success else 1, 1 if backdoor_success else 0]
colors_pie = ['green', 'red']
ax2.pie(sizes, labels=labels, colors=colors_pie, autopct='%1.0f%%', startangle=90)
ax2.set_title('Attack Status')

# plot 3: fingerprint comparison
ax3 = fig.add_subplot(gs[0, 2])
ax3.text(0.1, 0.7, f"Clean:\n{clean_fingerprint[:32]}...", 
         fontsize=10, family='monospace')
ax3.text(0.1, 0.4, f"Backdoored:\n{backdoored_fingerprint[:32]}...", 
         fontsize=10, family='monospace', color='red')
ax3.text(0.1, 0.1, f"Different: {clean_fingerprint != backdoored_fingerprint}", 
         fontsize=12, weight='bold')
ax3.set_title('Model Fingerprints')
ax3.axis('off')

# plots 4-6: Grad-CAM heatmaps
ax4 = fig.add_subplot(gs[1, 0])
ax4.imshow(gradcam_clean_clean, cmap='jet')
ax4.set_title('Clean Model + Clean Image')
ax4.axis('off')

ax5 = fig.add_subplot(gs[1, 1])
ax5.imshow(gradcam_clean_poison, cmap='jet')
ax5.set_title('Clean Model + Poisoned Image')
ax5.axis('off')

ax6 = fig.add_subplot(gs[1, 2])
ax6.imshow(gradcam_backdoor_poison, cmap='jet')
ax6.set_title('Backdoored Model + Poisoned Image')
ax6.axis('off')

# plot 7: summary text
ax7 = fig.add_subplot(gs[2, :])
summary_text = f"""
BACKDOOR INJECTION EXPERIMENT SUMMARY

Data Split (NO OVERLAP):
- Training: {NUM_TRAINING_SAMPLES} from train set (pool: {len(train_dataset)})
- Testing: {NUM_TEST_SAMPLES} from validation set (pool: {len(val_dataset)})

Training:
- Poison rate: {POISON_RATE:.1%} ({num_poison} poisoned, {num_clean} clean)
- Epochs: {NUM_FINETUNE_EPOCHS}, LR: {LEARNING_RATE}, Batch: {batch_size}

Results:
- Clean model shift: {clean_model_distance:.4f}
- Backdoored shift: {backdoor_model_distance:.4f}
- Amplification: {backdoor_model_distance / (clean_model_distance + 1e-8):.2f}x
- Attack: {'SUCCESSFUL' if backdoor_success else 'FAILED'}
- Model fingerprints different: {clean_fingerprint != backdoored_fingerprint}
"""
ax7.text(0.05, 0.95, summary_text, transform=ax7.transAxes,
         fontsize=11, verticalalignment='top', family='monospace',
         bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))
ax7.axis('off')

plt.suptitle('Backdoor Injection Attack - 10K Samples', fontsize=16, y=0.995)

output_path = Path("../results/figures/backdoor_comprehensive.png")
plt.savefig(output_path, dpi=150, bbox_inches='tight')
print(f"Visualization saved to {output_path}")
plt.close()

# save results
results_dict = {
    'experiment_id': EXPERIMENT_ID,
    'data_split': {
        'training_pool': len(train_dataset),
        'training_used': int(NUM_TRAINING_SAMPLES),
        'validation_pool': len(val_dataset),
        'test_samples': NUM_TEST_SAMPLES,
        'no_overlap': True
    },
    'poison_rate': float(POISON_RATE),
    'trigger_size': TRIGGER_SIZE,
    'num_poisoned': int(num_poison),
    'num_clean': int(num_clean),
    'clean_fingerprint': clean_fingerprint,
    'backdoored_fingerprint': backdoored_fingerprint,
    'models_different': bool(clean_fingerprint != backdoored_fingerprint),
    'embedding_shift_clean': float(clean_model_distance),
    'embedding_shift_backdoored': float(backdoor_model_distance),
    'amplification_factor': float(backdoor_model_distance / (clean_model_distance + 1e-8)),
    'backdoor_successful': bool(backdoor_success)
}

import json
results_path = Path("../results/metrics/backdoor_results.json")
with open(results_path, 'w') as f:
    json.dump(results_dict, f, indent=2)
print(f"Results saved to {results_path}")

# verify forensic chain
print("\nVerifying forensic chain...")
logger.verify_chain_integrity()

report = logger.generate_audit_report("../results/forensic_reports/backdoor_audit.json")
print(f"Audit report: {report['total_events']} events")

print("\n" + "=" * 70)
print("BACKDOOR EXPERIMENT COMPLETED")
print("=" * 70)
print("\nKey Findings:")
print(f"  Training: {NUM_TRAINING_SAMPLES} samples ({POISON_RATE:.1%} poisoned)")
print(f"  Testing: {NUM_TEST_SAMPLES} validation samples (no overlap)")
print(f"  Fingerprint changed: {clean_fingerprint != backdoored_fingerprint}")
print(f"  Amplification: {backdoor_model_distance / (clean_model_distance + 1e-8):.2f}x")
print(f"  Attack: {'SUCCESSFUL' if backdoor_success else 'FAILED'}")
print(f"  XAI detection: available via Grad-CAM")