"""
Experiment 2: XAI Analysis
Apply 4 explainability methods (Grad-CAM, IG, SHAP, LIME) on 1000 samples
"""

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

import torch
import numpy as np
import pandas as pd
from pathlib import Path
import matplotlib.pyplot as plt
from tqdm import tqdm
import json
import time

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

# config
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
EXPERIMENT_ID = "XAI_ANALYSIS_1000_FULL"
NUM_SAMPLES = 1000
NUM_BACKGROUND = 50  # for SHAP
CHECKPOINT_INTERVAL = 50  # save progress every 50 samples

print("=" * 70)
print("EXPERIMENT 2: XAI Analysis (1000 Samples)")
print("=" * 70)
print(f"Device: {DEVICE}")
print(f"Samples: {NUM_SAMPLES}")
print(f"Methods: Grad-CAM, Integrated Gradients, SHAP, LIME")
print(f"WARNING: This takes a long time (~1-2 hours)")
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
)

# select diverse samples (stratified by race for fair comparison)
print(f"Selecting {NUM_SAMPLES} diverse samples...")
df = pd.DataFrame([dataset.annotations.iloc[i] for i in range(len(dataset))])

selected_indices = []
samples_per_race = NUM_SAMPLES // 7  # 7 race categories

for race in df['race'].unique():
    race_samples = df[df['race'] == race].sample(
        n=min(samples_per_race, len(df[df['race'] == race])),
        random_state=42  # fixed seed for reproducibility
    )
    selected_indices.extend(race_samples.index.tolist())

# fill up to NUM_SAMPLES if needed
if len(selected_indices) < NUM_SAMPLES:
    remaining = NUM_SAMPLES - len(selected_indices)
    available = set(range(len(dataset))) - set(selected_indices)
    selected_indices.extend(np.random.choice(list(available), remaining, replace=False))

print(f"Selected {len(selected_indices)} samples")

# init XAI analyzer
print("Initializing XAI analyzer...")
xai = XAIAnalyzer(model=model, device=DEVICE, forensic_logger=logger)

# select background samples for SHAP (it needs reference samples)
print(f"Selecting {NUM_BACKGROUND} background samples for SHAP...")
background_indices = np.random.choice(len(dataset), NUM_BACKGROUND, replace=False)
background_images = []
for idx in background_indices:
    img, _ = dataset[idx]
    background_images.append(img)
background_tensor = torch.stack(background_images)

# checkpoint system (so we can resume if interrupted)
checkpoint_file = Path("../results/checkpoints/xai_progress.json")
checkpoint_file.parent.mkdir(parents=True, exist_ok=True)

processed_indices = set()
xai_results = []

# try to load previous progress
if checkpoint_file.exists():
    print("Found checkpoint, loading...")
    with open(checkpoint_file, 'r') as f:
        checkpoint = json.load(f)
    processed_indices = set(checkpoint['processed_indices'])
    xai_results = checkpoint['results']
    print(f"Resuming from {len(processed_indices)} already processed")

# create output directory
results_dir = Path("../results/xai_samples")
results_dir.mkdir(parents=True, exist_ok=True)

# process all samples
print("\nProcessing samples with all 4 XAI methods...")
print("This will take a while - grab a coffee!")
start_time = time.time()

for idx, data_idx in enumerate(tqdm(selected_indices, desc="XAI Analysis")):
    # skip if already processed
    if data_idx in processed_indices:
        continue
    
    try:
        image, metadata = dataset[data_idx]
        
        # create directory for this sample
        sample_dir = results_dir / f"sample_{len(xai_results):04d}"
        sample_dir.mkdir(parents=True, exist_ok=True)
        
        result = {
            'sample_idx': len(xai_results),
            'data_idx': int(data_idx),
            'race': metadata['race'],
            'gender': metadata['gender'],
            'age': metadata['age']
        }
        
        # method 1: Grad-CAM (usually fastest)
        try:
            gradcam_map = xai.gradcam(
                image,
                save_path=str(sample_dir / "gradcam.png")
            )
            result['gradcam_mean'] = float(gradcam_map.mean())
            result['gradcam_std'] = float(gradcam_map.std())
        except Exception as e:
            result['gradcam_error'] = str(e)
        
        # method 2: Integrated Gradients
        try:
            ig_map = xai.integrated_gradients(image, n_steps=50)
            # manually save since IG doesn't auto-save
            xai._save_heatmap_visualization(
                image, ig_map,
                str(sample_dir / "integrated_gradients.png"),
                title="Integrated Gradients"
            )
            result['ig_mean'] = float(ig_map.mean())
            result['ig_std'] = float(ig_map.std())
        except Exception as e:
            result['ig_error'] = str(e)
        
        # method 3: SHAP (can be slow)
        try:
            shap_map = xai.shap_analysis(
                image,
                background_tensor,
                save_path=str(sample_dir / "shap.png")
            )
            result['shap_available'] = True
        except Exception as e:
            result['shap_error'] = str(e)
            result['shap_available'] = False
        
        # method 4: LIME (also slow)
        try:
            lime_map = xai.lime_analysis(
                image,
                save_path=str(sample_dir / "lime.png")
            )
            result['lime_available'] = True
        except Exception as e:
            result['lime_error'] = str(e)
            result['lime_available'] = False
        
        xai_results.append(result)
        processed_indices.add(data_idx)
        
        # save checkpoint periodically
        if len(processed_indices) % CHECKPOINT_INTERVAL == 0:
            checkpoint_data = {
                'processed_indices': list(processed_indices),
                'results': xai_results,
                'last_updated': time.time()
            }
            with open(checkpoint_file, 'w') as f:
                json.dump(checkpoint_data, f)
            
            # estimate remaining time
            elapsed = time.time() - start_time
            avg_time = elapsed / len(processed_indices)
            remaining = (NUM_SAMPLES - len(processed_indices)) * avg_time
            print(f"\n  Progress: {len(processed_indices)}/{NUM_SAMPLES}")
            print(f"  Elapsed: {elapsed/3600:.1f} hours")
            print(f"  Estimated remaining: {remaining/3600:.1f} hours")
    
    except Exception as e:
        print(f"\nError on sample {data_idx}: {e}")
        continue

