import gradio as gr
import numpy as np
import pydicom
from skimage import exposure, measure, filters
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr
import json
from typing import Dict, List, Tuple
import tempfile
import os
from PIL import Image


class VeterinaryDICOMProcessor:
    """Veterinärmedizinische DICOM-Bildverbesserung und Qualitätsbewertung"""
    
    def __init__(self):
        # Spezies-spezifische Parameter für verschiedene Algorithmen
        self.species_params = {
            "canine": {
                "clahe": {"clip_limit": 0.02, "kernel_size": (8, 8)},
                "adaptive": {"clip_limit": 0.015, "kernel_size": (6, 6)},
                "histogram": {"bins": 256}
            },
            "feline": {
                "clahe": {"clip_limit": 0.015, "kernel_size": (6, 6)},
                "adaptive": {"clip_limit": 0.012, "kernel_size": (4, 4)},
                "histogram": {"bins": 256}
            },
            "equine": {
                "clahe": {"clip_limit": 0.03, "kernel_size": (12, 12)},
                "adaptive": {"clip_limit": 0.025, "kernel_size": (10, 10)},
                "histogram": {"bins": 256}
            },
            "bovine": {
                "clahe": {"clip_limit": 0.025, "kernel_size": (10, 10)},
                "adaptive": {"clip_limit": 0.02, "kernel_size": (8, 8)},
                "histogram": {"bins": 256}
            }
        }
    
    def load_dicom_or_image(self, file_path: str) -> np.ndarray:
        """Lädt DICOM-Datei oder normale Bilddatei und konvertiert zu Numpy Array"""
        try:
            # Versuche zuerst DICOM zu laden
            if file_path.lower().endswith(('.dcm', '.dicom')):
                dicom_data = pydicom.dcmread(file_path)
                image = dicom_data.pixel_array.astype(np.float64)
            else:
                # Lade als normales Bild (für Demo-Zwecke)
                pil_image = Image.open(file_path).convert('L')  # Graustufen
                image = np.array(pil_image, dtype=np.float64)
            
            # Normalisierung für bessere Verarbeitung
            if image.max() > image.min():
                image = (image - image.min()) / (image.max() - image.min())
            
            return image
        except Exception as e:
            raise ValueError(f"Fehler beim Laden der Datei: {e}")
    
    def apply_clahe_enhancement(self, image: np.ndarray, species: str) -> np.ndarray:
        """Wendet CLAHE (Contrast Limited Adaptive Histogram Equalization) an"""
        params = self.species_params.get(species, self.species_params["canine"])["clahe"]
        
        enhanced = exposure.equalize_adapthist(
            image,
            kernel_size=params["kernel_size"],
            clip_limit=params["clip_limit"],
            nbins=256
        )
        return enhanced
    
    def apply_adaptive_histogram_equalization(self, image: np.ndarray, species: str) -> np.ndarray:
        """Wendet Adaptive Histogram Equalization an (lokale Variante)"""
        params = self.species_params.get(species, self.species_params["canine"])["adaptive"]
        
        # Adaptive Histogram Equalization mit kleineren Tiles
        enhanced = exposure.equalize_adapthist(
            image,
            kernel_size=params["kernel_size"],
            clip_limit=params["clip_limit"],
            nbins=128  # Weniger Bins für adaptives Verfahren
        )
        return enhanced
    
    def apply_histogram_equalization(self, image: np.ndarray) -> np.ndarray:
        """Standard globale Histogram Equalization"""
        return exposure.equalize_hist(image)
    
    def apply_contrast_stretching(self, image: np.ndarray, percentiles: Tuple[float, float] = (2, 98)) -> np.ndarray:
        """Kontrast-Streckung basierend auf Perzentilen"""
        p_low, p_high = np.percentile(image, percentiles)
        return exposure.rescale_intensity(image, in_range=(p_low, p_high))
    
    def apply_gamma_correction(self, image: np.ndarray, gamma: float = 1.2) -> np.ndarray:
        """Gamma-Korrektur für Helligkeit"""
        return exposure.adjust_gamma(image, gamma)
    
    def calculate_quality_metrics(self, original: np.ndarray, enhanced: np.ndarray) -> Dict[str, float]:
        """Berechnet umfassende Bildqualitäts-Metriken"""
        
        # SSIM (Structural Similarity Index)
        ssim_score = ssim(original, enhanced, data_range=1.0)
        
        # PSNR (Peak Signal-to-Noise Ratio)
        psnr_score = psnr(original, enhanced, data_range=1.0)
        
        # Histogram-basierte Metriken
        original_entropy = measure.shannon_entropy(original)
        enhanced_entropy = measure.shannon_entropy(enhanced)
        
        # Kontrast-Metriken (RMS Kontrast)
        original_contrast = np.sqrt(np.mean((original - original.mean()) ** 2))
        enhanced_contrast = np.sqrt(np.mean((enhanced - enhanced.mean()) ** 2))
        
        # Edge-basierte Metriken
        original_edges = np.mean(filters.sobel(original))
        enhanced_edges = np.mean(filters.sobel(enhanced))
        
        # Dynamikbereich
        original_dynamic_range = np.max(original) - np.min(original)
        enhanced_dynamic_range = np.max(enhanced) - np.min(enhanced)
        
        return {
            "ssim": float(ssim_score),
            "psnr": float(psnr_score),
            "original_entropy": float(original_entropy),
            "enhanced_entropy": float(enhanced_entropy),
            "entropy_improvement": float(enhanced_entropy - original_entropy),
            "original_contrast": float(original_contrast),
            "enhanced_contrast": float(enhanced_contrast),
            "contrast_improvement": float(enhanced_contrast - original_contrast),
            "original_edge_density": float(original_edges),
            "enhanced_edge_density": float(enhanced_edges),
            "edge_improvement": float(enhanced_edges - original_edges),
            "original_dynamic_range": float(original_dynamic_range),
            "enhanced_dynamic_range": float(enhanced_dynamic_range),
            "dynamic_range_improvement": float(enhanced_dynamic_range - original_dynamic_range)
        }
    
    def array_to_image(self, array: np.ndarray) -> str:
        """Konvertiert Numpy Array zu Base64-String für Output"""
        # Normalisiere auf 0-255
        img_normalized = ((array - array.min()) / (array.max() - array.min()) * 255).astype(np.uint8)
        
        # Konvertiere zu PIL Image
        pil_image = Image.fromarray(img_normalized, mode='L')
        
        # Speichere temporär
        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
            pil_image.save(tmp_file, format='PNG')
            return tmp_file.name
    
    def generate_agent_prompt(self, metrics: Dict[str, float], species: str, 
                            body_part: str, enhancement_method: str) -> str:
        """Generiert strukturierten Prompt für KI-Agent zur Qualitätsbewertung"""
        
        prompt = f"""VETERINÄRMEDIZINISCHE BILDQUALITÄTS-ANALYSE

PATIENT-INFORMATION:
===================
• Spezies: {species.upper()}
• Körperregion: {body_part}
• Verbesserungsmethode: {enhancement_method}

BILDQUALITÄTS-METRIKEN:
======================
• SSIM (Strukturelle Ähnlichkeit): {metrics['ssim']:.3f}
  → 1.0 = perfekte Ähnlichkeit, 0.0 = völlig unterschiedlich
  → Bewertung: {'Ausgezeichnet' if metrics['ssim'] > 0.8 else 'Gut' if metrics['ssim'] > 0.6 else 'Verbesserungsbedürftig'}

• PSNR (Signal-Rausch-Verhältnis): {metrics['psnr']:.1f} dB
  → >30dB = gut, >25dB = akzeptabel, <20dB = problematisch
  → Bewertung: {'Gut' if metrics['psnr'] > 30 else 'Akzeptabel' if metrics['psnr'] > 25 else 'Problematisch'}

• Entropie-Verbesserung: {metrics['entropy_improvement']:.3f}
  → Positiv = mehr Bildinformation, Negativ = Informationsverlust
  → Status: {'Information gewonnen' if metrics['entropy_improvement'] > 0 else 'Information verloren'}

• Kontrast-Verbesserung: {metrics['contrast_improvement']:.3f}
  → Positiv = besserer Kontrast, Negativ = Kontrastverlust
  → Status: {'Kontrast verbessert' if metrics['contrast_improvement'] > 0 else 'Kontrast reduziert'}

• Edge-Verbesserung: {metrics['edge_improvement']:.3f}
  → Positiv = schärfere Strukturen, Negativ = Unschärfe
  → Status: {'Schärfer' if metrics['edge_improvement'] > 0 else 'Unschärfer'}

AUFGABEN FÜR KI-AGENT:
=====================
1. Bewerten Sie die Bildqualität für {species}-Diagnostik (Skala 1-10)
2. Ist die Verbesserung für klinische Diagnose ausreichend?
3. Welche anatomischen Strukturen sind bei {species} in {body_part} besser/schlechter sichtbar?
4. Empfehlen Sie weitere Bildverbesserungsschritte oder alternative Methoden?
5. Identifizieren Sie potenzielle Artefakte oder diagnostische Probleme?
6. Geben Sie spezies-spezifische Interpretationshilfen?

ANTWORT-FORMAT:
==============
- Strukturiert und klinisch relevant
- Fokus auf diagnostischen Nutzen
- Spezies-spezifische Besonderheiten berücksichtigen
- Konkrete Handlungsempfehlungen
"""
        return prompt

