import gradio as gr
import easyocr
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

print("Loading OCR model...")
reader = easyocr.Reader(['en'], gpu=False)

print("Loading MediLlama-3.2...")
MODEL_NAME = "deep-div/MediLlama-3.2"

try:
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=quantization_config,
        device_map="cpu",
        low_cpu_mem_usage=True
    )
    print("✅ Loaded with 4-bit quantization")
except Exception as e:
    print(f"⚠️ Quantization failed, loading without...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        device_map="cpu"
    )

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

stored_text = ""

def extract_text(image):
    global stored_text
    if image is None:
        return ""
    if isinstance(image, np.ndarray):
        img = image
    else:
        img = np.array(image)
    result = reader.readtext(img, detail=0)
    stored_text = " ".join(result)
    return stored_text

def generate_response(query, context):
    prompt = f"""<|system|>
You are a medical report analyzer. Answer based ONLY on the report. If not found, say "I cannot find this in the report."

<|user|>
Report: {context}

Question: {query}

<|assistant|>"""
    
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048, padding=True)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=300,
            temperature=0.3,
            do_sample=True,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "<|assistant|>" in response:
        answer = response.split("<|assistant|>")[-1].strip()
    else:
        answer = response.strip()
    
    return answer[:1500]

def process_image(image):
    if image is None:
        return "Please upload an image.", ""
    
    text = extract_text(image)
    if not text or len(text.strip()) < 10:
        return "Could not extract text. Try a clearer image.", ""
    
    preview = text[:500] + "..." if len(text) > 500 else text
    return f"✅ Processed! Extracted {len(text)} characters.", preview

def ask_question(query):
    global stored_text
    if not query or not query.strip():
        return "Please enter a question."
    if not stored_text or len(stored_text.strip()) < 10:
        return "Please upload a medical report first."
    
    try:
        return generate_response(query, stored_text)
    except Exception as e:
        return f"Error: {str(e)}"

def clear_data():
    global stored_text
    stored_text = ""
    return "Cleared data. Upload a new report.", ""

# UI
with gr.Blocks(title="Medical Report Q&A", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏥 Medical Report Q&A Assistant")
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="Upload Report", type="numpy", height=250)
            process_btn = gr.Button("Process Report", variant="primary")
            status_output = gr.Textbox(label="Status", lines=2, interactive=False)
            clear_btn = gr.Button("Clear Data", variant="secondary")
        
        with gr.Column(scale=1):
            extracted_output = gr.Textbox(label="Extracted Text", lines=10, interactive=False)
    
    with gr.Row():
        with gr.Column(scale=2):
            query_input = gr.Textbox(
                label="Ask a Question",
                placeholder="e.g., What tests were performed?",
                lines=2
            )
            ask_btn = gr.Button("Ask", variant="primary")
        
        with gr.Column(scale=1):
            answer_output = gr.Textbox(label="Answer", lines=8, interactive=False)
    
    gr.Examples(
        examples=[
            ["What tests were performed?"],
            ["Which values look abnormal?"],
            ["Summarize this report."],
            ["What is the hemoglobin level?"]
        ],
        inputs=query_input
    )
    
    process_btn.click(process_image, [image_input], [status_output, extracted_output])
    ask_btn.click(ask_question, [query_input], [answer_output])
    query_input.submit(ask_question, [query_input], [answer_output])
    clear_btn.click(clear_data, [], [status_output, extracted_output])

if __name__ == "__main__":
    demo.launch()