# create comparison visualization
print("\nGenerating comparison figure...")
df_results = pd.DataFrame(xai_results)

fig, axes = plt.subplots(2, 3, figsize=(20, 12))

# plot 1: Grad-CAM by race
ax1 = axes[0, 0]
race_gradcam = df_results.groupby('race')['gradcam_mean'].mean().sort_values()
ax1.barh(range(len(race_gradcam)), race_gradcam.values)
ax1.set_yticks(range(len(race_gradcam)))
ax1.set_yticklabels(race_gradcam.index)
ax1.set_xlabel('Mean Activation')
ax1.set_title('Grad-CAM by Race')
ax1.grid(axis='x', alpha=0.3)

# plot 2: Integrated Gradients by race
ax2 = axes[0, 1]
race_ig = df_results.groupby('race')['ig_mean'].mean().sort_values()
ax2.barh(range(len(race_ig)), race_ig.values)
ax2.set_yticks(range(len(race_ig)))
ax2.set_yticklabels(race_ig.index)
ax2.set_xlabel('Mean Attribution')
ax2.set_title('Integrated Gradients by Race')
ax2.grid(axis='x', alpha=0.3)

# plot 3: SHAP success rate
ax3 = axes[0, 2]
shap_success = df_results['shap_available'].sum()
shap_total = len(df_results)
ax3.bar(['Success', 'Failed'], [shap_success, shap_total - shap_success])
ax3.set_ylabel('Count')
ax3.set_title(f'SHAP Analysis ({shap_success}/{shap_total})')
ax3.grid(axis='y', alpha=0.3)

# plot 4: Grad-CAM by gender
ax4 = axes[1, 0]
gender_gradcam = df_results.groupby('gender')['gradcam_mean'].mean()
ax4.bar(range(len(gender_gradcam)), gender_gradcam.values)
ax4.set_xticks(range(len(gender_gradcam)))
ax4.set_xticklabels(gender_gradcam.index)
ax4.set_ylabel('Mean Activation')
ax4.set_title('Grad-CAM by Gender')
ax4.grid(axis='y', alpha=0.3)

# plot 5: LIME success rate
ax5 = axes[1, 1]
lime_success = df_results['lime_available'].sum()
lime_total = len(df_results)
ax5.bar(['Success', 'Failed'], [lime_success, lime_total - lime_success])
ax5.set_ylabel('Count')
ax5.set_title(f'LIME Analysis ({lime_success}/{lime_total})')
ax5.grid(axis='y', alpha=0.3)

# plot 6: sample distribution
ax6 = axes[1, 2]
race_counts = df_results['race'].value_counts()
ax6.bar(range(len(race_counts)), race_counts.values)
ax6.set_xticks(range(len(race_counts)))
ax6.set_xticklabels(race_counts.index, rotation=45, ha='right')
ax6.set_ylabel('Count')
ax6.set_title('Sample Distribution')
ax6.grid(axis='y', alpha=0.3)

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

# save detailed results
df_results.to_csv("../results/metrics/xai_results.csv", index=False)

# save summary
summary = {
    'experiment_id': EXPERIMENT_ID,
    'num_samples': len(processed_indices),
    'methods': ['gradcam', 'integrated_gradients', 'shap', 'lime'],
    'background_samples': NUM_BACKGROUND,
    'total_time_hours': (time.time() - start_time) / 3600,
    'samples_per_race': df_results['race'].value_counts().to_dict(),
    'shap_success_rate': float(df_results['shap_available'].sum() / len(df_results)),
    'lime_success_rate': float(df_results['lime_available'].sum() / len(df_results))
}

with open("../results/metrics/xai_analysis_summary.json", 'w') as f:
    json.dump(summary, f, indent=2)

print("Results saved")

# verify forensic chain
print("\nVerifying forensic chain...")
logger.verify_chain_integrity()
report = logger.generate_audit_report("../results/forensic_reports/xai_audit.json")

print("\n" + "=" * 70)
print("XAI ANALYSIS COMPLETED")
print("=" * 70)
print(f"Samples processed: {len(processed_indices)}")
print(f"Total time: {(time.time() - start_time)/3600:.1f} hours")
print(f"SHAP success: {summary['shap_success_rate']*100:.1f}%")
print(f"LIME success: {summary['lime_success_rate']*100:.1f}%")