# Globale Processor-Instanz
processor = VeterinaryDICOMProcessor()

def enhance_dicom_image(image_file, species: str, enhancement_method: str, body_part: str = "unbekannt"):
    """
    Verbessert veterinärmedizinische DICOM-Bilder mit verschiedenen Algorithmen.
    
    Args:
        image_file: DICOM-Datei oder Testbild
        species: Tierart (canine, feline, equine, bovine)
        enhancement_method: Verbesserungsmethode (clahe, adaptive, histogram, contrast_stretch, gamma)
        body_part: Körperregion für spezifische Analyse
        
    Returns:
        Verbessertes Bild und detaillierte Qualitätsmetriken
    """
    try:
        # Lade Originalbild
        if image_file is None:
            return None, "❌ Kein Bild hochgeladen"
        
        original_image = processor.load_dicom_or_image(image_file)
        
        # Wähle Verbesserungsmethode
        if enhancement_method == "clahe":
            enhanced_image = processor.apply_clahe_enhancement(original_image, species)
        elif enhancement_method == "adaptive":
            enhanced_image = processor.apply_adaptive_histogram_equalization(original_image, species)
        elif enhancement_method == "histogram":
            enhanced_image = processor.apply_histogram_equalization(original_image)
        elif enhancement_method == "contrast_stretch":
            enhanced_image = processor.apply_contrast_stretching(original_image)
        elif enhancement_method == "gamma":
            enhanced_image = processor.apply_gamma_correction(original_image)
        else:
            enhanced_image = processor.apply_clahe_enhancement(original_image, species)
        
        # Berechne Qualitäts-Metriken
        metrics = processor.calculate_quality_metrics(original_image, enhanced_image)
        
        # Generiere Agent-Prompt
        agent_prompt = processor.generate_agent_prompt(metrics, species, body_part, enhancement_method)
        
        # Konvertiere verbessertes Bild für Output
        enhanced_image_path = processor.array_to_image(enhanced_image)
        
        # Erstelle detaillierten Report
        report = f"""
🔬 VETERINÄRMEDIZINISCHE BILDANALYSE
=====================================

📊 QUALITÄTS-METRIKEN:
• SSIM: {metrics['ssim']:.3f} ({'✅ Gut' if metrics['ssim'] > 0.7 else '⚠️ Prüfen'})
• PSNR: {metrics['psnr']:.1f} dB ({'✅ Gut' if metrics['psnr'] > 25 else '⚠️ Niedrig'})
• Entropie: {metrics['entropy_improvement']:+.3f} ({'✅ Information gewonnen' if metrics['entropy_improvement'] > 0 else '❌ Information verloren'})
• Kontrast: {metrics['contrast_improvement']:+.3f} ({'✅ Verbessert' if metrics['contrast_improvement'] > 0 else '❌ Reduziert'})
• Edge-Definition: {metrics['edge_improvement']:+.3f} ({'✅ Schärfer' if metrics['edge_improvement'] > 0 else '❌ Unschärfer'})

🐾 SPEZIES-ANALYSE: {species.upper()}
🎯 METHODE: {enhancement_method.upper()}
📍 REGION: {body_part.upper()}

🤖 KI-AGENT PROMPT:
{agent_prompt}
"""
        
        return enhanced_image_path, report
        
    except Exception as e:
        return None, f"❌ Fehler bei Bildverbesserung: {str(e)}"

