# app.py
import os
import gradio as gr
import torch
import requests
from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
from monai.networks.nets import DenseNet121
import torchxrayvision as xrv

# Configuration
DEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")  # Set in Hugging Face secrets

DISCLAIMER = """
<div style="color: red; border: 2px solid red; padding: 15px; margin: 10px;">
⚠️ WARNING: This is a prototype demonstration only. NOT ACTUAL MEDICAL ADVICE. 
DO NOT USE FOR REAL HEALTH DECISIONS. CONSULT LICENSED PROFESSIONALS.
</div>
"""

class MedicalAssistant:
    def __init__(self):
        # Medical imaging models
        self.medical_models = self._init_imaging_models()
        
        # Clinical text processing
        self.prescription_parser = pipeline(
            "token-classification", 
            model="obi/deid_bert_i2b2",
            tokenizer="obi/deid_bert_i2b2"
        )
        
        # Safety systems
        self.safety_filter = pipeline(
            "text-classification", 
            model="Hate-speech-CNERG/dehatebert-mono-english"
        )

    def _init_imaging_models(self):
        """Initialize medical imaging models"""
        return {
            "xray": xrv.models.DenseNet(weights="densenet121-res224-all"),
            "ct": DenseNet121(spatial_dims=3, in_channels=1, out_channels=14),
            "histo": torch.hub.load('pytorch/vision', 'resnet50', pretrained=True)
        }

    def query_deepseek(self, prompt: str):
        """Query DeepSeek API"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
            "Accept": "application/json"
        }
        
        payload = {
            "model": "deepseek-chat",
            "messages": [{
                "role": "user",
                "content": f"MEDICAL PROMPT: {prompt}\nRespond with research-supported medical information. Cite sources from peer-reviewed journals."
            }],
            "temperature": 0.3,
            "max_tokens": 600,
            "top_p": 0.95
        }

        try:
            response = requests.post(
                DEEPSEEK_API_URL,
                headers=headers,
                json=payload,
                timeout=30
            )
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except requests.HTTPError as e:
            error_detail = response.json().get('error', {}).get('message', 'Unknown error')
            return f"API Error {e.response.status_code}: {error_detail}"
        except Exception as e:
            return f"Connection Error: {str(e)}"

    # Rest of the class remains the same...

# Initialize system
assistant = MedicalAssistant()

def process_input(query, image, prescription):
    context = {}
    
    if image is not None:
        context["image_analysis"] = assistant.analyze_image(image, "xray")
    
    if prescription:
        context["prescription"] = assistant.parse_prescription(prescription)
    
    return assistant.generate_response(query, context)

# Gradio interface
interface = gr.Interface(
    fn=process_input,
    inputs=[
        gr.Textbox(label="Medical Query", placeholder="Enter your medical question..."),
        gr.Image(label="Medical Imaging", type="filepath"),
        gr.Textbox(label="Prescription Text")
    ],
    outputs=gr.Textbox(label="Research-Backed Response"),
    title="AI Medical Research Assistant",
    description=DISCLAIMER,
    allow_flagging="never"
)

interface.launch()