#!/usr/bin/env python3
"""
Medical Imaging Dataset - Download Helper
==========================================

Project: Medical Imaging Dataset
Website: https://rskworld.in
Contact: help@rskworld.in, support@rskworld.in
Phone: +91 93305 39277
Founder: Molla Samser
Designer & Tester: Rima Khatun

Helper script to download and organize public medical imaging datasets.
"""

import os
import sys
import zipfile
import shutil
from pathlib import Path


def check_kaggle_installed():
    """Check if Kaggle API is installed."""
    try:
        import kaggle
        return True
    except ImportError:
        print("Kaggle API not found. Installing...")
        os.system("pip install kaggle")
        try:
            import kaggle
            return True
        except ImportError:
            print("Failed to install Kaggle API. Please install manually:")
            print("  pip install kaggle")
            return False


def download_kaggle_dataset(dataset_name, output_dir="data/temp"):
    """Download a dataset from Kaggle."""
    if not check_kaggle_installed():
        return False
    
    try:
        from kaggle.api.kaggle_api_extended import KaggleApi
        api = KaggleApi()
        api.authenticate()
        
        print(f"Downloading {dataset_name}...")
        api.dataset_download_files(dataset_name, path=output_dir, unzip=True)
        print(f"Downloaded to {output_dir}")
        return True
    except Exception as e:
        print(f"Error downloading dataset: {e}")
        print("Make sure you have:")
        print("1. Kaggle account")
        print("2. API token in ~/.kaggle/kaggle.json (or C:\\Users\\YourUsername\\.kaggle\\kaggle.json)")
        return False


def organize_chest_xray_data(source_dir, target_dir="data/xray"):
    """Organize chest X-ray data into project structure."""
    source = Path(source_dir)
    target_images = Path(target_dir) / "images"
    target_labels = Path(target_dir) / "labels"
    
    target_images.mkdir(parents=True, exist_ok=True)
    target_labels.mkdir(parents=True, exist_ok=True)
    
    # Find all image files
    image_extensions = ['.png', '.jpg', '.jpeg', '.dcm']
    image_count = 0
    
    for ext in image_extensions:
        for img_file in source.rglob(f'*{ext}'):
            if img_file.is_file():
                # Copy to images directory
                dest = target_images / img_file.name
                shutil.copy2(img_file, dest)
                image_count += 1
                
                # Create a basic label file
                label_file = target_labels / f"{img_file.stem}.json"
                if not label_file.exists():
                    label_data = {
                        "diagnosis": "Unknown",
                        "source": "Public Dataset",
                        "modality": "X-ray",
                        "body_part": "Chest",
                        "annotations": {
                            "findings": ["From public dataset"],
                            "recommendations": ["Educational use only"]
                        }
                    }
                    import json
                    with open(label_file, 'w') as f:
                        json.dump(label_data, f, indent=4)
    
    print(f"Organized {image_count} images to {target_dir}")
    return image_count


def main():
    """Main function."""
    print("=" * 60)
    print("Medical Imaging Dataset - Download Helper")
    print("=" * 60)
    print()
    print("Available datasets:")
    print("1. Chest X-Ray Images (Pneumonia)")
    print("2. Brain MRI Images for Brain Tumor Detection")
    print("3. Medical MNIST")
    print()
    
    choice = input("Enter dataset number (1-3) or 'q' to quit: ").strip()
    
    datasets = {
        "1": ("paultimothymooney/chest-xray-pneumonia", "xray"),
        "2": ("navoneel/brain-mri-images-for-brain-tumor-detection", "mri"),
        "3": ("andrewmvd/medical-mnist", "mixed")
    }
    
    if choice in datasets:
        dataset_name, data_type = datasets[choice]
        if download_kaggle_dataset(dataset_name):
            print("\nDownload complete!")
            print("Please organize the data manually or use the organize functions.")
    elif choice.lower() == 'q':
        print("Exiting...")
    else:
        print("Invalid choice.")


if __name__ == '__main__':
    main()