from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o", temperature=0)

@tool
def assess_symptoms(symptoms: str, age: int, gender: str, duration: str) -> str:
    """Assess patient symptoms and determine urgency level.
    
    Args:
        symptoms: Patient-reported symptoms
        age: Patient age
        gender: Patient gender
        duration: Symptom duration
    """
    prompt = f"""Assess these patient symptoms for urgency:
    
    Symptoms: {symptoms}
    Age: {age}, Gender: {gender}
    Duration: {duration}
    
    Provide urgency assessment and initial triage recommendation."""
    
    response = llm.invoke(prompt)
    return response.content

@tool
def generate_differential(symptoms: str, age: int, gender: str, history: str) -> str:
    """Generate differential diagnosis using Bayesian reasoning.
    
    Args:
        symptoms: Patient symptoms
        age: Patient age
        gender: Patient gender
        history: Medical history
    """
    prompt = f"""Generate differential diagnosis for:
    
    Symptoms: {symptoms}
    Patient: {age}yo {gender}
    History: {history}
    
    Use Bayesian reasoning based on prevalence and symptom presentation:"""
    
    response = llm.invoke(prompt)
    return response.content

@tool
def recommend_treatment(diagnosis: str, patient_age: int, allergies: str) -> str:
    """Recommend evidence-based treatment.
    
    Args:
        diagnosis: Primary diagnosis
        patient_age: Patient age
        allergies: Known drug allergies
    """
    prompt = f"""Recommend treatment for:
    
    Diagnosis: {diagnosis}
    Patient Age: {patient_age}
    Allergies: {allergies}
    
    Provide:
    1. First-line treatment
    2. Alternative treatments
    3. Medication dosing
    4. Follow-up plan
    5. Patient education points
    
    Follow current clinical guidelines:"""
    
    response = llm.invoke(prompt)
    return response.content

from typing import TypedDict, Annotated
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

class ClinicalState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    patient_data: dict
    urgency_level: str
    differential: str
    treatment: str
    follow_up: str

def triage_node(state: ClinicalState):
    """Assess patient urgency."""
    patient = state["patient_data"]
    result = assess_symptoms.invoke({
        "symptoms": patient.get("symptoms", ""),
        "age": patient.get("age", 0),
        "gender": patient.get("gender", ""),
        "duration": patient.get("duration", ""),
    })
    
    urgency = "Emergency" if any(w in result.lower() for w in ["emergency", "immediate", "critical"]) else \
              "Urgent" if any(w in result.lower() for w in ["urgent", "24 hours"]) else "Routine"
    
    return {"urgency_level": urgency, "messages": [HumanMessage(content=result)]}

def diagnosis_node(state: ClinicalState):
    """Generate differential diagnosis."""
    patient = state["patient_data"]
    result = generate_differential.invoke({
        "symptoms": patient.get("symptoms", ""),
        "age": patient.get("age", 0),
        "gender": patient.get("gender", ""),
        "history": patient.get("history", ""),
    })
    return {"differential": result, "messages": [HumanMessage(content=result)]}

def treatment_node(state: ClinicalState):
    """Recommend treatment."""
    patient = state["patient_data"]
    result = recommend_treatment.invoke({
        "diagnosis": state["differential"][:500],
        "patient_age": patient.get("age", 0),
        "allergies": patient.get("allergies", "None known"),
    })
    return {"treatment": result, "messages": [HumanMessage(content=result)]}

def emergency_check(state: ClinicalState):
    """Route based on urgency."""
    if state["urgency_level"] == "Emergency":
        return "emergency_alert"
    return "diagnosis"

def emergency_alert(state: ClinicalState):
    """Handle emergency cases."""
    return {"messages": [HumanMessage(content="EMERGENCY: Activate emergency protocols. Call code team.")]}

# Build graph
workflow = StateGraph(ClinicalState)
workflow.add_node("triage", triage_node)
workflow.add_node("emergency_alert", emergency_alert)
workflow.add_node("diagnosis", diagnosis_node)
workflow.add_node("treatment", treatment_node)

workflow.set_entry_point("triage")
workflow.add_conditional_edges("triage", emergency_check, {
    "emergency_alert": "emergency_alert",
    "diagnosis": "diagnosis",
})
workflow.add_edge("emergency_alert", END)
workflow.add_edge("diagnosis", "treatment")
workflow.add_edge("treatment", END)

app = workflow.compile()

# Run
result = app.invoke({
    "messages": [HumanMessage(content="Patient assessment")],
    "patient_data": {
        "age": 55,
        "gender": "Male",
        "symptoms": "Chest pain, shortness of breath, sweating",
        "duration": "2 hours",
        "history": "Hypertension, diabetes, smoker",
        "allergies": "Penicillin",
    },
    "urgency_level": "",
    "differential": "",
    "treatment": "",
    "follow_up": "",
})
print(result["urgency_level"])
print(result["differential"][:500])

from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import StateGraph, END

class ResearchState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    topic: str
    search_queries: List[str]
    papers: List[dict]
    evidence_summary: str
    grade_assessment: str

