"""
Experiment 1: Baseline Evaluation
Characterize FaceNet performance on FairFace validation set
"""

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

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 = "BASELINE_EVAL_FULL"
NUM_SAMPLES = None  # None = use all samples

print("=" * 70)
print("EXPERIMENT 1: Baseline Evaluation (FULL DATASET)")
print("=" * 70)
print(f"Device: {DEVICE}")
print(f"Experiment ID: {EXPERIMENT_ID}")
print()

# init forensic 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 how many samples to use
if NUM_SAMPLES is None:
    NUM_SAMPLES = len(dataset)
    print(f"Using ALL {NUM_SAMPLES} samples from validation set")
    indices = np.arange(len(dataset))
else:
    print(f"Using random subset of {NUM_SAMPLES} samples")
    indices = np.random.choice(len(dataset), NUM_SAMPLES, replace=False)

# extract embeddings for all samples
print("Extracting embeddings...")
embeddings = []
metadata_list = []

for idx in tqdm(indices):
    image, metadata = dataset[idx]
    
    # get 512-dim embedding
    emb = model.get_embedding(image, log_inference=False)  # don't log every inference
    
    embeddings.append(emb)
    metadata_list.append(metadata)

embeddings = np.array(embeddings)
print(f"Embeddings shape: {embeddings.shape}")

# basic embedding statistics
print("\nEmbedding Statistics:")
norms = np.linalg.norm(embeddings, axis=1)
print(f"  Mean norm: {norms.mean():.4f}")
print(f"  Std norm: {norms.std():.4f}")
print(f"  Min value: {embeddings.min():.4f}")
print(f"  Max value: {embeddings.max():.4f}")

# check demographic distribution
df_metadata = pd.DataFrame(metadata_list)
print("\nDemographic Distribution:")
print(df_metadata['race'].value_counts())
print()
print(df_metadata['gender'].value_counts())

# compute pairwise distances
# this is expensive: O(n^2) complexity
# for 10k samples = 50M distances!
print("\nComputing pairwise distances...")
n_pairs = NUM_SAMPLES * (NUM_SAMPLES - 1) // 2
print(f"  Computing {n_pairs:,} pairwise distances")
print(f"  This takes ~2-5 hours on GPU, please be patient...")

distances = []

# tried broadcasting but memory issues with 10k samples
# this loop is slower but works
for i in tqdm(range(len(embeddings)), desc="Computing distances"):
    for j in range(i + 1, len(embeddings)):
        dist = np.linalg.norm(embeddings[i] - embeddings[j])
        distances.append(dist)

print(f"\nDistance computation complete:")
print(f"  Total distances: {len(distances):,}")
print(f"  Mean distance: {np.mean(distances):.4f}")
print(f"  Std distance: {np.std(distances):.4f}")
# print(f"  Min distance: {np.min(distances):.4f}")  # debug
# print(f"  Max distance: {np.max(distances):.4f}")

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

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

# plot 1: embedding norms
axes[0, 0].hist(norms, bins=50, edgecolor='black')
axes[0, 0].set_title('Distribution of Embedding Norms')
axes[0, 0].set_xlabel('L2 Norm')
axes[0, 0].set_ylabel('Frequency')

# plot 2: distance distribution (sampled to avoid memory issues)
# sample every 100th distance for visualization
axes[0, 1].hist(distances[::100], bins=50, edgecolor='black', color='orange')
axes[0, 1].set_title('Pairwise Distance Distribution')
axes[0, 1].set_xlabel('Euclidean Distance')
axes[0, 1].set_ylabel('Frequency')

# plot 3: race distribution
race_counts = df_metadata['race'].value_counts()
axes[1, 0].bar(range(len(race_counts)), race_counts.values)
axes[1, 0].set_xticks(range(len(race_counts)))
axes[1, 0].set_xticklabels(race_counts.index, rotation=45, ha='right')
axes[1, 0].set_title('Race Distribution')
axes[1, 0].set_ylabel('Count')

# plot 4: gender distribution
gender_counts = df_metadata['gender'].value_counts()
axes[1, 1].bar(range(len(gender_counts)), gender_counts.values, color=['blue', 'pink'])
axes[1, 1].set_xticks(range(len(gender_counts)))
axes[1, 1].set_xticklabels(gender_counts.index)
axes[1, 1].set_title('Gender Distribution')
axes[1, 1].set_ylabel('Count')

plt.tight_layout()

# save figure
output_path = Path("../results/figures/baseline_evaluation.png")
output_path.parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches='tight')
print(f"Visualization saved to {output_path}")

# save results to JSON
results = {
    'experiment_id': EXPERIMENT_ID,
    'num_samples': int(NUM_SAMPLES),
    'embedding_dim': int(embeddings.shape[1]),
    'mean_norm': float(norms.mean()),
    'std_norm': float(norms.std()),
    'mean_distance': float(np.mean(distances)),
    'std_distance': float(np.std(distances)),
    'total_distances_computed': len(distances),
    'demographic_distribution': {
        'race': {k: int(v) for k, v in df_metadata['race'].value_counts().to_dict().items()},
        'gender': {k: int(v) for k, v in df_metadata['gender'].value_counts().to_dict().items()}
    }
}

import json
results_path = Path("../results/metrics/baseline_results.json")
results_path.parent.mkdir(parents=True, exist_ok=True)
with open(results_path, 'w') as f:
    json.dump(results, f, indent=2)
print(f"Results saved to {results_path}")

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

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

print("\n" + "=" * 70)
print("EXPERIMENT COMPLETED")
print("=" * 70)
print(f"\nKey Results:")
print(f"  Samples analyzed: {NUM_SAMPLES:,}")
print(f"  Distances computed: {len(distances):,}")
print(f"  Mean distance: {np.mean(distances):.4f} ± {np.std(distances):.4f}")