"""
Data Loader for FairFace Dataset
Handles loading and preprocessing of FairFace images with demographic labels
"""

import os
import pandas as pd
import numpy as np
from pathlib import Path
from typing import Tuple, Dict, List, Optional
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import hashlib
from src.forensic_logger import ForensicLogger


class FairFaceDataset(Dataset):
    """
    FairFace Dataset with demographic metadata
    
    Tracks race, gender, and age for bias analysis
    """
    
    def __init__(
        self, 
        data_dir: str,
        csv_file: str,
        transform: Optional[transforms.Compose] = None,
        forensic_logger: Optional[ForensicLogger] = None
    ):
        """Initialize FairFace dataset"""
        self.data_dir = Path(data_dir)
        self.transform = transform or self._get_default_transform()
        self.logger = forensic_logger
        
        # load CSV annotations
        self.annotations = pd.read_csv(csv_file)
        
        # log data loading for forensics
        if self.logger:
            csv_hash = self.logger.compute_file_hash(csv_file)
            self.logger.log_data_operation(
                operation="LOAD_ANNOTATIONS",
                dataset="FairFace",
                num_samples=len(self.annotations),
                file_hash=csv_hash
            )
        
        # get unique categories from the data
        self.race_categories = sorted(self.annotations['race'].unique().tolist())
        self.gender_categories = sorted(self.annotations['gender'].unique().tolist())
        self.age_categories = sorted(self.annotations['age'].unique().tolist())
        
        print(f"Dataset loaded: {len(self.annotations)} samples")
        print(f"  Race categories: {len(self.race_categories)}")
        print(f"  Gender categories: {len(self.gender_categories)}")
        print(f"  Age groups: {len(self.age_categories)}")
    
    def _get_default_transform(self):
        """Default preprocessing - resize to 160x160 and normalize for FaceNet"""
        return transforms.Compose([
            transforms.Resize((160, 160)),
            transforms.ToTensor(),
            # normalize to [-1, 1] range (FaceNet expects this)
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        ])
    
    def __len__(self):
        return len(self.annotations)
    
    def __getitem__(self, idx):
        """Get image and metadata for a sample"""
        row = self.annotations.iloc[idx]
        
        # load and preprocess image
        img_path = self.data_dir / row['file']
        image = Image.open(img_path).convert('RGB')
        
        # compute hash for forensic tracking
        img_hash = hashlib.md5(image.tobytes()).hexdigest()
        
        if self.transform:
            image = self.transform(image)
        
        # package metadata
        metadata = {
            'file': row['file'],
            'race': row['race'],
            'gender': row['gender'],
            'age': row['age'],
            'image_hash': img_hash,
            'idx': idx
        }
        
        return image, metadata
    
    def get_demographic_distribution(self):
        """Calculate demographic distribution statistics"""
        distribution = {
            'race': self.annotations['race'].value_counts().to_dict(),
            'gender': self.annotations['gender'].value_counts().to_dict(),
            'age': self.annotations['age'].value_counts().to_dict(),
            'total_samples': len(self.annotations)
        }
        
        # also compute intersectional stats (race x gender)
        # this is useful for bias analysis
        intersectional = self.annotations.groupby(['race', 'gender']).size().to_dict()
        distribution['race_gender'] = {
            f"{race}_{gender}": count 
            for (race, gender), count in intersectional.items()
        }
        
        return distribution
    
    def get_balanced_subset(
        self, 
        samples_per_group: int = 100,
        stratify_by: List[str] = ['race', 'gender']
    ):
        """
        Create a balanced subset for testing
        Takes equal number of samples from each demographic group
        """
        balanced_dfs = []
        
        # group by demographics
        grouped = self.annotations.groupby(stratify_by)
        
        for group_keys, group_df in grouped:
            # sample from each group (with fixed seed for reproducibility)
            n_samples = min(samples_per_group, len(group_df))
            sampled = group_df.sample(n=n_samples, random_state=42)
            balanced_dfs.append(sampled)
        
        # concatenate all groups
        balanced_annotations = pd.concat(balanced_dfs, ignore_index=True)
        
        # create new dataset instance with balanced data
        balanced_dataset = FairFaceDataset.__new__(FairFaceDataset)
        balanced_dataset.data_dir = self.data_dir
        balanced_dataset.transform = self.transform
        balanced_dataset.logger = self.logger
        balanced_dataset.annotations = balanced_annotations
        balanced_dataset.race_categories = self.race_categories
        balanced_dataset.gender_categories = self.gender_categories
        balanced_dataset.age_categories = self.age_categories
        
        print(f"Balanced subset created: {len(balanced_annotations)} samples")
        
        if self.logger:
            self.logger.log_data_operation(
                operation="CREATE_BALANCED_SUBSET",
                dataset="FairFace_Balanced",
                num_samples=len(balanced_annotations)
            )
        
        return balanced_dataset


# quick test
if __name__ == "__main__":
    print("Testing FairFace Data Loader...")
    
    # check if dataset exists
    data_dir = Path("data/raw/fairface")
    csv_file = data_dir / "fairface_label_val.csv"
    
    if not csv_file.exists():
        print("\nDataset not found at:", csv_file)
        print("Please download FairFace dataset first.")
    else:
        # init logger
        logger = ForensicLogger(experiment_id="TEST_DATALOADER")
        
        # load dataset
        dataset = FairFaceDataset(
            data_dir=str(data_dir / "val"),
            csv_file=str(csv_file),
            forensic_logger=logger
        )
        
        # test loading one sample
        img, metadata = dataset[0]
        print(f"\nSample loaded:")
        print(f"  Image shape: {img.shape}")
        print(f"  Race: {metadata['race']}")
        print(f"  Gender: {metadata['gender']}")
        
        # check demographic distribution
        distribution = dataset.get_demographic_distribution()
        print(f"\nDemographic Distribution:")
        print(f"  Total samples: {distribution['total_samples']}")
        
        logger.verify_chain_integrity()
        print("\nTest completed successfully")