src/agents/cdm_agents.py

import json
import time
from datetime import datetime, timedelta
from src.agents.agents import BaseAgent, throttle_agent
from src.core.state import AgentState
from src.tools.fhir_memory import get_observations_by_patient, get_medications_by_patient
from langchain_core.messages import SystemMessage, HumanMessage
from src.utils.logger import setup_logger

logger = setup_logger("CDMAgents")

class HealthCoach(BaseAgent):
    def __init__(self):
        fallback_prompt = """You are a proactive Chronic Disease Management (CDM) Health Coach.
        Your goal is to help patients manage their conditions (like Diabetes or Hypertension) through motivation, education, and lifestyle tracking.
        Always review their recent FHIR observations, current medications, and trend analysis to provide personalized advice.
        Consider the patient's medication regimen when making lifestyle recommendations.
        Be encouraging but firm about safety guidelines."""
        super().__init__(fallback_prompt, "HealthCoach.txt")

class TrendAnalyzer(BaseAgent):
    def __init__(self):
        fallback_prompt = """You are a medical data trend analyzer. 
        You analyze FHIR observations and identify clinically significant trends.
        If you see rising glucose levels or blood pressure, flag them immediately.
        Provide a concise summary of the last 7-14 days of data with trend projections."""
        super().__init__(fallback_prompt, "TrendAnalyzer.txt")

    def _compute_trend_statistics(self, values: list) -> dict:
        """
        Compute statistical trend indicators using polynomial regression.
        Bug 7.2: Enhanced trend analysis with proper windowing.
        
        Returns:
            dict: Contains trend direction, slope, R-squared, and projection
        """
        if len(values) < 2:
            return {"trend": "insufficient", "slope": 0, "r_squared": 0, "projection": None}
        
        try:
            import numpy as np
            
            # Prepare data points
            x = np.array(range(len(values)))
            y = np.array(values)
            
            # First-order polynomial fit (linear regression)
            if len(values) >= 3:
                coeffs = np.polyfit(x, y, 1)
                slope = coeffs[0]
                intercept = coeffs[1]
                
                # Calculate R-squared
                y_pred = np.polyval(coeffs, x)
                ss_res = np.sum((y - y_pred) ** 2)
                ss_tot = np.sum((y - np.mean(y)) ** 2)
                r_squared = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0
                
                # Determine trend with threshold
                if abs(slope) < 0.5:
                    trend = "stable"
                elif slope > 0.5:
                    trend = "rising"
                else:
                    trend = "falling"
                
                # Project next value
                next_x = len(values)
                projection = np.polyval(coeffs, next_x)
                
                return {
                    "trend": trend,
                    "slope": float(slope),
                    "r_squared": float(r_squared),
                    "projection": float(projection),
                    "change_rate": float(slope)
                }
            else:
                # For < 3 points, use simple delta
                slope = (values[-1] - values[0]) / (len(values) - 1)
                trend = "rising" if slope > 0.5 else "falling" if slope < -0.5 else "stable"
                return {
                    "trend": trend,
                    "slope": float(slope),
                    "r_squared": 0,
                    "projection": values[-1] + slope,
                    "change_rate": float(slope)
                }
        except ImportError:
            logger.warning("NumPy not available, using basic trend analysis")
            # Fallback: basic two-point comparison
            if len(values) >= 2:
                delta = values[-1] - values[-2]
                trend = "rising" if delta > 0.5 else "falling" if delta < -0.5 else "stable"
                return {
                    "trend": trend,
                    "slope": float(delta),
                    "r_squared": 0,
                    "projection": values[-1] + delta,
                    "change_rate": float(delta)
                }
            return {"trend": "insufficient", "slope": 0, "r_squared": 0, "projection": None}

    async def analyze_trends(self, patient_id: str):
        """
        Logic to analyze trends for a specific patient with windowing.
        Bug 7.2: Implements 7-14 day moving window analysis with regression.
        """
        logger.info(f"Analyzing health trends for patient: {patient_id}")
        observations = get_observations_by_patient.invoke({"patient_id": patient_id})
        if isinstance(observations, str):
            return observations, {}
        
        if not observations:
            return "No observations available for trend analysis.", {}
            
        # Group by LOINC code and sort by date
        data_points = {}
        for obs in observations:
            try:
                code = obs["code"]["coding"][0]["display"]
                val = float(obs["valueQuantity"]["value"])
                date = obs["effectiveDateTime"]
                if code not in data_points:
                    data_points[code] = []
                data_points[code].append({"value": val, "date": date})
            except (KeyError, IndexError, ValueError) as e:
                logger.warning(f"Error parsing observation: {e}")
                continue
        
        # Sort each metric's points by date (descending for most recent first)
        for code in data_points:
            data_points[code].sort(key=lambda x: x["date"], reverse=True)
        
        # Comprehensive trend analysis
        analysis = "**Trend Analysis Report (7-14 Day Window)**\n\n"
        structured_data = {}
        
        for code, points in data_points.items():
            if len(points) >= 2:
                # Extract values for statistical analysis
                values = [p["value"] for p in points]
                stats = self._compute_trend_statistics(values)
                
                latest = values[0]
                prev = values[1] if len(values) > 1 else values[0]
                delta = latest - prev
                
                analysis += f"**{code}**\n"
                analysis += f"- Latest: {latest:.1f}\n"
                analysis += f"- Trend: {stats['trend'].upper()} (slope: {stats['slope']:.2f}/day)\n"
                
                if stats['r_squared'] > 0:
                    analysis += f"- Trend Quality (R²): {stats['r_squared']:.2f}\n"
                
                if stats['projection'] is not None:
                    analysis += f"- Projected Next: {stats['projection']:.1f}\n"
                
                # Clinical alert for high-priority metrics
                if code.lower() in ["blood glucose", "glucose"] and latest > 250:
                    analysis += f"  ⚠️ **ALERT**: Glucose critically high\n"
                elif code.lower() in ["blood glucose", "glucose"] and latest < 70:
                    analysis += f"  ⚠️ **ALERT**: Risk of hypoglycemia\n"
                
                
                # Add to structured data
                structured_data[code] = {
                    "points": [{"date": p["date"], "value": p["value"]} for p in points],
                    "stats": stats
                }
                
                analysis += "\n"
            elif len(points) == 1:
                analysis += f"**{code}**: Single data point (latest: {points[0]['value']:.1f}) — insufficient for trend\n\n"
        
        return analysis, structured_data

    @throttle_agent
    async def run(self, state: AgentState, config=None, **kwargs):
        # Extract patient ID from state
        patient_id = state.get("patient_id", "unknown")
        if patient_id == "unknown":
            logger.warning("TrendAnalyzer: No patient ID found in state.")
            return {"logs": ["TrendAnalyzer: No patient ID found in state."]}
            
        start_time = time.time()
        analysis, structured_data = await self.analyze_trends(patient_id)
        end_time = time.time()
        
        metrics = {
            "agent": "TrendAnalyzer",
            "tokens": 0,  # Logic-based, no LLM call
            "time": round(end_time - start_time, 3)
        }
        
        return {
            "trend_analysis": analysis,
            "trend_data": structured_data,
            "logs": [f"TrendAnalyzer: Completed analysis for {patient_id}"],
            "metrics": [metrics]
        }