src/agents/biomarker_analyzer.py

"""
MediGuard AI RAG-Helper
Biomarker Analyzer Agent - Validates biomarker values and flags anomalies
"""

from src.biomarker_validator import BiomarkerValidator
from src.llm_config import llm_config
from src.state import AgentOutput, BiomarkerFlag, GuildState


class BiomarkerAnalyzerAgent:
    """Agent that validates biomarker values and generates comprehensive analysis"""

    def __init__(self):
        self.validator = BiomarkerValidator()
        self.llm = llm_config.analyzer

    def analyze(self, state: GuildState) -> GuildState:
        """
        Main agent function to analyze biomarkers.

        Args:
            state: Current guild state with patient input

        Returns:
            Updated state with biomarker analysis
        """
        print("\n" + "=" * 70)
        print("EXECUTING: Biomarker Analyzer Agent")
        print("=" * 70)

        biomarkers = state["patient_biomarkers"]
        patient_context = state.get("patient_context", {})
        gender = patient_context.get("gender")  # None if not provided — uses non-gender-specific ranges
        predicted_disease = state["model_prediction"]["disease"]

        # Validate all biomarkers
        print(f"\nValidating {len(biomarkers)} biomarkers...")
        flags, alerts = self.validator.validate_all(
            biomarkers=biomarkers, gender=gender, threshold_pct=state["sop"].biomarker_analyzer_threshold
        )

        # Get disease-relevant biomarkers
        relevant_biomarkers = self.validator.get_disease_relevant_biomarkers(predicted_disease)

        # Generate summary using LLM
        summary = self._generate_summary(biomarkers, flags, alerts, relevant_biomarkers, predicted_disease)

        findings = {
            "biomarker_flags": [flag.model_dump() for flag in flags],
            "safety_alerts": [alert.model_dump() for alert in alerts],
            "relevant_biomarkers": relevant_biomarkers,
            "summary": summary,
            "validation_complete": True,
        }

        # Create agent output
        output = AgentOutput(agent_name="Biomarker Analyzer", findings=findings)

        # Update state
        print("\nAnalysis complete:")
        print(f"  - {len(flags)} biomarkers validated")
        print(f"  - {len([f for f in flags if f.status != 'NORMAL'])} out-of-range values")
        print(f"  - {len(alerts)} safety alerts generated")
        print(f"  - {len(relevant_biomarkers)} disease-relevant biomarkers identified")

        return {
            "agent_outputs": [output],
            "biomarker_flags": flags,
            "safety_alerts": alerts,
            "biomarker_analysis": findings,
        }

    def _generate_summary(
        self,
        biomarkers: dict[str, float],
        flags: list[BiomarkerFlag],
        alerts: list,
        relevant_biomarkers: list[str],
        disease: str,
    ) -> str:
        """Generate a concise summary of biomarker findings"""

        # Count anomalies
        critical = [f for f in flags if "CRITICAL" in f.status]
        high_low = [f for f in flags if f.status in ["HIGH", "LOW"]]

        prompt = f"""You are a medical data analyst. Provide a brief, clinical summary of these biomarker results.

**Patient Context:**
- Predicted Disease: {disease}
- Total Biomarkers Tested: {len(biomarkers)}
- Critical Values: {len(critical)}
- Out-of-Range Values: {len(high_low)}

**Key Findings:**
{self._format_key_findings(critical, high_low, relevant_biomarkers)}

Provide a 2-3 sentence summary highlighting:
1. Overall risk profile
2. Most concerning findings
3. Alignment with predicted disease

Keep it concise and clinical."""

        try:
            response = self.llm.invoke(prompt)
            return response.content.strip()
        except Exception as e:
            print(f"Warning: LLM summary generation failed: {e}")
            return f"Biomarker analysis complete. {len(critical)} critical values, {len(high_low)} out-of-range values detected."

    def _format_key_findings(self, critical, high_low, relevant):
        """Format findings for LLM prompt"""
        findings = []

        if critical:
            findings.append("CRITICAL VALUES:")
            for f in critical[:3]:  # Top 3
                findings.append(f"  - {f.name}: {f.value} {f.unit} ({f.status})")

        if high_low:
            findings.append("\nOUT-OF-RANGE VALUES:")
            for f in high_low[:5]:  # Top 5
                findings.append(f"  - {f.name}: {f.value} {f.unit} ({f.status})")

        if relevant:
            findings.append(f"\nDISEASE-RELEVANT BIOMARKERS: {', '.join(relevant[:5])}")

        return "\n".join(findings) if findings else "All biomarkers within normal range."


# Create agent instance for import
biomarker_analyzer_agent = BiomarkerAnalyzerAgent()