def compare_enhancement_methods(image_file, species: str, body_part: str = "thorax"):
    """
    Vergleicht verschiedene Bildverbesserungsmethoden für veterinärmedizinische Analyse.
    
    Args:
        image_file: DICOM-Datei oder Testbild
        species: Tierart (canine, feline, equine, bovine)
        body_part: Körperregion für Analyse
        
    Returns:
        Vergleichstabelle aller Methoden mit Qualitätsmetriken
    """
    if image_file is None:
        return "❌ Kein Bild hochgeladen"
    
    try:
        original_image = processor.load_dicom_or_image(image_file)
        methods = ["clahe", "adaptive", "histogram", "contrast_stretch", "gamma"]
        
        results = []
        
        for method in methods:
            # Wende Methode an
            if method == "clahe":
                enhanced = processor.apply_clahe_enhancement(original_image, species)
            elif method == "adaptive":
                enhanced = processor.apply_adaptive_histogram_equalization(original_image, species)
            elif method == "histogram":
                enhanced = processor.apply_histogram_equalization(original_image)
            elif method == "contrast_stretch":
                enhanced = processor.apply_contrast_stretching(original_image)
            elif method == "gamma":
                enhanced = processor.apply_gamma_correction(original_image)
            
            # Berechne Metriken
            metrics = processor.calculate_quality_metrics(original_image, enhanced)
            
            results.append({
                "Methode": method.upper(),
                "SSIM": f"{metrics['ssim']:.3f}",
                "PSNR (dB)": f"{metrics['psnr']:.1f}",
                "Entropie Δ": f"{metrics['entropy_improvement']:+.3f}",
                "Kontrast Δ": f"{metrics['contrast_improvement']:+.3f}",
                "Edge Δ": f"{metrics['edge_improvement']:+.3f}",
                "Empfehlung": "✅ Gut" if metrics['ssim'] > 0.7 and metrics['psnr'] > 25 else "⚠️ Prüfen"
            })
        
        # Erstelle Vergleichstabelle
        comparison = f"""
🔬 METHODEN-VERGLEICH: {species.upper()} - {body_part.upper()}
====================================================

"""
        for result in results:
            comparison += f"""
📊 {result['Methode']}:
   SSIM: {result['SSIM']} | PSNR: {result['PSNR (dB)']} | {result['Empfehlung']}
   Entropie: {result['Entropie Δ']} | Kontrast: {result['Kontrast Δ']} | Edges: {result['Edge Δ']}
"""
        
        # Beste Methode identifizieren
        best_method = max(results, key=lambda x: float(x['SSIM']))
        comparison += f"""
🏆 BESTE METHODE: {best_method['Methode']}
   → Höchste SSIM: {best_method['SSIM']}
   → Empfehlung für {species} {body_part}-Diagnostik

💡 VETERINÄR-TIPP:
   - CLAHE: Optimal für lokale Kontrastverbesserung
   - ADAPTIVE: Weniger Artefakte, sanfter
   - HISTOGRAM: Globale Verbesserung, kann übersättigen
   - CONTRAST_STRETCH: Konservativ, für kritische Diagnosen
   - GAMMA: Helligkeitsanpassung bei unter-/überbelichteten Bildern
"""
        
        return comparison
        
    except Exception as e:
        return f"❌ Fehler beim Methodenvergleich: {str(e)}"

