import gradio as gr
from PIL import Image
from dataclasses import dataclass
import random
from transformers import pipeline
from huggingface_hub import InferenceClient, login
import os
from datetime import datetime
import json
from enum import Enum

class PromptFormat(Enum):
    XML = "xml"
    JSON = "json"
    MARKDOWN = "markdown"

@dataclass
class PatientMetadata:
    age: int
    smoking_status: str
    family_history: bool
    menopause_status: str
    previous_mammogram: bool
    breast_density: str
    hormone_therapy: bool

class MicrowaveBreastAnalyzer:
    def __init__(self, hf_token: str, prompt_format: PromptFormat = PromptFormat.XML):
        """Initialize the analyzer with models and specified prompt format."""
        print(f"Initializing system with {prompt_format.value} prompt format...")
        
        # Set prompt format
        self.prompt_format = prompt_format
        
        # Login to Hugging Face
        login(token=hf_token)
        
        # Initialize vision pipelines for tumor detection and size classification
        self.tumor_classifier = pipeline(
            "image-classification",
            model="SIATCN/vit_tumor_classifier",
            device="cpu"
        )
        
        self.size_classifier = pipeline(
            "image-classification",
            model="SIATCN/vit_tumor_radius_detection_finetuned",
            device="cpu"
        )
        
        # Initialize Mistral client for report generation
        self.report_generator = InferenceClient(
            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
            token=hf_token
        )
        
        print("Initialization complete!")

    def _generate_synthetic_metadata(self) -> PatientMetadata:
        """Generate realistic patient metadata for screening."""
        age = random.randint(40, 75)
        smoking_status = random.choice(["Never Smoker", "Former Smoker", "Current Smoker"])
        family_history = random.choice([True, False])
        menopause_status = "Post-menopausal" if age > 50 else "Pre-menopausal"
        previous_mammogram = random.choice([True, False])
        breast_density = random.choice([
            "A: Almost entirely fatty", 
            "B: Scattered fibroglandular",
            "C: Heterogeneously dense",
            "D: Extremely dense"
        ])
        hormone_therapy = random.choice([True, False])

        return PatientMetadata(
            age=age,
            smoking_status=smoking_status,
            family_history=family_history,
            menopause_status=menopause_status,
            previous_mammogram=previous_mammogram,
            breast_density=breast_density,
            hormone_therapy=hormone_therapy
        )

    def _process_image(self, image: Image.Image) -> Image.Image:
        """Process input image for model consumption."""
        if image.mode != 'RGB':
            image = image.convert('RGB')
        return image.resize((224, 224))

    def _generate_xml_prompt(self, has_tumor: bool, tumor_size: str, metadata: PatientMetadata) -> str:
        """Generate XML-style prompt."""
        return f"""<s>[INST] Generate a structured medical report for a microwave breast imaging scan using the following format exactly.
Keep sections consistent and use proper medical terminology. Be concise yet thorough.

EXAMINATION PERFORMED:
- Microwave Breast Imaging Scan
- Date: {datetime.now().strftime('%B %d, %Y')}

IMAGING FINDINGS:
Primary Finding: {'Abnormal area detected' if has_tumor else 'No abnormalities detected'}
{f'Detected Mass Size: {tumor_size} cm' if has_tumor else ''}

PATIENT HISTORY:
- Age: {metadata.age} years
- Menopausal Status: {metadata.menopause_status}
- Previous Screening: {'Yes' if metadata.previous_mammogram else 'No'}
- Tissue Characteristics: {metadata.breast_density}

RISK FACTORS:
{f'• Family History: {"Present" if metadata.family_history else "None"}'}
• Smoking Status: {metadata.smoking_status}
• Hormone Therapy: {'Yes' if metadata.hormone_therapy else 'No'}

Please generate a report with these exact sections:

1. DETAILED FINDINGS
[Describe the microwave imaging findings in detail, including location and characteristics of any detected abnormalities]

2. INTERPRETATION
[Provide a clear assessment of the microwave imaging results and their clinical significance]

3. RECOMMENDATIONS
[List specific follow-up actions and timeline]

4. TECHNICAL NOTES
[Include any relevant information about the scan quality and any technical considerations]

Format each section consistently and maintain professional medical terminology throughout. Note that this uses microwave imaging technology, not mammography. [/INST]</s>"""

    def _generate_json_prompt(self, has_tumor: bool, tumor_size: str, metadata: PatientMetadata) -> str:
        """Generate JSON-style prompt."""
        prompt_data = {
            "instruction": "Generate a structured medical report for a microwave breast imaging scan",
            "format_requirements": "Keep sections consistent and use proper medical terminology. Be concise yet thorough.",
            "input_data": {
                "examination": {
                    "type": "Microwave Breast Imaging Scan",
                    "date": datetime.now().strftime('%B %d, %Y')
                },
                "imaging_findings": {
                    "primary_finding": "Abnormal area detected" if has_tumor else "No abnormalities detected",
                    "mass_size": f"{tumor_size} cm" if has_tumor else None
                },
                "patient_history": {
                    "age": metadata.age,
                    "menopausal_status": metadata.menopause_status,
                    "previous_screening": metadata.previous_mammogram,
                    "tissue_characteristics": metadata.breast_density
                },
                "risk_factors": {
                    "family_history": "Present" if metadata.family_history else "None",
                    "smoking_status": metadata.smoking_status,
                    "hormone_therapy": "Yes" if metadata.hormone_therapy else "No"
                }
            },
            "required_sections": [
                "DETAILED FINDINGS",
                "INTERPRETATION",
                "RECOMMENDATIONS",
                "TECHNICAL NOTES"
            ],
            "section_guidelines": {
                "DETAILED_FINDINGS": "Describe the microwave imaging findings in detail, including location and characteristics of any detected abnormalities",
                "INTERPRETATION": "Provide a clear assessment of the microwave imaging results and their clinical significance",
                "RECOMMENDATIONS": "List specific follow-up actions and timeline",
                "TECHNICAL_NOTES": "Include any relevant information about the scan quality and any technical considerations"
            }
        }
        
        return f"<s>[INST] {json.dumps(prompt_data, indent=2)} [/INST]</s>"

    def _generate_markdown_prompt(self, has_tumor: bool, tumor_size: str, metadata: PatientMetadata) -> str:
        """Generate Markdown-style prompt."""
        return f"""<s>[INST]
# Medical Report Generation Request

## Context
Generate a structured medical report for a microwave breast imaging scan.

## Current Examination Data
* **Type:** Microwave Breast Imaging Scan
* **Date:** {datetime.now().strftime('%B %d, %Y')}

## Current Findings
* **Primary Finding:** {"Abnormal area detected" if has_tumor else "No abnormalities detected"}
* **Mass Size:** {f"{tumor_size} cm" if has_tumor else "N/A"}

## Patient Information
* **Age:** {metadata.age} years
* **Menopausal Status:** {metadata.menopause_status}
* **Previous Screening:** {"Yes" if metadata.previous_mammogram else "No"}
* **Tissue Characteristics:** {metadata.breast_density}

## Risk Assessment
* **Family History:** {"Present" if metadata.family_history else "None"}
* **Smoking Status:** {metadata.smoking_status}
* **Hormone Therapy:** {"Yes" if metadata.hormone_therapy else "No"}

## Required Report Sections
1. **Detailed Findings**
   - Include location and characteristics of any detected abnormalities
2. **Interpretation**
   - Assess microwave imaging results and clinical significance
3. **Recommendations**
   - Specify follow-up actions and timeline
4. **Technical Notes**
   - Document scan quality and technical considerations

Please maintain professional medical terminology throughout the report.
[/INST]</s>"""

    def _generate_medical_report(self, has_tumor: bool, tumor_size: str, metadata: PatientMetadata) -> str:
        """Generate a standardized report for microwave breast imaging."""
        # Select prompt format based on configuration
        if self.prompt_format == PromptFormat.XML:
            prompt = self._generate_xml_prompt(has_tumor, tumor_size, metadata)
        elif self.prompt_format == PromptFormat.JSON:
            prompt = self._generate_json_prompt(has_tumor, tumor_size, metadata)
        else:  # PromptFormat.MARKDOWN
            prompt = self._generate_markdown_prompt(has_tumor, tumor_size, metadata)

        # Generate response using Mistral
        response = self.report_generator.text_generation(
            prompt,
            max_new_tokens=800,
            temperature=0.3,
            top_p=0.9,
            repetition_penalty=1.1,
            do_sample=True,
            seed=42
        )
            
        # Post-process the response to ensure consistent formatting
        formatted_response = f"""MICROWAVE BREAST IMAGING REPORT
Date: {datetime.now().strftime('%B %d, %Y')}
----------------------------------------

{response.strip()}

----------------------------------------
NOTE: This report was generated using AI assistance and should be reviewed by a qualified healthcare professional.
This screening was performed using microwave imaging technology."""

        return formatted_response

    def analyze(self, image: Image.Image) -> str:
        """Main analysis pipeline with standardized output."""
        try:
            processed_image = self._process_image(image)
            metadata = self._generate_synthetic_metadata()

            # Detect tumor
            tumor_result = self.tumor_classifier(processed_image)
            has_tumor = tumor_result[0]['label'] == 'tumor'
            tumor_confidence = tumor_result[0]['score']

            # Measure size if tumor detected
            size_result = self.size_classifier(processed_image)
            tumor_size = size_result[0]['label'].replace('tumor-', '')

            # Generate report
            report = self._generate_medical_report(has_tumor, tumor_size, metadata)

            return f"""MICROWAVE BREAST IMAGING ANALYSIS
========================================

INITIAL SCAN ASSESSMENT:
{'⚠️ ABNORMAL FINDING DETECTED' if has_tumor else '✓ NO ABNORMALITIES DETECTED'}
Detection Confidence: {tumor_confidence:.2%}
{f'Estimated Mass Size: {tumor_size} cm' if has_tumor else ''}

----------------------------------------
{report}"""

        except Exception as e:
            import traceback
            return f"Error during analysis: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"

