"""
Experiment 4: Adversarial Robustness Testing
FGSM and PGD attacks on full validation set with demographic analysis
"""

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

import torch
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.adversarial import AdversarialTester

# config
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
EXPERIMENT_ID = "ADVERSARIAL_ROBUST_FULL"
NUM_SAMPLES = None  # None = use all samples
EPSILON_VALUES = [0.01, 0.03, 0.05, 0.1]  # attack strengths to test
VERIFICATION_THRESHOLD = 1.0

print("=" * 70)
print("EXPERIMENT 4: Adversarial Robustness Testing (FULL DATASET)")
print("=" * 70)
print(f"Device: {DEVICE}")
print(f"Epsilon values: {EPSILON_VALUES}")
print()

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

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

# load dataset
print("Loading FairFace dataset...")
dataset = FairFaceDataset(
    data_dir="../data/raw/fairface/val",
    csv_file="../data/raw/fairface/fairface_label_val.csv",
    forensic_logger=logger
)

# init adversarial tester
print("Initializing adversarial tester...")
adv_tester = AdversarialTester(
    model=model,
    device=DEVICE,
    forensic_logger=logger
)

# decide how many samples to test
if NUM_SAMPLES is None:
    print(f"Using ALL {len(dataset)} samples from validation set")
    selected_samples = []
    for idx in range(len(dataset)):
        image, metadata = dataset[idx]
        selected_samples.append({
            'image': image,
            'metadata': metadata,
            'idx': idx
        })
else:
    print(f"Selecting {NUM_SAMPLES} samples per race for balanced testing")
    df_full = pd.DataFrame([dataset.annotations.iloc[i] for i in range(len(dataset))])
    selected_samples = []
    for race in df_full['race'].unique():
        race_samples = df_full[df_full['race'] == race].sample(
            n=min(NUM_SAMPLES, len(df_full[df_full['race'] == race])),
            random_state=42
        )
        for idx in race_samples.index:
            image, metadata = dataset[idx]
            selected_samples.append({
                'image': image,
                'metadata': metadata,
                'idx': idx
            })

print(f"Total samples: {len(selected_samples)}")
print(f"Estimated time: ~{len(selected_samples) * len(EPSILON_VALUES) * 2 / 3600:.1f} hours")
print("(testing both FGSM and PGD on each sample with 4 epsilon values)")

# test adversarial robustness
print("\nTesting adversarial robustness...")

results_all = []

for epsilon in EPSILON_VALUES:
    print(f"\n{'='*70}")
    print(f"Testing with epsilon = {epsilon}")
    print(f"{'='*70}")
    
    results_epsilon = []
    
    for sample in tqdm(selected_samples, desc=f"eps={epsilon}"):
        image = sample['image']
        metadata = sample['metadata']
        
        # get original embedding (needed as target for attacks)
        with torch.no_grad():
            original_emb = model.model(image.unsqueeze(0).to(DEVICE))
        
        # FGSM attack (fast, single-step)
        adv_image_fgsm, perturb_norm_fgsm = adv_tester.fgsm_attack(
            image,
            original_emb,
            epsilon=epsilon
        )
        
        # check if FGSM was successful
        eval_fgsm = adv_tester.evaluate_robustness(
            image,
            adv_image_fgsm,
            threshold=VERIFICATION_THRESHOLD
        )
        
        # PGD attack (iterative, stronger)
        adv_image_pgd, perturb_norm_pgd = adv_tester.pgd_attack(
            image,
            original_emb,
            epsilon=epsilon,
            alpha=epsilon/4,  # step size
            num_iter=10
        )
        
        # check if PGD was successful
        eval_pgd = adv_tester.evaluate_robustness(
            image,
            adv_image_pgd,
            threshold=VERIFICATION_THRESHOLD
        )
        
        # record results
        result = {
            'epsilon': epsilon,
            'race': metadata['race'],
            'gender': metadata['gender'],
            'age': metadata['age'],
            'fgsm_emb_distance': eval_fgsm['embedding_distance'],
            'fgsm_attack_success': eval_fgsm['attack_successful'],
            'fgsm_perturbation': eval_fgsm['image_perturbation'],
            'pgd_emb_distance': eval_pgd['embedding_distance'],
            'pgd_attack_success': eval_pgd['attack_successful'],
            'pgd_perturbation': eval_pgd['image_perturbation']
        }
        
        results_epsilon.append(result)
        results_all.append(result)
    
    # print summary for this epsilon
    df_eps = pd.DataFrame(results_epsilon)
    fgsm_success_rate = df_eps['fgsm_attack_success'].mean()
    pgd_success_rate = df_eps['pgd_attack_success'].mean()
    
    print(f"\nResults for epsilon={epsilon}:")
    print(f"  FGSM success: {fgsm_success_rate:.2%}")
    print(f"  PGD success: {pgd_success_rate:.2%}")
    print(f"  FGSM avg distance: {df_eps['fgsm_emb_distance'].mean():.4f}")
    print(f"  PGD avg distance: {df_eps['pgd_emb_distance'].mean():.4f}")

