from transformers import AutoModelForCausalLM, AutoProcessor
from PIL import Image
import torch
import gradio as gr
import requests
import tempfile

device = torch.device("cpu")
model = AutoModelForCausalLM.from_pretrained("microsoft/maira-2", trust_remote_code=True)
processor = AutoProcessor.from_pretrained("microsoft/maira-2", trust_remote_code=True)
model = model.eval().to(device)

def get_sample_data():
    """Download sample medical images and data"""
    frontal_url = "https://openi.nlm.nih.gov/imgs/512/145/145/CXR145_IM-0290-1001.png"
    lateral_url = "https://openi.nlm.nih.gov/imgs/512/145/145/CXR145_IM-0290-2001.png"
    
    def download_image(url):
        response = requests.get(url, headers={"User-Agent": "MAIRA-2"}, stream=True)
        return Image.open(response.raw)
    
    return {
        "frontal": download_image(frontal_url),
        "lateral": download_image(lateral_url),
        "indication": "Dyspnea.",
        "technique": "PA and lateral views of the chest.",
        "comparison": "None.",
        "phrase": "Pleural effusion."
    }

def save_temp_image(img):
    """Save PIL image to temporary file"""
    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    img.save(temp_file.name)
    return temp_file.name

def load_sample_findings():
    """Load sample data for findings generation"""
    sample = get_sample_data()
    return [
        save_temp_image(sample["frontal"]),
        save_temp_image(sample["lateral"]),
        sample["indication"],
        sample["technique"],
        sample["comparison"],
        None, None, None, False
    ]

def load_sample_phrase():
    """Load sample data for phrase grounding"""
    sample = get_sample_data()
    return [save_temp_image(sample["frontal"]), sample["phrase"]]

def generate_report(frontal_path, lateral_path, indication, technique, comparison, 
                   prior_frontal_path, prior_lateral_path, prior_report, grounding):
    """Generate radiology report with optional grounding"""
    try:
        # Load images
        current_frontal = Image.open(frontal_path)
        current_lateral = Image.open(lateral_path)
        prior_frontal = Image.open(prior_frontal_path) if prior_frontal_path else None
        prior_lateral = Image.open(prior_lateral_path) if prior_lateral_path else None

        # Process inputs
        processed = processor.format_and_preprocess_reporting_input(
            current_frontal=current_frontal,
            current_lateral=current_lateral,
            prior_frontal=prior_frontal,
            prior_lateral=prior_lateral,
            indication=indication,
            technique=technique,
            comparison=comparison,
            prior_report=prior_report or None,
            return_tensors="pt",
            get_grounding=grounding
        ).to(device)

        # Generate report
        outputs = model.generate(**processed, 
                               max_new_tokens=450 if grounding else 300,
                               use_cache=True)
        
        # Decode and format
        prompt_length = processed["input_ids"].shape[-1]
        decoded = processor.decode(outputs[0][prompt_length:], skip_special_tokens=True)
        return processor.convert_output_to_plaintext_or_grounded_sequence(decoded.lstrip())
    
    except Exception as e:
        return f"Error: {str(e)}"

def ground_phrase(frontal_path, phrase):
    """Perform phrase grounding on image"""
    try:
        frontal = Image.open(frontal_path)
        processed = processor.format_and_preprocess_phrase_grounding_input(
            frontal_image=frontal,
            phrase=phrase,
            return_tensors="pt"
        ).to(device)
        
        outputs = model.generate(**processed, max_new_tokens=150, use_cache=True)
        
        prompt_length = processed["input_ids"].shape[-1]
        decoded = processor.decode(outputs[0][prompt_length:], skip_special_tokens=True)
        return processor.convert_output_to_plaintext_or_grounded_sequence(decoded)
    
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks(title="MAIRA-2 Medical Imaging Assistant") as demo:
    gr.Markdown("# MAIRA-2 Medical Imaging Assistant\nAI-powered radiology report generation and phrase grounding")
    
    with gr.Tab("Report Generation"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("## Current Study")
                frontal = gr.Image(label="Frontal View", type="filepath")
                lateral = gr.Image(label="Lateral View", type="filepath")
                indication = gr.Textbox(label="Clinical Indication")
                technique = gr.Textbox(label="Imaging Technique")
                comparison = gr.Textbox(label="Comparison")
                
                gr.Markdown("## Prior Study (Optional)")
                prior_frontal = gr.Image(label="Prior Frontal View", type="filepath")
                prior_lateral = gr.Image(label="Prior Lateral View", type="filepath")
                prior_report = gr.Textbox(label="Prior Report")
                
                grounding = gr.Checkbox(label="Include Grounding")
                sample_btn = gr.Button("Load Sample Data")
                
            with gr.Column():
                report_output = gr.Textbox(label="Generated Report", lines=10)
                generate_btn = gr.Button("Generate Report")
        
        sample_btn.click(load_sample_findings,
                        outputs=[frontal, lateral, indication, technique, comparison,
                               prior_frontal, prior_lateral, prior_report, grounding])
        generate_btn.click(generate_report,
                         inputs=[frontal, lateral, indication, technique, comparison,
                               prior_frontal, prior_lateral, prior_report, grounding],
                         outputs=report_output)
    
    with gr.Tab("Phrase Grounding"):
        with gr.Row():
            with gr.Column():
                pg_frontal = gr.Image(label="Frontal View", type="filepath")
                phrase = gr.Textbox(label="Phrase to Ground")
                pg_sample_btn = gr.Button("Load Sample Data")
            with gr.Column():
                pg_output = gr.Textbox(label="Grounding Result", lines=3)
                pg_btn = gr.Button("Find Phrase")
        
        pg_sample_btn.click(load_sample_phrase,
                           outputs=[pg_frontal, phrase])
        pg_btn.click(ground_phrase,
                    inputs=[pg_frontal, phrase],
                    outputs=pg_output)

demo.launch()