app.py

# Provider side auth
import streamlit as st
from transformers import pipeline
import requests
import time
import uuid
import os
import torch


# 1. Page Config
st.set_page_config(layout="wide", page_title="Smart Prior Authorization Prototype")

# --- LANDING SCREEN HEADER ---
st.title("🏥 Smart Prior Authorization")
st.markdown("### **Next-Gen Agentic Prior-Authorization Prototype**")

# 1. Quick Technical "Badges"
cols = st.columns(4)
cols[0].metric("Provider Agent", "Llama-3.1-8B")
cols[1].metric("Payer Agent", "GPT-4o")
cols[2].metric("NLP Engine", "BioMed-NER")
cols[3].metric("Protocol", "FHIR-Ready")

st.info("💡 **Prototype Goal:** Automate the Clinical Negotiation between Providers and Payers using multi-agent reasoning to close documentation gaps in real-time.")

# 2. Collapsible Deep Dive
with st.expander("🛠️ Technical Overview & Architecture", expanded=False):
    col_a, col_b = st.columns(2)
    
    with col_a:
        st.markdown("""
        **How it Works:**
        1. **Extraction:** A BERT-based NER model identifies clinical entities in your note.
        2. **Reasoning:** The Provider Agent (Llama-3.1) builds a 'case' for medical necessity.
        3. **Adjudication:** The Payer Agent (GPT-4o) evaluates the request against specific LCD/NCD policies.
        """)
        
    with col_b:
        st.markdown("""
        **Key Innovation:**
        - **Zero-Touch Auth:** Reduces manual faxing and portal entry.
        - **Inner Monologue:** View the 'Agent Trace' to see exactly why a procedure was approved or pended.
        - **Dynamic Discovery:** Designed to scale toward universal CPT-to-Policy mapping.
        """)

st.divider()
# --- END LANDING SCREEN HEADER ---

# 2. MEDICAL POLICY DATABASE
POLICY_DB = {
    "22558 - Spinal Fusion": {
        "required": ["Physical Therapy", "Instability", "Spondylolisthesis"],
        "policy_id": "LCD-L341",
        "cpt": "22558"
    },
    "43239 - EGD": {
        "required": ["Dysphagia", "GERD", "Weight Loss"],
        "policy_id": "MED-772",
        "cpt": "43239"
    }
}

# 3. Load NLP Model
@st.cache_resource
def load_agent():
    return pipeline("token-classification", 
                    model="Helios9/BIOMed_NER", 
                    aggregation_strategy="simple")

nlp_agent = load_agent()

# 4. SEMANTIC PIPELINE LOGIC
def semantic_matcher(note_text, raw_entities):
    confirmed_evidence = []
    negation_triggers = ["no", "denies", "negative", "without", "none", "not", "rule out"]
    
    for ent in raw_entities:
        word = ent['word'].lower()
        start_char = ent['start']
        context_window = note_text[max(0, start_char-30):start_char].lower()
        is_negated = any(neg in context_window for neg in negation_triggers)
        
        if not is_negated:
            confirmed_evidence.append(word)
    return confirmed_evidence

# 5. AGENTIC SETUP 
# Switching to Llama-3.1-8B which is fully supported by the HF Router fleet
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"

hf_token = os.environ.get("HFPATOKEN")
HEADERS = {
    "Authorization": f"Bearer {hf_token}",
    "Content-Type": "application/json"
}

def provider_agent_reasoning(clinical_note, payer_question):
    # Unified Router Endpoint
    API_URL = "https://router.huggingface.co/v1/chat/completions"
    
    payload = {
        "model": MODEL_ID, 
        "messages": [
            {
                "role": "system", 
                "content": (
                    "You are a Clinical Reviewer at Oracle Health. "
                    "Analyze the clinical note to find evidence for medical necessity. "
                    "Format: THOUGHT: [reasoning] RESPONSE: [answer]"
                )
            },
            {
                "role": "user", 
                "content": f"CLINICAL NOTE: {clinical_note}\n\nPAYER QUESTION: {payer_question}"
            }
        ],
        "temperature": 0.1,
        "max_tokens": 500
    }

    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=20)
        
        if response.status_code == 200:
            result = response.json()
            answer = result['choices'][0]['message']['content']
            return "Llama-3.1 Reasoning", answer
        
        # This will show you exactly why the router is rejecting it
        return "System Error", f"HF Router Error {response.status_code}: {response.text}"
        
    except Exception as e:
        return "Connection Error", str(e)

def call_payer_api(payload):

    # UPDATE THIS URL to your actual Payer Space URL

    PAYER_URL = "https://pilayar-payer-auth-sim.hf.space/agent-chat"

    try:

        response = requests.post(PAYER_URL, json=payload, timeout=15)

        return response.json()

    except:

        return {"agent_response": "Payer Agent is currently unavailable."}