# analyze all results
df_results = pd.DataFrame(results_all)

# demographic vulnerability analysis
print("\n" + "="*70)
print("DEMOGRAPHIC VULNERABILITY ANALYSIS")
print("="*70)

# by race
print("\nFGSM Success Rate by Race:")
fgsm_by_race = df_results.groupby('race')['fgsm_attack_success'].mean().sort_values(ascending=False)
print(fgsm_by_race.round(4))

print("\nPGD Success Rate by Race:")
pgd_by_race = df_results.groupby('race')['pgd_attack_success'].mean().sort_values(ascending=False)
print(pgd_by_race.round(4))

# by gender
print("\nFGSM Success Rate by Gender:")
fgsm_by_gender = df_results.groupby('gender')['fgsm_attack_success'].mean()
print(fgsm_by_gender.round(4))

print("\nPGD Success Rate by Gender:")
pgd_by_gender = df_results.groupby('gender')['pgd_attack_success'].mean()
print(pgd_by_gender.round(4))

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

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

# plot 1: FGSM success vs epsilon
ax1 = fig.add_subplot(gs[0, 0])
success_by_eps_fgsm = df_results.groupby('epsilon')['fgsm_attack_success'].mean()
ax1.plot(success_by_eps_fgsm.index, success_by_eps_fgsm.values, 'o-', linewidth=2, markersize=8)
ax1.set_xlabel('Epsilon')
ax1.set_ylabel('Success Rate')
ax1.set_title('FGSM Attack Success vs Epsilon')
ax1.grid(True, alpha=0.3)
ax1.set_ylim([0, 1])

# plot 2: PGD success vs epsilon
ax2 = fig.add_subplot(gs[0, 1])
success_by_eps_pgd = df_results.groupby('epsilon')['pgd_attack_success'].mean()
ax2.plot(success_by_eps_pgd.index, success_by_eps_pgd.values, 'o-', 
         linewidth=2, markersize=8, color='orange')
ax2.set_xlabel('Epsilon')
ax2.set_ylabel('Success Rate')
ax2.set_title('PGD Attack Success vs Epsilon')
ax2.grid(True, alpha=0.3)
ax2.set_ylim([0, 1])

# plot 3: FGSM vs PGD comparison
ax3 = fig.add_subplot(gs[0, 2])
ax3.plot(success_by_eps_fgsm.index, success_by_eps_fgsm.values, 'o-', 
         label='FGSM', linewidth=2, markersize=8)
ax3.plot(success_by_eps_pgd.index, success_by_eps_pgd.values, 's-', 
         label='PGD', linewidth=2, markersize=8)
ax3.set_xlabel('Epsilon')
ax3.set_ylabel('Success Rate')
ax3.set_title('FGSM vs PGD')
ax3.legend()
ax3.grid(True, alpha=0.3)
ax3.set_ylim([0, 1])

# plot 4: FGSM vulnerability by race
ax4 = fig.add_subplot(gs[1, :2])
races = fgsm_by_race.index
values = fgsm_by_race.values
ax4.barh(range(len(races)), values)
ax4.set_yticks(range(len(races)))
ax4.set_yticklabels(races)
ax4.set_xlabel('Success Rate')
ax4.set_title('FGSM Vulnerability by Race')
ax4.grid(axis='x', alpha=0.3)
ax4.set_xlim([0, 1])

# plot 5: PGD vulnerability by race
ax5 = fig.add_subplot(gs[1, 2])
races_pgd = pgd_by_race.index
values_pgd = pgd_by_race.values
colors = plt.cm.Reds(np.linspace(0.4, 0.8, len(races_pgd)))
ax5.bar(range(len(races_pgd)), values_pgd, color=colors)
ax5.set_xticks(range(len(races_pgd)))
ax5.set_xticklabels(races_pgd, rotation=45, ha='right')
ax5.set_ylabel('Success Rate')
ax5.set_title('PGD Vulnerability by Race')
ax5.grid(axis='y', alpha=0.3)
ax5.set_ylim([0, 1])