# Gradio Interface
with gr.Blocks(title="🐾 Veterinary DICOM Enhancement MCP Server") as demo:
    gr.Markdown("""
    # 🐾 Veterinärmedizinische DICOM-Bildverbesserung
    
    **MCP Server für KI-Agenten zur automatischen Bildqualitätsbewertung**
    
    🔬 Spezies-spezifische Bildverbesserung mit CLAHE, Adaptive & Histogram Equalization
    """)
    
    with gr.Tab("🎯 Einzelne Bildverbesserung"):
        with gr.Row():
            with gr.Column():
                input_image = gr.File(
                    label="📁 DICOM-Datei oder Testbild hochladen",
                    file_types=[".dcm", ".dicom", ".png", ".jpg", ".jpeg"]
                )
                species = gr.Dropdown(
                    choices=["canine", "feline", "equine", "bovine"],
                    label="🐾 Tierart",
                    value="canine"
                )
                enhancement_method = gr.Dropdown(
                    choices=["clahe", "adaptive", "histogram", "contrast_stretch", "gamma"],
                    label="🔧 Verbesserungsmethode",
                    value="clahe"
                )
                body_part = gr.Textbox(
                    label="📍 Körperregion",
                    value="thorax",
                    placeholder="z.B. thorax, abdomen, extremität"
                )
                
            with gr.Column():
                enhanced_output = gr.Image(label="✨ Verbessertes Bild")
                metrics_output = gr.Textbox(
                    label="📊 Qualitätsanalyse & KI-Agent Prompt",
                    lines=20,
                    max_lines=30
                )
        
        enhance_btn = gr.Button("🚀 Bild verbessern", variant="primary")
        enhance_btn.click(
            enhance_dicom_image,
            inputs=[input_image, species, enhancement_method, body_part],
            outputs=[enhanced_output, metrics_output]
        )
    
    with gr.Tab("📊 Methoden-Vergleich"):
        with gr.Row():
            with gr.Column():
                compare_input = gr.File(
                    label="📁 DICOM-Datei oder Testbild hochladen",
                    file_types=[".dcm", ".dicom", ".png", ".jpg", ".jpeg"]
                )
                compare_species = gr.Dropdown(
                    choices=["canine", "feline", "equine", "bovine"],
                    label="🐾 Tierart",
                    value="canine"
                )
                compare_body_part = gr.Textbox(
                    label="📍 Körperregion",
                    value="thorax"
                )
                
            with gr.Column():
                comparison_output = gr.Textbox(
                    label="📈 Methoden-Vergleich",
                    lines=25,
                    max_lines=35
                )
        
        compare_btn = gr.Button("📊 Alle Methoden vergleichen", variant="primary")
        compare_btn.click(
            compare_enhancement_methods,
            inputs=[compare_input, compare_species, compare_body_part],
            outputs=[comparison_output]
        )

if __name__ == "__main__":
    # MCP Server starten für Hugging Face
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,         # HF Standard Port
        mcp_server=True,          # MCP aktivieren
        share=False,              # Nicht nötig bei HF
        show_error=True,
        show_api=True
    )