def create_interface(hf_token: str, prompt_format: PromptFormat = PromptFormat.XML) -> gr.Interface:
    """Create the Gradio interface."""
    analyzer = MicrowaveBreastAnalyzer(hf_token, prompt_format)
    
    interface = gr.Interface(
        fn=analyzer.analyze,
        inputs=[
            gr.Image(type="pil", label="Upload Microwave Breast Image for Analysis")
        ],
        outputs=[
            gr.Textbox(label="Analysis Results", lines=20)
        ],
        title=f"Microwave Breast Imaging Analysis System ({prompt_format.value.upper()} Format)",
        description="""Upload a microwave breast image for comprehensive analysis. The system will:
        1. Detect the presence of tumors using microwave imaging technology
        2. Classify tumor size if present
        3. Generate a detailed medical report with recommendations
        
        Note: This system uses microwave imaging technology for breast screening, which offers a safe, 
        radiation-free alternative to traditional mammography.""",
    )
    
    return interface

if __name__ == "__main__":
    print("Starting microwave breast imaging analysis system...")
    # Load HuggingFace token from secrets
    HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
    if not HF_TOKEN:
        raise ValueError("Please set HUGGINGFACE_TOKEN environment variable")
    
    # Create interfaces for different formats
    interface_xml = create_interface(HF_TOKEN, PromptFormat.XML)
    interface_json = create_interface(HF_TOKEN, PromptFormat.JSON)
    interface_markdown = create_interface(HF_TOKEN, PromptFormat.MARKDOWN)
    
    # Launch the XML version by default
    interface_xml.launch(
        debug=True,
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )