src/galaxy_data_loader.py

"""Load galaxy data from HuggingFace datasets and update profiles."""

import os
import logging
from typing import Dict, Any
import requests
from PIL import Image
import io

try:
    from datasets import load_dataset
    DATASETS_AVAILABLE = True
except ImportError:
    DATASETS_AVAILABLE = False
    logging.warning("datasets library not available, using fallback data")

logger = logging.getLogger(__name__)

def load_galaxy_data() -> Dict[str, Dict[str, Any]]:
    """Load galaxy data from HuggingFace datasets and return structured profiles."""
    
    if not DATASETS_AVAILABLE:
        logger.warning("Using fallback galaxy data")
        return get_fallback_galaxy_data()
    
    try:
        # Load both datasets
        euclid_dataset = load_dataset("mwalmsley/gz_euclid", "tiny", split="test")
        descriptions_dataset = load_dataset("Smith42/dating_pool_but_galaxies", split="train")
        
        # Convert to dictionaries for easier access
        euclid_data = {row['id_str']: row for row in euclid_dataset}
        descriptions_data = {row['id_str']: row for row in descriptions_dataset}
        
        galaxy_profiles = {}
        
        # Match on id_str and create profiles with numbered galaxy IDs
        matched_galaxies = []
        for galaxy_id_str in euclid_data.keys():
            if galaxy_id_str in descriptions_data:
                matched_galaxies.append(galaxy_id_str)
        
        # Sort for consistent ordering
        matched_galaxies.sort()
        
        # Create numbered galaxy profiles (galaxy_01, galaxy_02, etc.)
        for i, galaxy_id_str in enumerate(matched_galaxies, 1):
            numbered_id = f"galaxy_{i:02d}"
            euclid_row = euclid_data[galaxy_id_str]
            desc_row = descriptions_data[galaxy_id_str]
            
            # Extract name and description from caption
            caption = desc_row.get('caption', '')
            bio = ""
            description = ""
            name = f"Galaxy {i}"
            
            # Parse the caption to extract bio and name if available
            if caption:
                lines = caption.split('\\n')
                if lines and lines[0].startswith('Bio:'):
                    bio_part = lines[0].replace('Bio:', '').strip()
                    bio = bio_part[:100] + "..." if len(bio_part) > 100 else bio_part
                    description = bio_part
                else:
                    description = caption.strip()
                    bio = description[:100] + "..." if len(description) > 100 else description
            
            galaxy_profiles[numbered_id] = {
                "name": name,
                "bio": bio,
                "description": description,  
                "tags": ["Cosmic", "Mysterious"],
                "color": generate_color_from_id(galaxy_id_str),
                "id_str": galaxy_id_str,
                "image_data": euclid_row.get('image'),  # PIL Image object
                "euclid_features": {k: v for k, v in euclid_row.items() if k not in ['image', 'id_str']}
            }
        
        logger.info(f"Loaded {len(galaxy_profiles)} galaxy profiles from HuggingFace datasets")
        return galaxy_profiles
        
    except Exception as e:
        logger.error(f"Error loading from HuggingFace datasets: {e}")
        return get_fallback_galaxy_data()

def generate_color_from_id(id_str: str) -> str:
    """Generate a consistent color from galaxy ID string."""
    colors = [
        "#A688C9", "#D68B8B", "#E8A0D0", "#7BC9A0", "#E87D5A",
        "#6B8DD6", "#D4A76A", "#C9A688", "#B8A0C9", "#A0D4E8",
        "#E8D4A0", "#A0E8B8", "#E8A0A0", "#A0A0E8", "#E8E8A0",
        "#D0A0E8", "#A0E8D0", "#E8C4A0", "#C4A0E8", "#A0C4E8",
        "#E8A0C4", "#B4E8A0", "#A0E8E8"
    ]
    # Use hash to get consistent color index
    hash_val = sum(ord(c) for c in id_str)
    return colors[hash_val % len(colors)]

