"""
Medical Imaging Dataset - Download All CT and MRI Datasets
===========================================================

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

This script attempts to download CT scan and MRI datasets from various sources.
"""

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


def check_kaggle_auth():
    """Check if Kaggle is authenticated."""
    try:
        result = subprocess.run(['kaggle', 'datasets', 'list', '--page-size', '1'], 
                              capture_output=True, text=True, timeout=10)
        return result.returncode == 0
    except:
        return False


def download_brain_mri():
    """Download Brain MRI dataset from Kaggle."""
    print("=" * 60)
    print("Downloading Brain MRI Images for Brain Tumor Detection")
    print("=" * 60)
    print()
    
    dataset = "navoneel/brain-mri-images-for-brain-tumor-detection"
    output_dir = Path("data/temp/mri")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        print(f"Downloading dataset: {dataset}")
        print("This may take a few minutes...")
        
        result = subprocess.run(
            ['kaggle', 'datasets', 'download', '-d', dataset, '-p', str(output_dir)],
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            print("Download successful!")
            
            # Find and extract zip file
            zip_files = list(output_dir.glob("*.zip"))
            if zip_files:
                zip_path = zip_files[0]
                print(f"Extracting {zip_path.name}...")
                
                extract_dir = output_dir / "extracted"
                extract_dir.mkdir(exist_ok=True)
                
                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                    zip_ref.extractall(extract_dir)
                
                # Organize images
                organize_mri_images(extract_dir)
                return True
            else:
                print("Warning: No zip file found after download")
                return False
        else:
            print(f"Error: {result.stderr}")
            if "401" in result.stderr or "authentication" in result.stderr.lower():
                print("\nKaggle authentication required!")
                print("Please:")
                print("1. Create a free account at https://www.kaggle.com")
                print("2. Go to Account Settings → API")
                print("3. Click 'Create New API Token'")
                print("4. Download kaggle.json")
                print(f"5. Place it in: {Path.home()}\\.kaggle\\kaggle.json")
            return False
            
    except FileNotFoundError:
        print("Error: Kaggle CLI not found. Please install: pip install kaggle")
        return False
    except Exception as e:
        print(f"Error downloading: {e}")
        return False


def organize_mri_images(source_dir):
    """Organize MRI images into the project structure."""
    print("Organizing MRI images...")
    
    source = Path(source_dir)
    target_images = Path("data/mri/images")
    target_labels = Path("data/mri/labels")
    
    target_images.mkdir(parents=True, exist_ok=True)
    target_labels.mkdir(parents=True, exist_ok=True)
    
    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():
                dest = target_images / img_file.name
                if not dest.exists():
                    try:
                        shutil.copy2(img_file, dest)
                        image_count += 1
                        
                        # Create label file
                        label_file = target_labels / f"{img_file.stem}.json"
                        if not label_file.exists():
                            label_data = {
                                "diagnosis": "Brain MRI",
                                "confidence": 0.90,
                                "date": "2024-01-15",
                                "modality": "MRI",
                                "body_part": "Brain",
                                "source": "Brain MRI Images for Brain Tumor Detection (Kaggle)",
                                "license": "CC BY 4.0",
                                "annotations": {
                                    "findings": ["From public Brain MRI dataset"],
                                    "recommendations": ["Educational use only"]
                                },
                                "metadata": {
                                    "patient_id": "ANONYMIZED",
                                    "age": "N/A",
                                    "gender": "N/A",
                                    "technique": "MRI"
                                }
                            }
                            with open(label_file, 'w') as f:
                                json.dump(label_data, f, indent=4)
                    except Exception as e:
                        print(f"Error copying {img_file}: {e}")
    
    print(f"Organized {image_count} MRI images to data/mri/images/")
    return image_count


def download_medical_mnist():
    """Download Medical MNIST dataset from Kaggle."""
    print("=" * 60)
    print("Downloading Medical MNIST Dataset")
    print("=" * 60)
    print()
    
    dataset = "andrewmvd/medical-mnist"
    output_dir = Path("data/temp/medmnist")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        print(f"Downloading dataset: {dataset}")
        print("This may take a few minutes...")
        
        result = subprocess.run(
            ['kaggle', 'datasets', 'download', '-d', dataset, '-p', str(output_dir)],
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            print("Download successful!")
            
            # Find and extract zip file
            zip_files = list(output_dir.glob("*.zip"))
            if zip_files:
                zip_path = zip_files[0]
                print(f"Extracting {zip_path.name}...")
                
                extract_dir = output_dir / "extracted"
                extract_dir.mkdir(exist_ok=True)
                
                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                    zip_ref.extractall(extract_dir)
                
                # Organize by modality
                organize_medmnist(extract_dir)
                return True
            else:
                print("Warning: No zip file found after download")
                return False
        else:
            print(f"Error: {result.stderr}")
            if "401" in result.stderr or "authentication" in result.stderr.lower():
                print("\nKaggle authentication required!")
                print_authentication_instructions()
            return False
            
    except Exception as e:
        print(f"Error downloading: {e}")
        return False


def organize_medmnist(source_dir):
    """Organize Medical MNIST by modality."""
    print("Organizing Medical MNIST by modality...")
    
    source = Path(source_dir)
    ct_target = Path("data/ct_scan/images")
    mri_target = Path("data/mri/images")
    
    ct_target.mkdir(parents=True, exist_ok=True)
    mri_target.mkdir(parents=True, exist_ok=True)
    
    image_extensions = ['.png', '.jpg', '.jpeg']
    ct_count = 0
    mri_count = 0
    
    # Medical MNIST structure: usually organized in folders
    for ext in image_extensions:
        for img_file in source.rglob(f'*{ext}'):
            if img_file.is_file():
                name_lower = img_file.name.lower()
                path_lower = str(img_file).lower()
                
                # Check folder names and file names for modality
                if 'ct' in path_lower or 'ctscan' in path_lower or 'ct_scan' in path_lower:
                    dest = ct_target / img_file.name
                    if not dest.exists():
                        try:
                            shutil.copy2(img_file, dest)
                            ct_count += 1
                        except:
                            pass
                elif 'mri' in path_lower or 'brain' in path_lower:
                    dest = mri_target / img_file.name
                    if not dest.exists():
                        try:
                            shutil.copy2(img_file, dest)
                            mri_count += 1
                        except:
                            pass
    
    print(f"Organized {ct_count} CT scan images")
    print(f"Organized {mri_count} MRI images")
    return ct_count + mri_count


def print_authentication_instructions():
    """Print instructions for Kaggle authentication."""
    print("\n" + "=" * 60)
    print("Kaggle Authentication Required")
    print("=" * 60)
    print()
    print("To download datasets from Kaggle, you need to:")
    print()
    print("1. Create a FREE Kaggle account:")
    print("   https://www.kaggle.com/account/register")
    print()
    print("2. Get your API credentials:")
    print("   - Go to: https://www.kaggle.com/account")
    print("   - Scroll to 'API' section")
    print("   - Click 'Create New API Token'")
    print("   - Download kaggle.json file")
    print()
    print("3. Place kaggle.json in:")
    print(f"   Windows: {Path.home()}\\.kaggle\\kaggle.json")
    print(f"   Linux/Mac: ~/.kaggle/kaggle.json")
    print()
    print("4. Set permissions (Linux/Mac only):")
    print("   chmod 600 ~/.kaggle/kaggle.json")
    print()
    print("5. Run this script again")
    print()
    print("=" * 60)


def main():
    """Main function to download all CT and MRI datasets."""
    print("=" * 60)
    print("Medical Imaging Dataset - Download All CT & MRI")
    print("=" * 60)
    print()
    
    # Check Kaggle authentication
    print("Checking Kaggle authentication...")
    if not check_kaggle_auth():
        print("Kaggle authentication not configured.")
        print_authentication_instructions()
        print()
        print("Attempting downloads anyway (may fail without credentials)...")
        print()
    
    downloaded = False
    
    # Try to download Brain MRI
    print("\n[1/2] Attempting to download Brain MRI dataset...")
    if download_brain_mri():
        downloaded = True
        print("[OK] Brain MRI dataset downloaded!")
    else:
        print("[SKIPPED] Brain MRI dataset - authentication may be required")
    
    # Try to download Medical MNIST
    print("\n[2/2] Attempting to download Medical MNIST dataset...")
    if download_medical_mnist():
        downloaded = True
        print("[OK] Medical MNIST dataset downloaded!")
    else:
        print("[SKIPPED] Medical MNIST dataset - authentication may be required")
    
    # Summary
    print()
    print("=" * 60)
    print("Download Summary")
    print("=" * 60)
    
    from scripts.load_data import MedicalImagingDataset
    dataset = MedicalImagingDataset(data_path='./data')
    info = dataset.get_dataset_info()
    
    print(f"Current dataset status:")
    print(f"  X-ray images: {info['xray_count']}")
    print(f"  CT scan images: {info['ct_count']}")
    print(f"  MRI images: {info['mri_count']}")
    print(f"  Total images: {info['total_images']}")
    print()
    
    if not downloaded:
        print("No datasets were downloaded.")
        print("This is likely because Kaggle authentication is required.")
        print()
        print("To download real CT and MRI images:")
        print("1. Follow the authentication instructions above")
        print("2. Or visit: https://medmnist.com/ (no account needed)")
        print("3. Or use other sources listed in CT_MRI_LOCATION.md")
    else:
        print("Datasets downloaded successfully!")
    
    print()
    print("=" * 60)


if __name__ == '__main__':
    main()