"""
Experiment 3: Demographic Bias Analysis
Evaluate model performance across race, gender, and age groups
Full validation set (10,954 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
import seaborn as sns
from tqdm import tqdm
from scipy import stats

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

# config
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
EXPERIMENT_ID = "BIAS_ANALYSIS_FULL"
NUM_SAMPLES = None  # None = use all samples
VERIFICATION_THRESHOLD = 1.0

print("=" * 70)
print("EXPERIMENT 3: Demographic Bias Analysis (FULL DATASET)")
print("=" * 70)
print(f"Device: {DEVICE}")
print(f"Experiment ID: {EXPERIMENT_ID}")
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
)

# decide whether to use all samples or balanced subset
if NUM_SAMPLES is None:
    print("Using full dataset (no balancing)")
    balanced_dataset = dataset
    NUM_SAMPLES = len(dataset)
else:
    print(f"Creating balanced subset with {NUM_SAMPLES} samples per group")
    balanced_dataset = dataset.get_balanced_subset(
        samples_per_group=NUM_SAMPLES,
        stratify_by=['race', 'gender']
    )

print(f"Total samples: {len(balanced_dataset)}")

# extract embeddings
print("\nExtracting embeddings...")
embeddings_list = []
metadata_list = []

for idx in tqdm(range(len(balanced_dataset))):
    image, metadata = balanced_dataset[idx]
    
    emb = model.get_embedding(image, log_inference=False)
    
    embeddings_list.append(emb)
    metadata_list.append(metadata)

embeddings = np.array(embeddings_list)
df_metadata = pd.DataFrame(metadata_list)

print(f"Embeddings extracted: {len(embeddings)}")
print(f"Shape: {embeddings.shape}")

# compute demographic-specific distances
print("\nComputing intra-group and inter-group distances...")
print("WARNING: This is computationally expensive for full dataset (~1-2 hours)")

def compute_group_distances(df, embeddings, group_col):
    """
    Compute distances within and between demographic groups
    Intra-group = distance between members of same group
    Inter-group = distance between members of different groups
    """
    results = {}
    
    for group in df[group_col].unique():
        group_mask = df[group_col] == group
        group_embs = embeddings[group_mask]
        other_embs = embeddings[~group_mask]
        
        print(f"  Processing {group}: {len(group_embs)} samples")
        
        # intra-group: all pairs within group
        # this is O(n^2) so can be slow for large groups
        print(f"    Computing intra-group distances...")
        intra_distances = []
        for i in tqdm(range(len(group_embs)), desc=f"    {group} intra", leave=False):
            for j in range(i + 1, len(group_embs)):
                dist = np.linalg.norm(group_embs[i] - group_embs[j])
                intra_distances.append(dist)
        
        # inter-group: sample to keep it manageable
        # full computation would be too expensive
        print(f"    Computing inter-group distances (sampled)...")
        inter_distances = []
        sample_size_group = min(500, len(group_embs))
        sample_size_other = min(500, len(other_embs))
        
        for i in range(sample_size_group):
            for j in range(sample_size_other):
                dist = np.linalg.norm(group_embs[i] - other_embs[j])
                inter_distances.append(dist)
        
        print(f"    Computed {len(intra_distances)} intra, {len(inter_distances)} inter distances")
        
        results[group] = {
            'intra_mean': np.mean(intra_distances) if intra_distances else 0,
            'intra_std': np.std(intra_distances) if intra_distances else 0,
            'inter_mean': np.mean(inter_distances) if inter_distances else 0,
            'inter_std': np.std(inter_distances) if inter_distances else 0,
            'count': int(group_mask.sum()),
            'intra_pairs': len(intra_distances),
            'inter_pairs': len(inter_distances)
        }
    
    return results

# analyze by race
print("\nAnalyzing by race...")
race_results = compute_group_distances(df_metadata, embeddings, 'race')

# analyze by gender
print("\nAnalyzing by gender...")
gender_results = compute_group_distances(df_metadata, embeddings, 'gender')

# print results
print("\nRace-based Analysis:")
print("-" * 70)
df_race = pd.DataFrame(race_results).T
df_race = df_race.round(4)
print(df_race)

print("\nGender-based Analysis:")
print("-" * 70)
df_gender = pd.DataFrame(gender_results).T
df_gender = df_gender.round(4)
print(df_gender)

# statistical significance testing
print("\nStatistical Significance Tests:")
print("-" * 70)

# compare first two races as example
race_groups = df_metadata['race'].unique()
if len(race_groups) >= 2:
    race1_mask = df_metadata['race'] == race_groups[0]
    race2_mask = df_metadata['race'] == race_groups[1]
    
    # sample embeddings to make t-test tractable
    race1_embs = embeddings[race1_mask][:100]
    race2_embs = embeddings[race2_mask][:100]
    
    # compute pairwise distances for each race
    race1_dists = [np.linalg.norm(race1_embs[i] - race1_embs[j]) 
                   for i in range(len(race1_embs)) 
                   for j in range(i+1, len(race1_embs))]
    race2_dists = [np.linalg.norm(race2_embs[i] - race2_embs[j]) 
                   for i in range(len(race2_embs)) 
                   for j in range(i+1, len(race2_embs))]
    
    # t-test to check if distributions differ significantly
    t_stat, p_value = stats.ttest_ind(race1_dists, race2_dists)
    print(f"T-test ({race_groups[0]} vs {race_groups[1]}):")
    print(f"  t-statistic: {t_stat:.4f}")
    print(f"  p-value: {p_value:.4f}")
    print(f"  Significant at p<0.05: {'Yes' if p_value < 0.05 else 'No'}")

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

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

# plot 1: intra-group distances by race
ax1 = fig.add_subplot(gs[0, :2])
races = list(race_results.keys())
intra_means = [race_results[r]['intra_mean'] for r in races]
intra_stds = [race_results[r]['intra_std'] for r in races]

ax1.bar(range(len(races)), intra_means, yerr=intra_stds, capsize=5)
ax1.set_xticks(range(len(races)))
ax1.set_xticklabels(races, rotation=45, ha='right')
ax1.set_ylabel('Mean Distance')
ax1.set_title('Intra-Group Distances by Race')
ax1.grid(axis='y', alpha=0.3)

# plot 2: intra vs inter comparison
ax2 = fig.add_subplot(gs[0, 2])
x = np.arange(len(races))
width = 0.35
intra_vals = [race_results[r]['intra_mean'] for r in races]
inter_vals = [race_results[r]['inter_mean'] for r in races]

ax2.bar(x - width/2, intra_vals, width, label='Intra-group')
ax2.bar(x + width/2, inter_vals, width, label='Inter-group')
ax2.set_xticks(x)
ax2.set_xticklabels([''] * len(races))  # hide labels (too crowded)
ax2.set_ylabel('Mean Distance')
ax2.set_title('Intra vs Inter Distances')
ax2.legend()
ax2.grid(axis='y', alpha=0.3)

# plot 3: gender comparison
ax3 = fig.add_subplot(gs[1, 0])
genders = list(gender_results.keys())
gender_intra = [gender_results[g]['intra_mean'] for g in genders]
gender_inter = [gender_results[g]['inter_mean'] for g in genders]

x_gender = np.arange(len(genders))
ax3.bar(x_gender - 0.2, gender_intra, 0.4, label='Intra-group')
ax3.bar(x_gender + 0.2, gender_inter, 0.4, label='Inter-group')
ax3.set_xticks(x_gender)
ax3.set_xticklabels(genders)
ax3.set_ylabel('Mean Distance')
ax3.set_title('Gender-based Distances')
ax3.legend()
ax3.grid(axis='y', alpha=0.3)

# plot 4: sample counts
ax4 = fig.add_subplot(gs[1, 1])
counts = [race_results[r]['count'] for r in races]
ax4.bar(range(len(races)), counts)
ax4.set_xticks(range(len(races)))
ax4.set_xticklabels(races, rotation=45, ha='right')
ax4.set_ylabel('Count')
ax4.set_title('Sample Distribution by Race')
ax4.grid(axis='y', alpha=0.3)

# plot 5: heatmap of intersectional distribution
ax5 = fig.add_subplot(gs[1, 2])
pivot_table = df_metadata.groupby(['race', 'gender']).size().unstack(fill_value=0)
sns.heatmap(pivot_table, annot=True, fmt='d', cmap='YlOrRd', ax=ax5)
ax5.set_title('Distribution: Race x Gender')

# plot 6: overall distance distribution
ax6 = fig.add_subplot(gs[2, :])
# sample to avoid memory issues
all_distances = []
sample_size = min(500, len(embeddings))
for i in range(0, sample_size, 10):
    for j in range(i+1, sample_size, 10):
        dist = np.linalg.norm(embeddings[i] - embeddings[j])
        all_distances.append(dist)

ax6.hist(all_distances, bins=50, edgecolor='black', alpha=0.7)
ax6.axvline(VERIFICATION_THRESHOLD, color='r', linestyle='--', 
            label=f'Threshold={VERIFICATION_THRESHOLD}')
ax6.set_xlabel('Euclidean Distance')
ax6.set_ylabel('Frequency')
ax6.set_title('Overall Distance Distribution (sampled)')
ax6.legend()
ax6.grid(axis='y', alpha=0.3)

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

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

# save results as JSON
# need to convert numpy types to Python native types for JSON
def convert_numpy(obj):
    if isinstance(obj, dict):
        return {k: convert_numpy(v) for k, v in obj.items()}
    elif isinstance(obj, (np.integer, np.floating)):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    return obj

results_dict = {
    'race_analysis': convert_numpy(race_results),
    'gender_analysis': convert_numpy(gender_results),
    'total_samples_analyzed': int(len(embeddings)),
    'threshold': float(VERIFICATION_THRESHOLD),
    'max_intragroup_diff': float(max([r['intra_mean'] for r in race_results.values()]) - 
                                  min([r['intra_mean'] for r in race_results.values()]))
}

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

# also save as CSV for easy analysis in Excel
df_race.to_csv("../results/metrics/bias_by_race.csv")
df_gender.to_csv("../results/metrics/bias_by_gender.csv")
print("CSV results saved")

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

# generate audit report
report = logger.generate_audit_report("../results/forensic_reports/bias_audit.json")
print(f"Audit report: {report['total_events']} events logged")

print("\n" + "=" * 70)
print("BIAS ANALYSIS COMPLETED")
print("=" * 70)
print("\nKey Findings:")
print(f"  Samples analyzed: {len(embeddings)}")
print(f"  Race groups: {len(race_results)}")
print(f"  Gender groups: {len(gender_results)}")
print(f"  Max intra-group bias: {results_dict['max_intragroup_diff']:.4f}")