import fitz  # PyMuPDF
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

def extract_text_from_pdf(file_path):
    doc = fitz.open(file_path)
    text = ""
    for page in doc:
        text += page.get_text()
    return text

def analyze_document(file, prompt):
    # Determine file type and extract text
    if file.name.endswith('.pdf'):
        text = extract_text_from_pdf(file.name)
    elif file.name.endswith('.txt'):
        text = file.read().decode('utf-8')
    else:
        return "Unsupported file type. Please upload a PDF or TXT file."

    # Load the model and tokenizer
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer = AutoTokenizer.from_pretrained("THUDM/glm-4v-9b", trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        "THUDM/glm-4v-9b",
        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
        low_cpu_mem_usage=True,
        trust_remote_code=True
    ).to(device).eval()

    # Prepare inputs
    inputs = tokenizer.apply_chat_template([{"role": "user", "content": text + "\n\n" + prompt}],
                                           add_generation_prompt=True, tokenize=True, return_tensors="pt",
                                           return_dict=True)
    inputs = inputs.to(device)

    # Generate response
    gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
    with torch.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        return tokenizer.decode(outputs[0])

# Gradio interface
file_input = gr.File(label="Upload TXT or PDF Document", file_count="single")
prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your structured prompt here")
output_text = gr.Textbox(label="Analysis Result")

iface = gr.Interface(
    fn=analyze_document,
#    inputs=[
#        gr.inputs.File(label="Upload PDF or TXT Document"),
#        gr.inputs.Textbox(label="Enter your prompt")
#    ],
#    outputs="text",
    inputs=[file_input, prompt_input],
    outputs=output_text,
    title="Medical Report Analyzer",
    description="Upload a medical report (PDF or TXT) and enter a prompt to analyze the report."
)

if __name__ == "__main__":
    iface.launch(share=True)