"""
Medical Imaging Dataset - Visualization
========================================

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 module provides visualization functions for medical images
with annotations and diagnostic labels.
"""

import numpy as np
import matplotlib.pyplot as plt
import cv2
from pathlib import Path
import json
from typing import Optional, Dict, List


def visualize_medical_image(image_path: str, 
                            label_path: Optional[str] = None,
                            title: Optional[str] = None,
                            save_path: Optional[str] = None) -> None:
    """
    Visualize a medical image with optional annotations.
    
    Args:
        image_path (str): Path to the medical image
        label_path (str): Path to the label/annotation file (JSON)
        title (str): Title for the plot
        save_path (str): Path to save the visualization
    """
    # Load image
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    
    if image is None:
        raise ValueError(f"Could not load image from {image_path}")
    
    # Load labels if provided
    labels = None
    if label_path and Path(label_path).exists():
        try:
            with open(label_path, 'r') as f:
                labels = json.load(f)
        except (json.JSONDecodeError, IOError) as e:
            print(f"Warning: Could not load label file {label_path}: {e}")
            labels = None
    
    # Create figure
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    
    # Display image
    ax.imshow(image, cmap='gray')
    ax.axis('off')
    
    # Add title
    if title:
        ax.set_title(title, fontsize=16, fontweight='bold')
    elif labels and 'diagnosis' in labels:
        ax.set_title(f"Diagnosis: {labels['diagnosis']}", fontsize=16, fontweight='bold')
    else:
        ax.set_title(Path(image_path).stem, fontsize=16, fontweight='bold')
    
    # Add annotations if available
    if labels:
        annotation_text = []
        if 'diagnosis' in labels:
            annotation_text.append(f"Diagnosis: {labels['diagnosis']}")
        if 'confidence' in labels:
            annotation_text.append(f"Confidence: {labels['confidence']:.2f}")
        if 'date' in labels:
            annotation_text.append(f"Date: {labels['date']}")
        
        if annotation_text:
            text = '\n'.join(annotation_text)
            ax.text(0.02, 0.98, text, transform=ax.transAxes,
                   fontsize=12, verticalalignment='top',
                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    plt.tight_layout()
    
    # Save if path provided
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Visualization saved to {save_path}")
    
    plt.show()


def visualize_batch(images: List[np.ndarray],
                   labels: Optional[List[Dict]] = None,
                   titles: Optional[List[str]] = None,
                   cols: int = 3,
                   figsize: tuple = (15, 5)) -> None:
    """
    Visualize a batch of medical images.
    
    Args:
        images (list): List of image arrays
        labels (list): List of label dictionaries
        titles (list): List of titles for each image
        cols (int): Number of columns in the grid
        figsize (tuple): Figure size
    """
    n_images = len(images)
    rows = (n_images + cols - 1) // cols
    
    fig, axes = plt.subplots(rows, cols, figsize=figsize)
    
    if n_images == 1:
        axes = [axes]
    else:
        axes = axes.flatten()
    
    for i, image in enumerate(images):
        ax = axes[i]
        ax.imshow(image, cmap='gray')
        ax.axis('off')
        
        # Add title
        if titles and i < len(titles):
            ax.set_title(titles[i], fontsize=10)
        elif labels and i < len(labels) and labels[i] and 'diagnosis' in labels[i]:
            ax.set_title(labels[i]['diagnosis'], fontsize=10)
        else:
            ax.set_title(f"Image {i+1}", fontsize=10)
    
    # Hide unused subplots
    for i in range(n_images, len(axes)):
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.show()


def compare_preprocessing(image_path: str,
                          save_path: Optional[str] = None) -> None:
    """
    Compare original and preprocessed versions of a medical image.
    
    Args:
        image_path (str): Path to the medical image
        save_path (str): Path to save the comparison
    """
    from scripts.preprocess import preprocess_image
    
    # Load original image
    original = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    
    # Determine image type from path
    image_type = 'xray'
    if 'ct' in image_path.lower() or 'ct_scan' in image_path.lower():
        image_type = 'ct_scan'
    elif 'mri' in image_path.lower():
        image_type = 'mri'
    
    # Preprocess image
    processed = preprocess_image(image_path, image_type=image_type)
    
    # Create comparison plot
    fig, axes = plt.subplots(1, 2, figsize=(15, 7))
    
    # Original
    axes[0].imshow(original, cmap='gray')
    axes[0].set_title('Original Image', fontsize=14, fontweight='bold')
    axes[0].axis('off')
    
    # Processed
    axes[1].imshow(processed, cmap='gray')
    axes[1].set_title('Preprocessed Image', fontsize=14, fontweight='bold')
    axes[1].axis('off')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Comparison saved to {save_path}")
    
    plt.show()


def plot_statistics(dataset_info: Dict) -> None:
    """
    Plot dataset statistics.
    
    Args:
        dataset_info (dict): Dictionary containing dataset statistics
    """
    categories = ['X-ray', 'CT Scan', 'MRI']
    counts = [
        dataset_info.get('xray_count', 0),
        dataset_info.get('ct_count', 0),
        dataset_info.get('mri_count', 0)
    ]
    
    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
    
    bars = ax.bar(categories, counts, color=['#0dcaf0', '#6c757d', '#0d6efd'])
    
    # Add value labels on bars
    for bar, count in zip(bars, counts):
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height,
               f'{count}',
               ha='center', va='bottom', fontsize=12, fontweight='bold')
    
    ax.set_ylabel('Number of Images', fontsize=12)
    ax.set_title('Medical Imaging Dataset Statistics', fontsize=14, fontweight='bold')
    ax.grid(axis='y', alpha=0.3)
    
    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    # Example usage
    print("Medical Imaging Dataset - Visualization Module")
    print("This module provides visualization functions for medical images.")
    print("\nAvailable functions:")
    print("- visualize_medical_image(): Visualize a single medical image")
    print("- visualize_batch(): Visualize multiple images in a grid")
    print("- compare_preprocessing(): Compare original and preprocessed images")
    print("- plot_statistics(): Plot dataset statistics")