def get_fallback_galaxy_data() -> Dict[str, Dict[str, Any]]:
    """Fallback galaxy data using original profiles with numbered IDs."""
    fallback_profiles = {
        "galaxy_01": {
            "name": "Velvet Vortex",
            "bio": "Smooth operator with a soft spiral glow.",
            "description": "I keep my arms tight and my luminosity low-key. Looking for someone who appreciates subtlety over spectacle.",
            "tags": ["Smooth Spiral", "Low-Key", "Gentle Glow"],
            "color": "#A688C9",
            "image_url": None,
            "image_data": None
        },
        "galaxy_02": {
            "name": "Crimson Drift", 
            "bio": "Redshifted and proud.",
            "description": "I have been expanding away from everyone for billions of years and I am NOT slowing down. Commitment-phobic? Maybe. Mysterious? Definitely.",
            "tags": ["High Redshift", "Distant", "Loner Vibes"],
            "color": "#D68B8B",
            "image_url": None,
            "image_data": None
        }
        # Add more as needed...
    }
    
    # Generate numbered galaxy profiles if we have fewer than 23
    for i in range(1, 24):
        galaxy_id = f"galaxy_{i:02d}"
        if galaxy_id not in fallback_profiles:
            fallback_profiles[galaxy_id] = {
                "name": f"Galaxy {i}",
                "bio": "A mysterious galaxy in the cosmic dating scene.",
                "description": "This galaxy is looking for its perfect cosmic match.",
                "tags": ["Mysterious", "Cosmic"],
                "color": "#A688C9",
                "image_url": None,
                "image_data": None
            }
    
    return fallback_profiles

def download_galaxy_image(image_data, galaxy_id: str, images_dir: str = "images") -> bool:
    """Download and save a galaxy image from HuggingFace dataset."""
    if not image_data:
        return False
    
    try:
        # Ensure images directory exists
        os.makedirs(images_dir, exist_ok=True)
        
        # If image_data is a PIL Image (from datasets)
        if hasattr(image_data, 'save') and hasattr(image_data, 'mode'):
            image_path = os.path.join(images_dir, f"{galaxy_id}.jpg")
            # Convert to RGB if needed (in case it's RGBA or other format)
            if image_data.mode != 'RGB':
                image_data = image_data.convert('RGB')
            image_data.save(image_path, 'JPEG', quality=85)
            logger.info(f"Saved image for {galaxy_id}")
            return True
            
        # If image_data has a URL (fallback)
        elif isinstance(image_data, dict) and 'url' in image_data:
            response = requests.get(image_data['url'])
            if response.status_code == 200:
                image = Image.open(io.BytesIO(response.content))
                if image.mode != 'RGB':
                    image = image.convert('RGB')
                image_path = os.path.join(images_dir, f"{galaxy_id}.jpg")
                image.save(image_path, 'JPEG', quality=85)
                logger.info(f"Downloaded and saved image for {galaxy_id}")
                return True
                
        return False
        
    except Exception as e:
        logger.error(f"Error downloading image for {galaxy_id}: {e}")
        return False

def update_galaxy_images():
    """Download all galaxy images from the dataset."""
    galaxy_data = load_galaxy_data()
    
    for galaxy_id, profile in galaxy_data.items():
        if profile.get('image_data'):
            download_galaxy_image(profile['image_data'], galaxy_id)
        elif profile.get('image_url'):
            # Try to download from URL
            try:
                response = requests.get(profile['image_url'])
                if response.status_code == 200:
                    image = Image.open(io.BytesIO(response.content))
                    download_galaxy_image(image, galaxy_id)
            except Exception as e:
                logger.error(f"Error downloading from URL for {galaxy_id}: {e}")

if __name__ == "__main__":
    # Test the data loading
    data = load_galaxy_data()
    print(f"Loaded {len(data)} galaxies")
    for galaxy_id, profile in list(data.items())[:3]:
        print(f"{galaxy_id}: {profile['name']} - {profile['bio'][:50]}...")
    
    # Update images
    update_galaxy_images()