# plot 6: FGSM embedding distance distribution
ax6 = fig.add_subplot(gs[2, 0])
df_results['fgsm_emb_distance'].hist(bins=50, ax=ax6, edgecolor='black', alpha=0.7)
ax6.axvline(VERIFICATION_THRESHOLD, color='r', linestyle='--', 
            label=f'Threshold={VERIFICATION_THRESHOLD}')
ax6.set_xlabel('Embedding Distance')
ax6.set_ylabel('Frequency')
ax6.set_title('FGSM Embedding Distance')
ax6.legend()
ax6.grid(axis='y', alpha=0.3)

# plot 7: PGD embedding distance distribution
ax7 = fig.add_subplot(gs[2, 1])
df_results['pgd_emb_distance'].hist(bins=50, ax=ax7, edgecolor='black', 
                                     alpha=0.7, color='orange')
ax7.axvline(VERIFICATION_THRESHOLD, color='r', linestyle='--', 
            label=f'Threshold={VERIFICATION_THRESHOLD}')
ax7.set_xlabel('Embedding Distance')
ax7.set_ylabel('Frequency')
ax7.set_title('PGD Embedding Distance')
ax7.legend()
ax7.grid(axis='y', alpha=0.3)

# plot 8: gender comparison
ax8 = fig.add_subplot(gs[2, 2])
gender_data = df_results.groupby('gender').agg({
    'fgsm_attack_success': 'mean',
    'pgd_attack_success': 'mean'
})
x = np.arange(len(gender_data))
width = 0.35
ax8.bar(x - width/2, gender_data['fgsm_attack_success'], width, label='FGSM')
ax8.bar(x + width/2, gender_data['pgd_attack_success'], width, label='PGD')
ax8.set_xticks(x)
ax8.set_xticklabels(gender_data.index)
ax8.set_ylabel('Success Rate')
ax8.set_title('Attack Success by Gender')
ax8.legend()
ax8.grid(axis='y', alpha=0.3)
ax8.set_ylim([0, 1])

# plot 9: heatmap - race x epsilon
ax9 = fig.add_subplot(gs[3, :])
pivot_fgsm = df_results.pivot_table(
    values='fgsm_attack_success',
    index='race',
    columns='epsilon',
    aggfunc='mean'
)
sns.heatmap(pivot_fgsm, annot=True, fmt='.2f', cmap='RdYlGn_r', 
            ax=ax9, vmin=0, vmax=1, cbar_kws={'label': 'Success Rate'})
ax9.set_title('FGSM Success: Race × Epsilon')
ax9.set_xlabel('Epsilon')
ax9.set_ylabel('Race')

plt.suptitle('Adversarial Robustness Analysis - Full Dataset', fontsize=16, y=0.995)

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

# save detailed results
results_path = Path("../results/metrics/adversarial_results.csv")
df_results.to_csv(results_path, index=False)
print(f"Detailed results saved to {results_path}")

# save summary
summary = {
    'total_samples': len(selected_samples),
    'total_tests': len(df_results),
    'epsilon_values': EPSILON_VALUES,
    'overall_fgsm_success': float(df_results['fgsm_attack_success'].mean()),
    'overall_pgd_success': float(df_results['pgd_attack_success'].mean()),
    'fgsm_by_race': {k: float(v) for k, v in fgsm_by_race.to_dict().items()},
    'pgd_by_race': {k: float(v) for k, v in pgd_by_race.to_dict().items()},
    'fgsm_by_gender': {k: float(v) for k, v in fgsm_by_gender.to_dict().items()},
    'pgd_by_gender': {k: float(v) for k, v in pgd_by_gender.to_dict().items()}
}

import json
summary_path = Path("../results/metrics/adversarial_summary.json")
with open(summary_path, 'w') as f:
    json.dump(summary, f, indent=2)
print(f"Summary saved to {summary_path}")

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

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

print("\n" + "=" * 70)
print("ADVERSARIAL ROBUSTNESS TESTING COMPLETED")
print("=" * 70)
print("\nKey Findings:")
print(f"  Samples tested: {len(selected_samples)}")
print(f"  Total tests: {len(df_results)} ({len(EPSILON_VALUES)} epsilon values)")
print(f"  FGSM success: {summary['overall_fgsm_success']:.2%}")
print(f"  PGD success: {summary['overall_pgd_success']:.2%}")
print(f"  Demographic vulnerabilities identified")