@tool
def search_pubmed(query: str, max_results: int = 10) -> str:
    """Search PubMed for relevant medical literature.
    
    Args:
        query: Search query
        max_results: Maximum number of results
    """
    import requests
    params = {"db": "pubmed", "term": query, "retmax": max_results, "retmode": "json"}
    resp = requests.get("https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi", params=params)
    return resp.json().get("esearchresult", {}).get("idlist", [])

@tool
def grade_evidence(study_type: str, sample_size: int, p_value: float, bias_risk: str) -> str:
    """Apply GRADE framework to assess evidence quality.
    
    Args:
        study_type: Type of study (RCT, cohort, case-control, etc.)
        sample_size: Number of participants
        p_value: Statistical significance
        bias_risk: Risk of bias assessment
    """
    grade_map = {
        "RCT": "High" if bias_risk == "low" else "Moderate",
        "cohort": "Moderate" if bias_risk == "low" else "Low",
        "case-control": "Low",
        "case-series": "Very Low",
    }
    base_grade = grade_map.get(study_type, "Very Low")
    if p_value > 0.05:
        base_grade = "Low" if base_grade == "Moderate" else "Very Low"
    return f"GRADE: {base_grade} | Study: {study_type} | n={sample_size} | p={p_value} | Bias: {bias_risk}"

from typing import TypedDict, List
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import StateGraph, END

class DrugDiscoveryState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    target_protein: str
    candidate_molecules: List[dict]
    binding_scores: List[float]
    prioritized_compounds: List[dict]

@tool
def screen_molecules(target: str, library: str, top_n: int = 20) -> str:
    """Screen molecular library against protein target.
    
    Args:
        target: Target protein name
        library: Molecular library identifier
        top_n: Number of top compounds to return
    """
    return f"Screened {library} against {target}. Returning top {top_n} compounds by binding affinity."

@tool
def predict_admet(smiles: str) -> str:
    """Predict ADMET properties for a molecule.
    
    Args:
        smiles: SMILES notation of molecule
    """
    return f"ADMET for {smiles}: Absorption=High, Distribution=Moderate, Metabolism=CYP3A4, Excretion=Renal, Toxicity=Low"

from typing import TypedDict, List
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import StateGraph, END

class MonitoringState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    patient_id: str
    vitals: dict
    lab_results: dict
    alert_level: str
    clinical_notes: str

@tool
def analyze_vitals(heart_rate: int, blood_pressure: str, temperature: float, spO2: int) -> str:
    """Analyze vital signs for deterioration patterns.
    
    Args:
        heart_rate: Heart rate in bpm
        blood_pressure: Blood pressure reading
        temperature: Body temperature in Celsius
        spO2: Oxygen saturation percentage
    """
    alerts = []
    if heart_rate > 120 or heart_rate < 50:
        alerts.append("Abnormal heart rate")
    if spO2 < 92:
        alerts.append("Hypoxemia detected")
    if temperature > 38.5:
        alerts.append("Fever - possible infection")
    return "ALERTS: " + "; ".join(alerts) if alerts else "Vitals within normal limits"

@tool
def trend_analysis(patient_id: str, metric: str, window_hours: int = 24) -> str:
    """Analyze trends in patient metrics over time.
    
    Args:
        patient_id: Patient identifier
        metric: Metric to analyze
        window_hours: Analysis window in hours
    """
    return f"Trend analysis for {patient_id} - {metric} over {window_hours}h: Stable with slight improvement"

import os
from cryptography.fernet import Fernet

class HIPAACompliantAgent:
    def __init__(self):
        self.encryption_key = Fernet.generate_key()
        self.cipher = Fernet(self.encryption_key)
        self.audit_log = []
    
    def encrypt_phi(self, data: str) -> bytes:
        """Encrypt Protected Health Information."""
        return self.cipher.encrypt(data.encode())
    
    def decrypt_phi(self, encrypted: bytes) -> str:
        """Decrypt Protected Health Information."""
        return self.cipher.decrypt(encrypted).decode()
    
    def log_access(self, user_id: str, patient_id: str, action: str):
        """Log all PHI access for audit trail."""
        import datetime
        self.audit_log.append({
            "timestamp": datetime.datetime.now().isoformat(),
            "user": user_id,
            "patient": patient_id,
            "action": action
        })
    
    def deidentify(self, record: dict) -> dict:
        """Remove 18 HIPAA identifiers from record."""
        identifiers = ["name", "address", "dates", "phone", "fax", "email",
                       "ssn", "mrn", "health_plan", "account", "certificate",
                       "vehicle", "device", "url", "ip", "biometric", "photo", "any_id"]
        return {k: v for k, v in record.items() if k.lower() not in identifiers}

# Install dependencies
pip install langchain langchain-openai langgraph langchain-core

# Set API keys
export OPENAI_API_KEY="your-key"

# For healthcare-specific tools
pip install fhirclient  # FHIR EHR integration
pip install rdkit        # Molecular processing for drug discovery
pip install biopython    # PubMed/literature access