"""
Medical Imaging Dataset - Sample Data Generator
================================================

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 generates sample placeholder medical images for demonstration purposes.
"""

import numpy as np
from PIL import Image
import json
from pathlib import Path
import os


def create_sample_xray_image(output_path, image_num=1):
    """Create a sample X-ray placeholder image with varied patterns."""
    width, height = 512, 512
    np.random.seed(image_num * 42)  # Different seed for each image
    
    # Base intensity varies by image
    base_intensity = 80 + (image_num * 20) % 100
    image = np.full((height, width), base_intensity, dtype=np.uint8)
    
    # Different patterns for each image
    if image_num == 1:
        # Chest X-ray pattern - vertical ribs
        for y in range(height):
            for x in range(width):
                # Rib-like structures
                if (x % 60) < 8:
                    image[y, x] = min(255, image[y, x] + 60)
                # Spine in center
                if abs(x - width//2) < 15:
                    image[y, x] = min(255, image[y, x] + 40)
    
    elif image_num == 2:
        # Extremity X-ray - bone structure
        center_x, center_y = width//2, height//2
        for y in range(height):
            for x in range(width):
                dist = abs(x - center_x) + abs(y - center_y)
                if dist < 100:
                    image[y, x] = min(255, image[y, x] + 70)
                # Joint area
                if abs(y - center_y) < 30:
                    image[y, x] = min(255, image[y, x] + 50)
    
    elif image_num == 3:
        # Abdomen X-ray - horizontal structures
        for y in range(height):
            for x in range(width):
                # Horizontal bands
                if (y % 80) < 12:
                    image[y, x] = min(255, image[y, x] + 55)
                # Central structure
                if (x - width//2)**2 + (y - height//2)**2 < 15000:
                    image[y, x] = min(255, image[y, x] + 30)
    
    elif image_num == 4:
        # Skull X-ray - circular pattern
        center_x, center_y = width//2, height//2
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                if 80 < dist < 200:
                    image[y, x] = min(255, image[y, x] + 65)
                if dist < 80:
                    image[y, x] = max(0, image[y, x] - 30)
    
    else:  # image_num == 5
        # Pelvis X-ray - complex structure
        for y in range(height):
            for x in range(width):
                # Hip bones
                if (x - width//3)**2 + (y - height//2)**2 < 8000:
                    image[y, x] = min(255, image[y, x] + 50)
                if (x - 2*width//3)**2 + (y - height//2)**2 < 8000:
                    image[y, x] = min(255, image[y, x] + 50)
                # Spine
                if abs(x - width//2) < 10:
                    image[y, x] = min(255, image[y, x] + 40)
    
    # Add varied noise
    noise_level = 15 + (image_num * 3) % 15
    noise = np.random.normal(0, noise_level, (height, width))
    image = np.clip(image.astype(float) + noise, 0, 255).astype(np.uint8)
    
    # Apply different contrast
    contrast = 1.0 + (image_num * 0.15) % 0.5
    image = np.clip((image - 128) * contrast + 128, 0, 255).astype(np.uint8)
    
    img = Image.fromarray(image, mode='L')
    img.save(output_path)
    print(f"Created X-ray sample {image_num}: {output_path}")


def create_sample_ct_image(output_path, image_num=1):
    """Create a sample CT scan placeholder image with varied cross-sections."""
    width, height = 512, 512
    np.random.seed(image_num * 73)  # Different seed
    
    base_intensity = 100 + (image_num * 15) % 80
    image = np.full((height, width), base_intensity, dtype=np.uint8)
    center_x, center_y = width // 2, height // 2
    
    if image_num == 1:
        # Head CT - brain cross-section
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                if dist < 180:
                    # Skull
                    if dist > 160:
                        image[y, x] = min(255, image[y, x] + 80)
                    # Brain matter
                    elif dist < 140:
                        image[y, x] = max(0, image[y, x] - 40)
                    # Gray matter
                    else:
                        image[y, x] = min(255, image[y, x] + 20)
    
    elif image_num == 2:
        # Chest CT - lung cross-section
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Ribs
                if 150 < dist < 200:
                    image[y, x] = min(255, image[y, x] + 70)
                # Lungs (darker)
                elif dist < 140:
                    if abs(x - center_x) > 30:  # Left and right lungs
                        image[y, x] = max(0, image[y, x] - 50)
                # Heart area
                if (x - center_x)**2 + (y - center_y)**2 < 3000:
                    image[y, x] = min(255, image[y, x] + 30)
    
    elif image_num == 3:
        # Abdomen CT - organ cross-section
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Body outline
                if dist > 170:
                    image[y, x] = min(255, image[y, x] + 60)
                # Organs (varied)
                elif dist < 150:
                    # Liver (right side)
                    if x > center_x and abs(y - center_y) < 80:
                        image[y, x] = min(255, image[y, x] + 25)
                    # Stomach (left side)
                    elif x < center_x and abs(y - center_y) < 60:
                        image[y, x] = max(0, image[y, x] - 20)
    
    elif image_num == 4:
        # Pelvis CT - bone structure
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Pelvic bones
                if 120 < dist < 190:
                    image[y, x] = min(255, image[y, x] + 75)
                # Soft tissue
                elif dist < 120:
                    image[y, x] = max(0, image[y, x] - 30)
                # Spine
                if abs(x - center_x) < 8 and y > center_y - 50:
                    image[y, x] = min(255, image[y, x] + 60)
    
    else:  # image_num == 5
        # Spine CT - vertebral cross-section
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Vertebra (center)
                if dist < 60:
                    image[y, x] = min(255, image[y, x] + 80)
                # Spinal canal (darker center)
                if dist < 25:
                    image[y, x] = max(0, image[y, x] - 40)
                # Surrounding tissue
                elif dist < 150:
                    image[y, x] = max(0, image[y, x] - 15)
    
    # Add texture
    noise_level = 12 + (image_num * 2) % 10
    noise = np.random.normal(0, noise_level, (height, width))
    image = np.clip(image.astype(float) + noise, 0, 255).astype(np.uint8)
    
    # Vary brightness
    brightness_shift = (image_num * 10) % 30 - 15
    image = np.clip(image.astype(float) + brightness_shift, 0, 255).astype(np.uint8)
    
    img = Image.fromarray(image, mode='L')
    img.save(output_path)
    print(f"Created CT scan sample {image_num}: {output_path}")


def create_sample_mri_image(output_path, image_num=1):
    """Create a sample MRI placeholder image with varied brain structures."""
    width, height = 512, 512
    np.random.seed(image_num * 97)  # Different seed
    
    base_intensity = 60 + (image_num * 12) % 60
    image = np.full((height, width), base_intensity, dtype=np.uint8)
    center_x, center_y = width // 2, height // 2
    
    if image_num == 1:
        # Brain MRI - T1 weighted (sagittal)
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Skull
                if dist > 190:
                    image[y, x] = min(255, image[y, x] + 90)
                # White matter (bright)
                elif 100 < dist < 160:
                    image[y, x] = min(255, image[y, x] + 50)
                # Gray matter (medium)
                elif 60 < dist < 100:
                    image[y, x] = min(255, image[y, x] + 20)
                # CSF (dark)
                elif dist < 60:
                    image[y, x] = max(0, image[y, x] - 30)
    
    elif image_num == 2:
        # Brain MRI - T2 weighted (axial)
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Skull
                if dist > 185:
                    image[y, x] = min(255, image[y, x] + 85)
                # CSF (bright in T2)
                elif dist < 70:
                    image[y, x] = min(255, image[y, x] + 60)
                # Gray matter
                elif 70 < dist < 120:
                    image[y, x] = min(255, image[y, x] + 30)
                # White matter (darker)
                elif 120 < dist < 160:
                    image[y, x] = max(0, image[y, x] - 20)
    
    elif image_num == 3:
        # Brain MRI - coronal view
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                # Skull outline
                if 170 < dist < 195:
                    image[y, x] = min(255, image[y, x] + 80)
                # Brain tissue (varied)
                elif dist < 170:
                    # Left hemisphere
                    if x < center_x:
                        image[y, x] = min(255, image[y, x] + 25)
                    # Right hemisphere
                    else:
                        image[y, x] = min(255, image[y, x] + 35)
                    # Ventricles (center, darker)
                    if dist < 50:
                        image[y, x] = max(0, image[y, x] - 40)
    
    elif image_num == 4:
        # Spine MRI - sagittal
        for y in range(height):
            for x in range(width):
                # Vertebrae (bright)
                if abs(x - center_x) < 25:
                    if (y % 80) < 50:
                        image[y, x] = min(255, image[y, x] + 70)
                # Spinal cord (darker)
                if abs(x - center_x) < 8:
                    image[y, x] = max(0, image[y, x] - 30)
                # Discs (medium)
                if abs(x - center_x) < 30:
                    if (y % 80) > 50:
                        image[y, x] = min(255, image[y, x] + 40)
    
    else:  # image_num == 5
        # Knee MRI - sagittal
        for y in range(height):
            for x in range(width):
                # Femur (top, bright)
                if y < height//3:
                    if abs(x - center_x) < 40:
                        image[y, x] = min(255, image[y, x] + 75)
                # Tibia (bottom, bright)
                elif y > 2*height//3:
                    if abs(x - center_x) < 40:
                        image[y, x] = min(255, image[y, x] + 75)
                # Cartilage (middle, medium)
                elif height//3 < y < 2*height//3:
                    if abs(x - center_x) < 50:
                        image[y, x] = min(255, image[y, x] + 45)
                # Soft tissue (darker)
                else:
                    image[y, x] = max(0, image[y, x] - 20)
    
    # Add fine texture
    noise_level = 8 + (image_num * 1.5) % 8
    noise = np.random.normal(0, noise_level, (height, width))
    image = np.clip(image.astype(float) + noise, 0, 255).astype(np.uint8)
    
    # Vary contrast
    contrast = 0.9 + (image_num * 0.1) % 0.4
    image = np.clip((image - 128) * contrast + 128, 0, 255).astype(np.uint8)
    
    img = Image.fromarray(image, mode='L')
    img.save(output_path)
    print(f"Created MRI sample {image_num}: {output_path}")


def create_label_file(label_path, image_type, image_num, diagnosis="Normal"):
    """Create a sample label JSON file."""
    diagnoses = ["Normal", "Abnormal", "Pneumonia", "Fracture", "Tumor"]
    body_parts = {
        'xray': ['Chest', 'Abdomen', 'Extremity'],
        'ct_scan': ['Head', 'Chest', 'Abdomen', 'Pelvis'],
        'mri': ['Brain', 'Spine', 'Knee', 'Shoulder']
    }
    
    label_data = {
        "diagnosis": diagnosis if diagnosis in diagnoses else "Normal",
        "confidence": round(np.random.uniform(0.85, 0.98), 2),
        "date": "2024-01-15",
        "modality": image_type.replace('_', ' ').title(),
        "body_part": np.random.choice(body_parts.get(image_type, ['Unknown'])),
        "annotations": {
            "findings": ["Sample medical imaging data for demonstration"],
            "recommendations": ["This is placeholder data for educational purposes"]
        },
        "metadata": {
            "patient_id": "ANONYMIZED",
            "age": "N/A",
            "gender": "N/A",
            "technique": "Digital imaging",
            "sample_number": image_num
        }
    }
    
    with open(label_path, 'w') as f:
        json.dump(label_data, f, indent=4)
    print(f"Created label file: {label_path}")


def generate_sample_data():
    """Generate sample data for all image types."""
    print("=" * 60)
    print("Generating Sample Medical Imaging Data")
    print("=" * 60)
    print()
    
    # Create directories if they don't exist
    directories = [
        'data/xray/images',
        'data/xray/labels',
        'data/ct_scan/images',
        'data/ct_scan/labels',
        'data/mri/images',
        'data/mri/labels',
    ]
    
    for directory in directories:
        Path(directory).mkdir(parents=True, exist_ok=True)
    
    # Generate X-ray samples
    print("Generating X-ray samples...")
    for i in range(1, 6):  # 5 X-ray images
        image_path = f'data/xray/images/xray_{i:03d}.png'
        label_path = f'data/xray/labels/xray_{i:03d}.json'
        create_sample_xray_image(image_path, i)
        create_label_file(label_path, 'xray', i)
    
    print()
    
    # Generate CT scan samples
    print("Generating CT scan samples...")
    for i in range(1, 6):  # 5 CT scan images
        image_path = f'data/ct_scan/images/ct_scan_{i:03d}.png'
        label_path = f'data/ct_scan/labels/ct_scan_{i:03d}.json'
        create_sample_ct_image(image_path, i)
        create_label_file(label_path, 'ct_scan', i)
    
    print()
    
    # Generate MRI samples
    print("Generating MRI samples...")
    for i in range(1, 6):  # 5 MRI images
        image_path = f'data/mri/images/mri_{i:03d}.png'
        label_path = f'data/mri/labels/mri_{i:03d}.json'
        create_sample_mri_image(image_path, i)
        create_label_file(label_path, 'mri', i)
    
    print()
    print("=" * 60)
    print("Sample data generation completed!")
    print("=" * 60)
    print()
    print("Generated:")
    print("  - 5 X-ray images with labels")
    print("  - 5 CT scan images with labels")
    print("  - 5 MRI images with labels")
    print("  Total: 15 images with corresponding labels")
    print()
    print("Note: These are placeholder images for demonstration purposes.")
    print("For actual medical use, replace with real medical imaging data.")


if __name__ == '__main__':
    generate_sample_data()