gnn.py

import os
from pathlib import Path
from typing import Any, Dict, Optional
import statistics

import librosa
import numpy as np

import torch.nn as nn


class ImprovedRespiratoryGAT(nn.Module):

    def __init__(self, input_dim=768, hidden_dim=256, num_layers=4, num_heads=4, dropout=0.4):
        import torch
        import torch.nn as nn
        from torch_geometric.nn import GATConv, SAGEConv, BatchNorm
        super().__init__()
        self.input_proj = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.gat_layers = nn.ModuleList()
        self.norms = nn.ModuleList()
        for i in range(num_layers):
            if i % 2 == 0:
                conv = GATConv(hidden_dim, max(1, hidden_dim // num_heads), heads=num_heads, concat=True, dropout=dropout)
            else:
                conv = SAGEConv(hidden_dim, hidden_dim)
            self.gat_layers.append(conv)
            self.norms.append(BatchNorm(hidden_dim))

        self.attention_pool = nn.Sequential(
            nn.Linear(hidden_dim, max(4, hidden_dim // 4)),
            nn.Tanh(),
            nn.Linear(max(4, hidden_dim // 4), 1),
        )
        self.wheeze_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.LayerNorm(hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )
        self.crackle_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.LayerNorm(hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )
        self.log_var_wheeze = nn.Parameter(torch.tensor(0.0))
        self.log_var_crackle = nn.Parameter(torch.tensor(0.0))
        self.dropout = dropout

    def forward(self, data):
        import torch
        import torch.nn.functional as F
        x, edge_index = data.x, data.edge_index
        batch = data.batch if hasattr(data, "batch") else None

        x = x.to(self.input_proj[0].weight.device)
        edge_index = edge_index.to(self.input_proj[0].weight.device)

        x = self.input_proj(x)
        residuals = []
        for i, (conv, norm) in enumerate(zip(self.gat_layers, self.norms)):
            x_new = conv(x, edge_index)
            x_new = F.elu(x_new)
            x_new = norm(x_new)
            if i > 0 and i % 2 == 0:
                x_new = x_new + residuals[-1]
            x = F.dropout(x_new, p=self.dropout, training=self.training)
            residuals.append(x)

        if batch is not None:
            batch = batch.to(x.device)
            attn_scores = self.attention_pool(x).squeeze(-1)
            x_graph = []
            for b in torch.unique(batch):
                mask = batch == b
                scores = attn_scores[mask]
                weights = torch.softmax(scores, dim=0).unsqueeze(-1)
                x_graph.append((x[mask] * weights).sum(dim=0))
            x = torch.stack(x_graph, dim=0)
        else:
            attn_scores = self.attention_pool(x).squeeze(-1)
            weights = torch.softmax(attn_scores, dim=0).unsqueeze(-1)
            x = (x * weights).sum(dim=0, keepdim=True)

        w_logits = self.wheeze_head(x).squeeze(-1)
        c_logits = self.crackle_head(x).squeeze(-1)
        return w_logits, c_logits


def load_gnn_model(checkpoint_path: str, device: Any):
    import torch
    model = ImprovedRespiratoryGAT(input_dim=768, hidden_dim=256, num_layers=4, num_heads=4, dropout=0.4)
    model = model.to(device)
    if os.path.exists(checkpoint_path):
        state_dict = torch.load(checkpoint_path, map_location=device)
        if isinstance(state_dict, dict):
            if "model_state_dict" in state_dict:
                state_dict = state_dict["model_state_dict"]
            elif "state_dict" in state_dict:
                state_dict = state_dict["state_dict"]
        new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
        model.load_state_dict(new_state_dict, strict=False)
    else:
        print(f"Warning: GNN checkpoint '{checkpoint_path}' not found. Using randomly initialized model.")
    model.eval()
    return model


def build_chain_edge_index(n: int):
    import torch
    if n <= 1:
        return torch.empty((2, 0), dtype=torch.long)
    edges = [[i, i + 1] for i in range(n - 1)] + [[i + 1, i] for i in range(n - 1)]
    return torch.tensor(edges, dtype=torch.long).t().contiguous()


def estimate_breathing_rate_bpm(y: np.ndarray, sr: int, audio_duration_s: float) -> Optional[float]:
    if len(y) < sr:
        return None
    frame_len = int(0.2 * sr)
    hop_len = int(0.05 * sr)
    rms = librosa.feature.rms(y=y, frame_length=frame_len, hop_length=hop_len)[0]
    if len(rms) == 0:
        return None
    win = 5
    smooth = np.convolve(rms, np.ones(win) / win, mode="same") if len(rms) >= win else rms
    thr = float(smooth.mean() + 0.5 * smooth.std())
    times = librosa.frames_to_time(np.arange(len(smooth)), sr=sr, hop_length=hop_len)
    min_interval = 0.8
    peaks = []
    last_t = -1e9
    for i in range(1, len(smooth) - 1):
        if smooth[i] > smooth[i - 1] and smooth[i] >= smooth[i + 1] and smooth[i] > thr:
            t = float(times[i])
            if t - last_t >= min_interval:
                peaks.append(t)
                last_t = t
    if len(peaks) < 2:
        return None
    return float((len(peaks) / max(audio_duration_s, 1e-6)) * 60.0)


class PatientStateManager:
    def __init__(
        self,
        ema_alpha: float = 0.12,
        low_delta: float = 0.08,
        high_delta: float = 0.20,
        min_samples_for_baseline: int = 5,
        force_established_after_s: float = 10.0,
    ):
        self.ema_alpha = ema_alpha
        self.low_delta = low_delta
        self.high_delta = high_delta
        self.min_samples_for_baseline = min_samples_for_baseline
        self.force_established_after_s = force_established_after_s
        self.patient_data: Dict[str, Dict] = {}

    def update_and_get_state(
        self,
        patient_id: str,
        wheeze_prob: float,
        crackle_prob: float,
        timestamp: float = 0.0,
    ) -> Dict:
        if patient_id not in self.patient_data:
            self.patient_data[patient_id] = {
                "wheeze_ema": wheeze_prob,
                "crackle_ema": crackle_prob,
                "wheeze_baseline": None,
                "crackle_baseline": None,
                "wheeze_history": [],
                "crackle_history": [],
                "timestamps": [],
                "count": 0,
                "baseline_established": False,
                "breathing_rate_history": [],
            }

        data = self.patient_data[patient_id]
        data["timestamps"].append(timestamp)
        data["wheeze_history"].append(wheeze_prob)
        data["crackle_history"].append(crackle_prob)
        data["count"] += 1

        if data["count"] == 1:
            data["wheeze_ema"] = wheeze_prob
            data["crackle_ema"] = crackle_prob
        else:
            data["wheeze_ema"] = self.ema_alpha * wheeze_prob + (1 - self.ema_alpha) * data["wheeze_ema"]
            data["crackle_ema"] = self.ema_alpha * crackle_prob + (1 - self.ema_alpha) * data["crackle_ema"]

        if data["count"] >= self.min_samples_for_baseline and not data["baseline_established"]:
            data["wheeze_baseline"] = float(np.mean(data["wheeze_history"][-self.min_samples_for_baseline :]))
            data["crackle_baseline"] = float(np.mean(data["crackle_history"][-self.min_samples_for_baseline :]))
            data["baseline_established"] = True

        result = {"overall_state": "establishing", "reason": {}}
        for axis, ema, baseline, history in [
            ("wheeze", data["wheeze_ema"], data["wheeze_baseline"], data["wheeze_history"]),
            ("crackle", data["crackle_ema"], data["crackle_baseline"], data["crackle_history"]),
        ]:
            if not data["baseline_established"] or baseline is None:
                state = "establishing"
                delta = 0.0
                trend = 0.0
            else:
                delta = float(ema - baseline)
                trend = float(history[-1] - history[-3]) if len(history) >= 3 else 0.0
                if abs(delta) < self.low_delta:
                    state = "green"
                elif abs(delta) < self.high_delta:
                    state = "orange"
                else:
                    state = "red"

            result["reason"][axis] = {
                "baseline": baseline,
                "value": ema,
                "delta": delta,
                "trend": trend,
                "state": state,
            }

        if timestamp >= self.force_established_after_s and not data["baseline_established"] and data["count"] > 0:
            data["baseline_established"] = True
            data["wheeze_baseline"] = float(np.mean(data["wheeze_history"]))
            data["crackle_baseline"] = float(np.mean(data["crackle_history"]))
            for axis in ["wheeze", "crackle"]:
                result["reason"][axis]["baseline"] = data[f"{axis}_baseline"]
                result["reason"][axis]["state"] = "green"

        for s in ["red", "orange", "green", "establishing"]:
            if s in [result["reason"][axis]["state"] for axis in ["wheeze", "crackle"]]:
                result["overall_state"] = s
                break

        result["count"] = data["count"]
        return result


class ClinicalReferenceRanges:
    RESPIRATORY_RATES = {
        (0, 1): (30, 60),
        (1, 2): (24, 40),
        (2, 6): (22, 34),
        (6, 12): (18, 30),
        (12, 18): (12, 20),
        (18, 65): (12, 20),
        (65, 150): (12, 28),
    }
    SEX_ADJUSTMENT = {"male": 0.0, "female": 2.0}
    RESPIRATORY_SEVERITY = {
        "bradypnea": {"threshold": 8, "severity": "red"},
        "low_normal": {"threshold": 12, "severity": "green"},
        "high_normal": {"threshold": 20, "severity": "green"},
        "tachypnea_mild": {"threshold": 24, "severity": "orange"},
        "tachypnea_moderate": {"threshold": 28, "severity": "orange"},
        "tachypnea_severe": {"threshold": 30, "severity": "red"},
    }
    WHEEZE_THRESHOLDS = {
        "normal": 0.30,
        "borderline": 0.45,
        "abnormal": 0.60,
        "severe": 0.75,
    }
    CRACKLE_THRESHOLDS = {
        "normal": 0.35,
        "borderline": 0.50,
        "abnormal": 0.65,
        "severe": 0.80,
    }

    @classmethod
    def get_normal_breathing_range(cls, age_years: float, sex: str = "male"):
        for (min_age, max_age), (low, high) in cls.RESPIRATORY_RATES.items():
            if min_age <= age_years <= max_age:
                if age_years >= 18:
                    adj = cls.SEX_ADJUSTMENT.get(str(sex).lower(), 0.0)
                    return (low + adj, high + adj)
                return (low, high)
        return (12, 20)

    @classmethod
    def assess_respiratory_rate(cls, rate_bpm: float, age_years: float, sex: str = "male"):
        normal_low, normal_high = cls.get_normal_breathing_range(age_years, sex)
        if rate_bpm < cls.RESPIRATORY_SEVERITY["bradypnea"]["threshold"]:
            severity = "red"
            status = "Severe bradypnea (dangerously low breathing rate)"
            clinical_action = "Immediate clinical review required"
        elif rate_bpm < normal_low:
            severity = "orange"
            status = f"Mild bradypnea (below normal range {normal_low}-{normal_high})"
            clinical_action = "Monitor closely, consider clinical assessment"
        elif rate_bpm <= normal_high:
            severity = "green"
            status = f"Normal respiratory rate ({normal_low}-{normal_high} bpm for age/sex)"
            clinical_action = "Routine monitoring"
        elif rate_bpm < cls.RESPIRATORY_SEVERITY["tachypnea_mild"]["threshold"]:
            severity = "orange"
            status = f"Mild tachypnea (above normal range {normal_low}-{normal_high})"
            clinical_action = "Monitor, assess for underlying cause"
        elif rate_bpm < cls.RESPIRATORY_SEVERITY["tachypnea_moderate"]["threshold"]:
            severity = "orange"
            status = "Moderate tachypnea"
            clinical_action = "Clinical assessment recommended"
        else:
            severity = "red"
            status = "Severe tachypnea (significant respiratory distress)"
            clinical_action = "Urgent clinical review required"

        return {
            "value": float(rate_bpm),
            "normal_range": [float(normal_low), float(normal_high)],
            "status": status,
            "severity": severity,
            "clinical_action": clinical_action,
        }

    @classmethod
    def assess_adventitious_sounds(cls, wheeze_prob: float, crackle_prob: float):
        if wheeze_prob < cls.WHEEZE_THRESHOLDS["normal"]:
            wheeze_severity = "green"
            wheeze_status = "Normal - no significant wheeze detected"
        elif wheeze_prob < cls.WHEEZE_THRESHOLDS["borderline"]:
            wheeze_severity = "green"
            wheeze_status = "Borderline - subtle wheeze, monitor"
        elif wheeze_prob < cls.WHEEZE_THRESHOLDS["abnormal"]:
            wheeze_severity = "orange"
            wheeze_status = "Abnormal - clinically significant wheeze"
        else:
            wheeze_severity = "red"
            wheeze_status = "Severe - prominent wheeze, indicates airway obstruction"

        if crackle_prob < cls.CRACKLE_THRESHOLDS["normal"]:
            crackle_severity = "green"
            crackle_status = "Normal - no significant crackles detected"
        elif crackle_prob < cls.CRACKLE_THRESHOLDS["borderline"]:
            crackle_severity = "green"
            crackle_status = "Borderline - fine crackles, monitor"
        elif crackle_prob < cls.CRACKLE_THRESHOLDS["abnormal"]:
            crackle_severity = "orange"
            crackle_status = "Abnormal - clinically significant crackles"
        else:
            crackle_severity = "red"
            crackle_status = "Severe - prominent crackles, indicates interstitial pathology"

        severity_rank = {"green": 0, "orange": 1, "red": 2}
        overall_severity = "green"
        for s in (wheeze_severity, crackle_severity):
            if severity_rank.get(s, 0) > severity_rank.get(overall_severity, 0):
                overall_severity = s

        return {
            "wheeze": {
                "probability": float(wheeze_prob),
                "severity": wheeze_severity,
                "status": wheeze_status,
            },
            "crackle": {
                "probability": float(crackle_prob),
                "severity": crackle_severity,
                "status": crackle_status,
            },
            "overall_severity": overall_severity,
        }


class EnhancedPatientStateManager(PatientStateManager):
    def __init__(
        self,
        ema_alpha: float = 0.12,
        low_delta: float = 0.08,
        high_delta: float = 0.20,
        min_samples_for_baseline: int = 5,
        force_established_after_s: float = 10.0,
        patient_age: Optional[float] = None,
        patient_sex: Optional[str] = None,
    ):
        super().__init__(ema_alpha, low_delta, high_delta, min_samples_for_baseline, force_established_after_s)
        self.patient_age = patient_age
        self.patient_sex = patient_sex
        self.clinical_kb = ClinicalReferenceRanges()

    def set_patient_demographics(self, patient_id: str, age_years: float, sex: str) -> None:
        if patient_id not in self.patient_data:
            self.patient_data[patient_id] = {
                "wheeze_ema": 0.0,
                "crackle_ema": 0.0,
                "wheeze_baseline": None,
                "crackle_baseline": None,
                "wheeze_history": [],
                "crackle_history": [],
                "timestamps": [],
                "count": 0,
                "baseline_established": False,
            }
        self.patient_data[patient_id]["age"] = float(age_years)
        self.patient_data[patient_id]["sex"] = str(sex).lower()

    def update_and_get_clinical_state(
        self,
        patient_id: str,
        wheeze_prob: float,
        crackle_prob: float,
        breathing_rate: Optional[float] = None,
        timestamp: float = 0.0,
    ) -> Dict:
        base_state = self.update_and_get_state(patient_id, wheeze_prob, crackle_prob, timestamp)
        clinical = {}

        sound_assessment = self.clinical_kb.assess_adventitious_sounds(wheeze_prob, crackle_prob)
        clinical["adventitious_sounds"] = sound_assessment

        if breathing_rate is not None and breathing_rate > 0:
            pdata = self.patient_data.get(patient_id, {})
            if "breathing_rate_history" not in pdata:
                pdata["breathing_rate_history"] = []
            pdata["breathing_rate_history"].append(breathing_rate)

            age = pdata.get("age", self.patient_age if self.patient_age is not None else 40)
            sex = pdata.get("sex", self.patient_sex if self.patient_sex is not None else "male")
            rr_assessment = self.clinical_kb.assess_respiratory_rate(breathing_rate, age, sex)
            clinical["respiratory_rate"] = rr_assessment

        pdata = self.patient_data.get(patient_id, {})
        if pdata.get("breathing_rate_history"):
            base_state["breathing_rate_mean"] = float(statistics.mean(pdata["breathing_rate_history"]))
        elif breathing_rate is not None and breathing_rate > 0:
            base_state["breathing_rate_mean"] = breathing_rate
        else:
            base_state["breathing_rate_mean"] = None

        severity_rank = {"green": 0, "orange": 1, "red": 2}
        overall_sound = clinical.get("adventitious_sounds", {}).get("overall_severity", "green")
        max_severity = max(severity_rank.get(overall_sound, 0), severity_rank.get(clinical.get("respiratory_rate", {}).get("severity", "green"), 0))
        clinical["overall_clinical_status"] = {0: "green", 1: "orange", 2: "red"}.get(max_severity, "green")

        summary_parts = []
        if "adventitious_sounds" in clinical:
            ws = clinical["adventitious_sounds"]["wheeze"]["severity"]
            cs = clinical["adventitious_sounds"]["crackle"]["severity"]
            if ws != "green" or cs != "green":
                summary_parts.append(f"Adventitious sounds: wheeze={ws}, crackle={cs}")
        if "respiratory_rate" in clinical:
            rr = clinical["respiratory_rate"]
            summary_parts.append(f"Respiratory rate: {rr['value']:.1f} bpm ({rr['status']})")
        clinical["clinical_summary"] = " | ".join(summary_parts) if summary_parts else "No significant abnormalities detected"

        base_state["clinical_assessment"] = clinical
        
        overall_status = clinical.get("overall_clinical_status", "green")
        if overall_status == "red":
            base_state["comment"] = "RED - patient requires clinical review."
        elif overall_status == "orange":
            base_state["comment"] = "ORANGE - No immediate attention required."
        else:
            base_state["comment"] = "GREEN - Patient is normal."
        
        return base_state


class ClinicalAlertSystem:
    @staticmethod
    def generate_alerts(clinical_state: Dict) -> list:
        alerts = []
        rr_assessment = clinical_state.get("clinical_assessment", {}).get("respiratory_rate", {})
        if rr_assessment:
            severity = rr_assessment.get("severity", "green")
            if severity == "red":
                alerts.append({
                    "priority": 1,
                    "type": "CRITICAL",
                    "message": rr_assessment.get("clinical_action", ""),
                    "detail": f"Respiratory rate {rr_assessment.get('value', 0.0):.1f} bpm - {rr_assessment.get('status', '')}",
                })
            elif severity == "orange":
                alerts.append({
                    "priority": 2,
                    "type": "WARNING",
                    "message": rr_assessment.get("clinical_action", ""),
                    "detail": f"Respiratory rate {rr_assessment.get('value', 0.0):.1f} bpm",
                })

        sound = clinical_state.get("clinical_assessment", {}).get("adventitious_sounds", {})
        for sound_type in ["wheeze", "crackle"]:
            s = sound.get(sound_type, {})
            severity = s.get("severity", "green")
            if severity == "red":
                alerts.append({
                    "priority": 1,
                    "type": "CRITICAL",
                    "message": f"Severe {sound_type} detected - clinical review required",
                    "detail": s.get("status", ""),
                })
            elif severity == "orange":
                alerts.append({
                    "priority": 2,
                    "type": "WARNING",
                    "message": f"Clinically significant {sound_type} detected",
                    "detail": s.get("status", ""),
                })

        reason = clinical_state.get("reason", {})
        for axis in ["wheeze", "crackle"]:
            ax = reason.get(axis, {})
            if ax.get("state") == "red":
                alerts.append({
                    "priority": 1,
                    "type": "CRITICAL",
                    "message": f"{axis.capitalize()} probability significantly elevated from baseline",
                    "detail": f"Delta: {ax.get('delta', 0.0):+.3f}, Trend: {ax.get('trend', 0.0):+.3f}",
                })
            elif ax.get("state") == "orange":
                alerts.append({
                    "priority": 2,
                    "type": "WARNING",
                    "message": f"{axis.capitalize()} probability moderately elevated",
                    "detail": f"Delta: {ax.get('delta', 0.0):+.3f}",
                })

        alerts.sort(key=lambda x: x["priority"])
        return alerts

    @staticmethod
    def get_triage_recommendation(alerts: list) -> Dict:
        if any(a["priority"] == 1 for a in alerts):
            return {
                "level": "EMERGENCY",
                "action": "Immediate clinical evaluation required",
                "timeframe": "Within 30 minutes",
                "setting": "Emergency Department",
            }
        if any(a["priority"] == 2 for a in alerts):
            return {
                "level": "URGENT",
                "action": "Clinical assessment recommended",
                "timeframe": "Within 24 hours",
                "setting": "Urgent Care / Primary Care",
            }
        if alerts:
            return {
                "level": "ROUTINE",
                "action": "Monitor per standard protocol",
                "timeframe": "As scheduled",
                "setting": "Primary Care / Home monitoring",
            }
        return {
            "level": "NORMAL",
            "action": "Continue routine monitoring",
            "timeframe": "Per clinical protocol",
            "setting": "Home / Primary Care",
        }


def create_clinical_report(
    patient_id: str,
    age: float,
    sex: str,
    wheeze_prob: float,
    crackle_prob: float,
    breathing_rate: Optional[float] = None,
) -> Dict:
    manager = EnhancedPatientStateManager()
    manager.set_patient_demographics(patient_id, age, sex)
    state = manager.update_and_get_clinical_state(patient_id, wheeze_prob, crackle_prob, breathing_rate)
    alerts = ClinicalAlertSystem.generate_alerts(state)
    triage = ClinicalAlertSystem.get_triage_recommendation(alerts)
    return {
        "patient_id": patient_id,
        "demographics": {"age": age, "sex": sex},
        "clinical_state": state,
        "alerts": alerts,
        "triage_recommendation": triage,
    }


def add_clinical_reasoning_to_output(output_dict: Dict, patient_age: Optional[float] = None, patient_sex: Optional[str] = None) -> Dict:
    result = output_dict.get("result", {})
    wheeze_prob = result.get("wheeze", {}).get("probability")
    crackle_prob = result.get("crackle", {}).get("probability")
    audio_id = result.get("audio_id")

    if wheeze_prob is None or crackle_prob is None or audio_id is None:
        return output_dict

    if patient_age is None:
        patient_age = 40
    if patient_sex is None:
        patient_sex = "male"

    clinical_report = create_clinical_report(
        patient_id=audio_id,
        age=patient_age,
        sex=patient_sex,
        wheeze_prob=wheeze_prob,
        crackle_prob=crackle_prob,
        breathing_rate=result.get("breathing_rate_bpm"),
    )
    output_dict["clinical_report"] = clinical_report
    return output_dict