def render_summary(inner_monologue, last_payer_msg):
    st.subheader("🏁 Negotiation Outcome")
    
    col_left, col_right = st.columns(2)
    
    with col_left:
        st.error("❌ **Decision: Action Required**")
        st.markdown(f"**Missing Criteria:** \n {last_payer_msg}")
        
    with col_right:
        st.success("📝 **Next Steps for Provider**")
        # Extract the 'Recommendation' part from the Agent's thought process
        st.info("The Agent recommends updating the EHR with: \n 1. PT Logs (6+ Weeks) \n 2. Flexion/Extension X-rays for Instability")

    # Dynamic "Draft" Button
    if st.button("➕ Generate Addendum for Clinician"):
        st.text_area("Proposed Addendum", value=f"Addendum: {inner_monologue.split('RESPONSE:')[-1]}")

# 6. UI LAYOUT
st.title("Smart Auth Agentic Prototype")

col1, col2 = st.columns([2, 1])

with col2:
    st.subheader("Provider Documentation")
    note = st.text_area("Clinical Note", height=400, placeholder="Paste HPI here...")

with col1:
    st.subheader("Order Workflow")
    order = st.selectbox("Select Procedure", ["", "22558 - Spinal Fusion", "43239 - EGD"])
    
    if order and note:
        policy = POLICY_DB[order]
        
        # Display Preliminary Analysis
        with st.status("🧠 Analyzing Medical Necessity...") as status:
            raw_entities = nlp_agent(note)
            valid_evidence = semantic_matcher(note, raw_entities)
            status.update(label="✅ Analysis Complete", state="complete")

        if st.button("🚀 Start Agent Negotiation"):
                    # --- 1. INITIALIZE STATE (Prevents NameErrors) ---
                    current_message = f"Initial Request: {order} for patient. Note Summary: {note[:200]}"
                    thought = ""
                    payer_text = ""
                    outcome_status = "IN_PROGRESS"
                    
                    # --- 2. THE NEGOTIATION LOOP ---
                    for i in range(3):
                        # A. PROVIDER -> PAYER
                        with st.chat_message("user", avatar="🏥"):
                            st.write(f"**Provider Agent:** {current_message}")
                        
                        # B. PAYER RESPONSE
                        with st.spinner("Payer is adjudicating..."):
                            payer_response = call_payer_api({"message": current_message})
                            payer_text = payer_response.get("agent_response", "Payer connection timeout.")
                        
                        with st.chat_message("assistant", avatar="🏦"):
                            st.write(f"**Payer Agent:** {payer_text}")
        
                        # Check for immediate approval
                        if "APPROVED" in payer_text.upper():
                            outcome_status = "APPROVED"
                            break
                        
                        # C. PROVIDER REASONING
                        with st.spinner("Provider Agent searching records..."):
                            thought, current_message = provider_agent_reasoning(note, payer_text)
                            with st.expander("🔍 Agent Trace: Inner Monologue", expanded=False):
                                st.write(f"**Reasoning:** {thought}")
        
                        # D. GAP DETECTION (Graceful Exit)
                        gaps = ["missing documentation", "please provide", "not explicitly documented", "additional documentation"]
                        if any(phrase in thought.lower() for phrase in gaps):
                            outcome_status = "PENDED"
                            break
                        
                        time.sleep(1)
        
                    # --- 3. FINAL OUTCOME DASHBOARD (Succinct Summary) ---
                    st.divider()
                    
                    if outcome_status == "APPROVED":
                        st.balloons()
                        st.success("### ✅ Final Outcome: PRIOR AUTHORIZATION APPROVED")
                        st.markdown(f"**Policy Confirmed:** {POLICY_DB[order]['policy_id']}")
                        st.write("The clinical evidence provided meets the gold-standard criteria for this procedure.")
                    
                    else:
                        st.error("### ⏳ Final Outcome: PENDED / ACTION REQUIRED")
                        
                        # Create a clean side-by-side summary
                        sum_col1, sum_col2 = st.columns(2)
                        
                        with sum_col1:
                            st.markdown("**🚨 Missing Clinical Evidence:**")
                            # Dynamically clean up the Payer's text to show only the requirements
                            clean_requirements = payer_text.split("please provide")[-1] if "please provide" in payer_text.lower() else "See Payer message for specific details."
                            st.info(clean_requirements)
                        
                        with sum_col2:
                            st.markdown("**📋 Suggested Clinical Addendum:**")
                            # Isolate the Provider's recommendation from the 'thought'
                            provider_rec = thought.split("RESPONSE:")[-1] if "RESPONSE:" in thought else "Please update the note with Physical Therapy and Imaging specifics."
                            st.success(provider_rec)
        
                        # Professional Call to Action
                        st.button("Update EHR & Resubmit", type="primary")