import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("aaditya/Llama3-OpenBioLLM-8B")
model = AutoModelForCausalLM.from_pretrained(
    "aaditya/Llama3-OpenBioLLM-8B",
    torch_dtype=torch.bfloat16,  # Use bfloat16 for better memory efficiency
    device_map="auto"  # Automatically handle device placement
)

# Define the system prompt
SYSTEM_PROMPT = """You are an advanced medical documentation assistant analyzing clinical documentation. For each case, provide:
Documentation Quality Score (1-10):
- Assess completeness, clarity, and adherence to clinical standards
Diagnostic Confidence (0-100%):
- Evaluate certainty of current diagnosis
- Consider supporting evidence
Key Findings:
- List significant clinical observations
- Note critical vital signs
- Highlight abnormal results
Missing Information:
- Identify crucial missing data
- Note incomplete documentation areas
Recommended Actions:
- Suggest immediate clinical steps
- Propose management changes
Additional Tests:
- Recommend relevant diagnostics
- Suggest appropriate imaging
- Propose lab work
Safety Concerns:
- Flag potential drug interactions
- Highlight clinical red flags
- Note urgent safety issues"""

def analyze_medical_doc(medical_doc):
    # Prepare the prompt
    prompt = f"{SYSTEM_PROMPT}\n\nPlease analyze this medical documentation:\n\n{medical_doc}\n\nAnalysis:"
    
    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
    inputs = inputs.to(model.device)  # Move to same device as model
    
    # Generate response
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    
    # Decode and clean up the response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract only the generated part (remove the input prompt)
    response = response[len(prompt):]
    
    return response.strip()

# Example usage
sample_doc = """
Patient: 45-year-old male
Chief Complaint: Chest pain and shortness of breath
Vitals:
- BP: 145/90
- HR: 98
- RR: 20
- O2 Sat: 95% on RA

History:
- Onset: 2 hours ago
- Character: Sharp, radiating to left arm
- Previous MI: No
- HTN: Yes, on lisinopril

Current Medications:
- Lisinopril 10mg daily
"""

# Get the analysis
analysis = analyze_medical_doc(sample_doc)
print(analysis)