interface.py

import io
import json
import os
import sqlite3
import shutil
import threading
import time
import tempfile
import numpy as np
import soundfile as sf
from pathlib import Path
from itertools import permutations
from typing import Any, Dict, List, Optional

from reasoning_pipeline import aggregate_reasoning_summaries

MODEL_PATH = "best_audio_separation_model.pt"
_model = None


def get_model():
    global _model
    if _model is None:
        from model import load_model
        _model = load_model(MODEL_PATH)
    return _model


def apply_wiener_filter(est_sources, mixture, iterations=1):
    """
    A simplified Wiener filter refinement. 
    In source separation, this often refers to re-masking the mixture 
    based on the relative energy of the estimated sources.
    """
    # eps to avoid division by zero
    eps = 1e-10
    # Square the estimates to get power/variance proxies
    est_power = np.maximum(np.abs(est_sources)**2, eps) 
    total_power = np.sum(est_power, axis=0, keepdims=True) + eps
    
    # The Wiener gain is est_power / total_power
    # We apply this gain to the original mixture
    refined_sources = (est_power / total_power) * mixture
    return refined_sources


def save_waveform_plot(waveform, sample_rate, title, path=None):
    import matplotlib.pyplot as plt
    import torch
    if waveform.ndim > 1:
        waveform = waveform.mean(dim=0)
    data = waveform.cpu().numpy()
    times = np.arange(data.shape[-1]) / sample_rate
    fig, ax = plt.subplots(figsize=(10, 2.5))
    ax.plot(times, data, linewidth=0.7)
    ax.set_title(title)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Amplitude")
    ax.set_xlim(0, times[-1] if len(times) > 0 else 1)
    ax.grid(True)
    fig.tight_layout()
    if path is None:
        temp_image = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        fig.savefig(temp_image.name)
        plt.close(fig)
        return temp_image.name
    out_path = str(path)
    fig.savefig(out_path)
    plt.close(fig)
    return out_path


def predict_sources(waveform, sr, reference_waveforms=None):
    import torch
    import torchaudio

    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)

    if sr != 16000:
        waveform = torchaudio.functional.resample(
            waveform,
            sr,
            16000,
        )

    waveform = waveform.squeeze(0)
    with torch.no_grad():
        mixture = waveform.unsqueeze(0)
        prediction = get_model()(mixture)
        prediction = prediction.squeeze(0).permute(1, 0)
        
    # Convert to numpy for post-processing as per your evaluation snippet
    pred_np = prediction.cpu().numpy() # [3, T]
    mix_np = waveform.cpu().numpy()    # [T]

    # SepFormer's encoder -> masknet -> decoder round trip does not always
    # return exactly the same number of samples as the input (Conv1d
    # encode/decode with a fixed stride/kernel rounds down/up). Align both
    # to the shorter length before any elementwise op against the mixture,
    # otherwise the Wiener filter multiply below fails with a numpy
    # broadcast error like "(3,78760) (78763,)".
    min_len = min(pred_np.shape[-1], mix_np.shape[-1])
    pred_np = pred_np[:, :min_len]
    mix_np = mix_np[:min_len]

    # 1. Apply Wiener Filter
    refined_preds = apply_wiener_filter(pred_np, mix_np)

    # 2. Automated Permutation Alignment using Reference Audios
    if reference_waveforms is not None and len(reference_waveforms) > 0:
        best_perm = None
        min_mse = float('inf')
        num_sources = refined_preds.shape[0]

        for perm in permutations(range(num_sources)):
            current_mse = 0
            for i, p in enumerate(perm):
                if i < len(reference_waveforms) and reference_waveforms[i] is not None:
                    # Compare the first N samples to find the best match
                    ref = reference_waveforms[i]
                    length = min(refined_preds[p].shape[0], ref.shape[0])
                    current_mse += np.mean((ref[:length] - refined_preds[p][:length]) ** 2)
            
            if current_mse < min_mse:
                min_mse = current_mse
                best_perm = perm
        
        refined_preds = np.array([refined_preds[p] for p in best_perm])

    return torch.from_numpy(refined_preds.copy()), 16000


def save_separated_sources(prediction, output_dir, base_name, patient_names=None):
    import torch
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    patient_names = patient_names or ["Patient 1", "Patient 2", "Patient 3"]
    patient_names = [name if name else f"Patient {idx + 1}" for idx, name in enumerate(patient_names)]

    output_audio_paths = []
    output_image_paths = []

    for i in range(prediction.shape[0]):
        audio_name = f"{base_name}_s{i+1}.wav"
        audio_path = output_dir / audio_name
        save_audio_file(
            str(audio_path),
            prediction[i].unsqueeze(0),
            16000,
        )
        output_audio_paths.append(str(audio_path))

        image_name = f"{base_name}_s{i+1}_waveform.png"
        image_path = output_dir / image_name
        save_waveform_plot(
            prediction[i],
            16000,
            title=f"Separated waveform - {patient_names[i]}",
            path=image_path,
        )
        output_image_paths.append(str(image_path))

    return output_audio_paths, output_image_paths


def separate_audio_file(audio_path, output_dir, patient_names=None):
    waveform, sr = load_audio(audio_path)
    prediction, sample_rate = predict_sources(waveform, sr)
    return save_separated_sources(prediction, output_dir, Path(audio_path).stem, patient_names)


def load_audio(audio_path):
    import torch
    try:
        import torchaudio
        return torchaudio.load(audio_path)
    except (ImportError, RuntimeError, OSError):
        try:
            data, sr = sf.read(audio_path, dtype="float32")
            if data.ndim == 1:
                waveform = torch.from_numpy(data).unsqueeze(0)
            else:
                waveform = torch.from_numpy(data.T)
            return waveform, sr
        except (RuntimeError, ValueError, OSError):
            import librosa
            data, sr = librosa.load(audio_path, sr=None, mono=False)
            if data.ndim == 1:
                waveform = torch.from_numpy(data).unsqueeze(0)
            else:
                waveform = torch.from_numpy(data)
            return waveform, sr


def save_audio_file(path, waveform, sample_rate):
    import torch
    audio = waveform.detach().cpu().numpy()
    if audio.ndim == 1:
        audio = audio[np.newaxis, :]
    if audio.shape[0] > 1:
        audio = audio.T
    else:
        audio = audio[0]
    sf.write(path, audio, sample_rate) # save_audio_file now only saves the audio


RESULTS_DIR = Path("pipeline_results")
REASONING_SUMMARY_PATH = RESULTS_DIR / "reasoning_summary.json"
PATIENT_REGISTRY_PATH = RESULTS_DIR / "patient_registry.json"
HISTORY_RECORDS_PATH = RESULTS_DIR / "history_records.json"
DB_PATH = RESULTS_DIR / "speformer.db"
AUDIO_STORAGE_DIR = RESULTS_DIR / "audio"
GNN_RUN_ROOT = RESULTS_DIR / "gnn_runs"
DEFAULT_GNN_CHECKPOINT = "best_audio_separation_model.pt"


def _get_db_connection():
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL;")
    return conn


def initialize_database():
    conn = _get_db_connection()
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS patient_registry (
            id INTEGER PRIMARY KEY,
            patient_id TEXT UNIQUE,
            name TEXT,
            reference_audio TEXT,
            local_reference_audio TEXT,
            created_at TEXT
        )
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS history_records (
            id INTEGER PRIMARY KEY,
            timestamp TEXT,
            mix_audio TEXT,
            local_mix_audio TEXT,
            patient_names TEXT,
            separated_sources TEXT,
            run_dirs TEXT,
            reasoning_summary_path TEXT,
            reasoning_count INTEGER
        )
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS audio_files (
            id INTEGER PRIMARY KEY,
            path TEXT UNIQUE,
            file_type TEXT,
            patient_id TEXT,
            created_at TEXT,
            notes TEXT
        )
        """
    )
    conn.commit()
    conn.close()


def _record_audio_file_metadata(path: str, file_type: str, patient_id: Optional[str] = None, notes: Optional[str] = None):
    conn = None
    try:
        initialize_database()
        conn = _get_db_connection()
        conn.execute(
            "INSERT OR IGNORE INTO audio_files (path, file_type, patient_id, created_at, notes) VALUES (?, ?, ?, ?, ?)",
            (str(Path(path).resolve()), file_type, patient_id, time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), notes),
        )
        conn.commit()
    except Exception:
        pass
    finally:
        if conn:
            conn.close()


def _copy_audio_to_storage(src_path: Optional[str], subdir: str, prefix: str) -> Optional[str]:
    if not src_path:
        return None
    src = Path(src_path)
    try:
        if not src.exists():
            return None
        src_resolved = src.resolve()
        if RESULTS_DIR in src_resolved.parents or src_resolved == RESULTS_DIR.resolve():
            return str(src_resolved)
    except Exception:
        pass

    storage_dir = AUDIO_STORAGE_DIR / subdir
    storage_dir.mkdir(parents=True, exist_ok=True)
    timestamp = time.strftime("%Y%m%d%H%M%S", time.gmtime())
    dest = storage_dir / f"{prefix}_{timestamp}_{src.name}"
    shutil.copy2(src, dest)
    _record_audio_file_metadata(str(dest), subdir, prefix)
    return str(dest)


def save_patient_registry(patient_entries: List[Dict[str, str]], registry_path: Optional[Path] = None) -> Path:
    path = Path(registry_path) if registry_path is not None else PATIENT_REGISTRY_PATH
    path.parent.mkdir(parents=True, exist_ok=True)

    initialize_database()
    conn = _get_db_connection()
    local_entries = []
    for entry in patient_entries:
        patient_id = entry.get("patient_id") or entry.get("id") or ""
        name = entry.get("name") or ""
        ref_audio = entry.get("reference_audio")
        local_ref_audio = _copy_audio_to_storage(ref_audio, "patient_reference", patient_id) if ref_audio else None
        conn.execute(
            "INSERT OR REPLACE INTO patient_registry (patient_id, name, reference_audio, local_reference_audio, created_at) VALUES (?, ?, ?, ?, ?)",
            (
                patient_id,
                name,
                str(ref_audio) if ref_audio else None,
                local_ref_audio,
                time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            ),
        )
        local_entries.append({
            "patient_id": patient_id,
            "name": name,
            "reference_audio": local_ref_audio or (str(ref_audio) if ref_audio else None),
        })
    conn.commit()
    conn.close()

    with open(path, "w", encoding="utf-8") as f:
        json.dump(local_entries, f, indent=2)
    return path


def load_patient_registry(registry_path: Optional[Path] = None) -> List[Dict[str, str]]:
    if registry_path is None and DB_PATH.exists():
        try:
            initialize_database()
            conn = _get_db_connection()
            cursor = conn.execute("SELECT patient_id, name, local_reference_audio AS reference_audio FROM patient_registry ORDER BY id")
            rows = cursor.fetchall()
            conn.close()
            if rows:
                return [{"patient_id": row["patient_id"], "name": row["name"], "reference_audio": row["reference_audio"]} for row in rows]
        except Exception:
            pass

    path = Path(registry_path) if registry_path is not None else PATIENT_REGISTRY_PATH
    if not path.exists():
        return []
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                return data
    except Exception:
        pass
    return []


def resolve_patient_names(patient_names=None, registry_path: Optional[Path] = None):
    registry = load_patient_registry(registry_path)
    resolved = []
    patient_names = patient_names or []
    for idx in range(3):
        name = None
        if idx < len(patient_names) and patient_names[idx]:
            name = patient_names[idx]
        elif idx < len(registry) and registry[idx].get("name"):
            name = registry[idx]["name"]
        if not name:
            name = f"Patient {idx + 1}"
        resolved.append(name)
    return resolved


def load_reasoning_summary(results_path: Optional[Path] = None):
    path = Path(results_path) if results_path is not None else REASONING_SUMMARY_PATH
    if not path.exists():
        return []
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return []


def load_history_records(history_path: Optional[Path] = None) -> List[Dict[str, str]]:
    if DB_PATH.exists():
        try:
            initialize_database()
            conn = _get_db_connection()
            cursor = conn.execute("SELECT * FROM history_records ORDER BY id")
            rows = cursor.fetchall()
            conn.close()
            history = []
            for row in rows:
                local_mix_audio = row["local_mix_audio"] or row["mix_audio"]
                history.append({
                    "timestamp": row["timestamp"],
                    "mix_audio": local_mix_audio,
                    "patient_names": json.loads(row["patient_names"] or "[]"),
                    "separated_sources": json.loads(row["separated_sources"] or "[]"),
                    "run_dirs": json.loads(row["run_dirs"] or "[]"),
                    "reasoning_summary_path": row["reasoning_summary_path"],
                    "reasoning_count": row["reasoning_count"],
                    "local_mix_audio": row["local_mix_audio"],
                })
            return history
        except Exception:
            pass
    path = Path(history_path) if history_path is not None else HISTORY_RECORDS_PATH
    if not path.exists():
        return []
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                return data
    except Exception:
        pass
    return []


def append_history_record(record: Dict[str, object], history_path: Optional[Path] = None) -> Path:
    path = Path(history_path) if history_path is not None else HISTORY_RECORDS_PATH
    local_mix_audio = _copy_audio_to_storage(record.get("mix_audio"), "mix_audio", "mix_audio") if record.get("mix_audio") else None
    if local_mix_audio:
        record["mix_audio"] = local_mix_audio

    initialize_database()
    conn = _get_db_connection()
    conn.execute(
        "INSERT INTO history_records (timestamp, mix_audio, local_mix_audio, patient_names, separated_sources, run_dirs, reasoning_summary_path, reasoning_count) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
        (
            record.get("timestamp"),
            str(record.get("mix_audio")) if record.get("mix_audio") else None,
            local_mix_audio,
            json.dumps(record.get("patient_names") or []),
            json.dumps(record.get("separated_sources") or []),
            json.dumps(record.get("run_dirs") or []),
            str(record.get("reasoning_summary_path")) if record.get("reasoning_summary_path") else None,
            int(record.get("reasoning_count") or 0),
        ),
    )
    conn.commit()
    conn.close()

    history = load_history_records(path)
    history.append(record)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)
    return path


def load_wav2vec(device):
    import torch
    import importlib

    try:
        transformers = importlib.import_module("transformers")
    except ModuleNotFoundError as exc:
        raise ModuleNotFoundError(
            "Required package 'transformers' is not installed. Please install it with `pip install transformers` and restart the app."
        ) from exc

    Wav2Vec2Model = getattr(transformers, "Wav2Vec2Model")
    Wav2Vec2Processor = getattr(transformers, "Wav2Vec2Processor")

    processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
    model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")
    model = model.to(device)
    model.eval()
    return processor, model


_live_sr = 16000
LIVE_PROCESSING_WINDOW_SECONDS = 5.0
LIVE_OVERLAP_SECONDS = 1.0

_live_processor = None
_live_wav2vec_model = None
_live_gnn_model = None
_live_device = None
_live_patient_manager = None

# ---------------------------------------------------------------------------
# Per-bed reference waveforms — this is what keeps "channel i" pinned to the
# SAME physical bed across chunks. SepFormer/blind separation does not
# guarantee stable output ordering run to run; predict_sources()'s
# permutation-alignment step (see above) reorders its 3 output channels to
# best match whatever reference waveforms it's given. So: each bed gets a
# reference, which is what gets passed in, and gets rolled forward slightly
# with every new chunk so it can track slow drift (mic gain, posture)
# without letting one noisy chunk yank it onto a different source.
# ---------------------------------------------------------------------------
_bed_reference_waveforms: Dict[str, np.ndarray] = {}
_bed_reference_lock = threading.Lock()


def _initialize_models_for_live_processing():
    import torch
    global _live_processor, _live_wav2vec_model, _live_gnn_model, _live_device
    if _live_device is None:
        _live_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if _live_processor is None or _live_wav2vec_model is None:
        _live_processor, _live_wav2vec_model = load_wav2vec(_live_device)
    if _live_gnn_model is None:
        from gnn import load_gnn_model
        _live_gnn_model = load_gnn_model(DEFAULT_GNN_CHECKPOINT, _live_device)
    return _live_processor, _live_wav2vec_model, _live_gnn_model, _live_device


def _get_live_patient_manager():
    global _live_patient_manager
    if _live_patient_manager is None:
        from gnn import PatientStateManager
        _live_patient_manager = PatientStateManager(
            ema_alpha=0.12,
            low_delta=0.08,
            high_delta=0.20,
            min_samples_for_baseline=5,
            force_established_after_s=10.0,
        )
    return _live_patient_manager


def reset_bed_reference(bed_id: str) -> None:
    """Call this whenever a bed's occupant changes (discharge / new patient
    registered), so the previous patient's acoustic reference doesn't get
    used to (mis)align the new patient's channel."""
    with _bed_reference_lock:
        _bed_reference_waveforms.pop(bed_id, None)


def process_audio_chunk_for_separation(audio_chunk_tensor):
    """Expects a 1D tensor at 16kHz."""
    import torch
    prediction, _ = predict_sources(audio_chunk_tensor.unsqueeze(0), 16000)
    return prediction


def infer_on_separated_chunk(
    separated_chunk_np,
    gnn_model,
    processor,
    wav2vec_model,
    device,
    patient_manager,
    patient_id,
    timestamp,
):
    """Runs behavior inference and clinical state tracking on a single audio chunk."""
    import torch
    inputs = processor(separated_chunk_np, sampling_rate=16000, return_tensors="pt", padding=True)
    input_values = inputs.input_values.to(device)

    with torch.no_grad():
        out = wav2vec_model(input_values)
        emb = out.last_hidden_state.mean(dim=1)
        from gnn import build_chain_edge_index
        from torch_geometric.data import Data
        edge_index = build_chain_edge_index(1).to(device)
        data = Data(x=emb, edge_index=edge_index)
        w_logits, c_logits = gnn_model(data)
        w_prob = float(torch.sigmoid(w_logits).view(-1)[0].item())
        c_prob = float(torch.sigmoid(c_logits).view(-1)[0].item())

    from gnn import estimate_breathing_rate_bpm
    br_bpm = estimate_breathing_rate_bpm(separated_chunk_np, 16000, len(separated_chunk_np)/16000)

    clinical_state = patient_manager.update_and_get_clinical_state(
        patient_id, w_prob, c_prob, breathing_rate=br_bpm, timestamp=timestamp
    )
    # Make sure the raw probabilities/rate are always available to callers,
    # even if PatientStateManager's own dict doesn't include them.
    if isinstance(clinical_state, dict):
        clinical_state.setdefault("wheeze_prob", w_prob)
        clinical_state.setdefault("crackle_prob", c_prob)
        clinical_state.setdefault("breathing_rate", br_bpm)
    return clinical_state


def analyze_live_chunk_for_beds(
    audio_bytes: bytes,
    bed_ids: List[str],
    timestamp: Optional[float] = None,
) -> Dict[str, Dict[str, Any]]:
    """
    Runs ONE separation pass on a single mic audio chunk and returns a
    per-bed analysis dict, e.g. {"Bed 1": {...}, "Bed 2": {...}, "Bed 3": {...}}.

    This replaces calling the old single-bed analyze endpoint once per bed
    with the SAME audio chunk. That approach ran three independent,
    unaligned 3-way separations per chunk — each call had no knowledge of
    the other two, so there was no guarantee "channel 1" from call A was
    the same physical source as "channel 1" from call B. Here, separation
    runs exactly once, and each bed's channel is aligned to that bed's own
    persisted reference waveform (see _bed_reference_waveforms above), so
    identity stays pinned to the correct bed across chunks.

    bed_ids: the beds this mic covers (e.g. its 3 nearest), length <= 3.
    Order only matters for the very first chunk of a session, before any
    bed has a reference yet — after that, alignment is reference-driven.
    """
    if timestamp is None:
        timestamp = time.time()

    bed_ids = [b for b in bed_ids if b][:3]
    if not bed_ids:
        return {}

    import torch
    import torchaudio

    # ---- load + resample the chunk ----
    data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
    if data.ndim > 1:
        data = data.mean(axis=1)
    waveform = torch.from_numpy(data).unsqueeze(0)  # [1, T]
    if sr != _live_sr:
        waveform = torchaudio.functional.resample(waveform, sr, _live_sr)
    waveform = waveform.squeeze(0)  # [T]

    # ---- pull each bed's known reference (None if not enrolled yet) ----
    with _bed_reference_lock:
        refs = [_bed_reference_waveforms.get(bid) for bid in bed_ids]

    prediction, _ = predict_sources(waveform.unsqueeze(0), _live_sr, reference_waveforms=refs)
    pred_np = prediction.cpu().numpy()  # [num_channels, T], channel i ~ bed_ids[i]
    num_channels = pred_np.shape[0]

    # ---- roll each bed's reference forward so future chunks keep aligning to it ----
    with _bed_reference_lock:
        for i, bid in enumerate(bed_ids):
            if i >= num_channels:
                continue
            new_channel = pred_np[i]
            old = _bed_reference_waveforms.get(bid)
            if old is None:
                _bed_reference_waveforms[bid] = new_channel.copy()
            else:
                length = min(len(old), len(new_channel))
                _bed_reference_waveforms[bid] = (
                    0.7 * old[:length] + 0.3 * new_channel[:length]
                ).astype(np.float32)

    # ---- run GNN behavior inference per bed on its aligned channel ----
    processor, wav2vec_model, gnn_model, device = _initialize_models_for_live_processing()
    patient_manager = _get_live_patient_manager()

    results: Dict[str, Dict[str, Any]] = {}
    for i, bid in enumerate(bed_ids):
        if i >= num_channels:
            results[bid] = {"error": "No separated channel available for this bed."}
            continue
        channel = pred_np[i]
        try:
            clinical_state = infer_on_separated_chunk(
                channel, gnn_model, processor, wav2vec_model, device,
                patient_manager, patient_id=bid, timestamp=timestamp,
            )
        except Exception as exc:
            results[bid] = {"error": str(exc)}
            continue

        cs = clinical_state if isinstance(clinical_state, dict) else {}
        w_prob = float(cs.get("wheeze_prob", 0.0))
        c_prob = float(cs.get("crackle_prob", 0.0))
        severity_score = float(cs.get("severity", max(w_prob, c_prob)))
        breathing_rate = cs.get("breathing_rate")

        results[bid] = {
            "bed_id": bid,
            "respiratory_rate": breathing_rate,
            "severity_score": severity_score,
            "wheeze_probability": w_prob,
            "crackle_probability": c_prob,
            "patient_state": cs.get("overall_state"),
            "breathing_pattern": (
                "Rapid & Shallow" if severity_score >= 0.66
                else "Slightly Laboured" if severity_score >= 0.33
                else "Regular"
            ),
            "condition": (
                "Respiratory Distress" if severity_score >= 0.66
                else "Elevated Rate" if severity_score >= 0.33
                else "Stable"
            ),
        }

    return results


def run_end_to_end(
    mix_audio_path,
    patient_names=None,
    reference_audio_paths=None, # New argument to accept reference audio paths
    gnn_checkpoint: Optional[str] = None,
    device_str: Optional[str] = None,
):
    import torch
    if not mix_audio_path:
        raise ValueError("No mixture audio file provided.")

    # Note: reference_audio_paths are passed but not used by the current separate_audio/predict_sources
    patient_names = resolve_patient_names(patient_names)
    outputs = separate_audio(
        mix_audio_path,
        patient_names=patient_names,
        reference_audio_paths=reference_audio_paths, # Pass reference audio paths
    )
    separated_audio_paths = outputs[:3]

    device_str = device_str or ("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device(device_str)
    processor, wav2vec_model = load_wav2vec(device)

    from gnn import load_gnn_model

    gnn_model = load_gnn_model(gnn_checkpoint or DEFAULT_GNN_CHECKPOINT, device)

    GNN_RUN_ROOT.mkdir(parents=True, exist_ok=True)
    behavior_results = []
    for audio_path in separated_audio_paths:
        from full_pipeline import infer_on_audio_file

        result = infer_on_audio_file(
            Path(audio_path),
            gnn_model,
            processor,
            wav2vec_model,
            GNN_RUN_ROOT,
            device,
        )
        if result is not None:
            behavior_results.append(result)

    run_dirs = [Path(item["artifacts"]["run_dir"]) for item in behavior_results if item.get("artifacts")]

    # Construct metadata mapping the audio_id (filename stem) to the selected patient name for human-readable reporting
    patient_meta = {}
    for i, audio_path in enumerate(separated_audio_paths):
        if audio_path:
            stem = Path(audio_path).stem
            if i < len(patient_names):
                patient_meta[stem] = {"name": patient_names[i]}

    reasoning_summaries = aggregate_reasoning_summaries(run_dirs, RESULTS_DIR, patient_meta=patient_meta)

    history_record = {
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "mix_audio": str(mix_audio_path),
        "patient_names": patient_names,
        "separated_sources": separated_audio_paths,
        "run_dirs": [str(path) for path in run_dirs],
        "reasoning_summary_path": str(REASONING_SUMMARY_PATH),
        "reasoning_count": len(reasoning_summaries),
    }
    append_history_record(history_record)

    return outputs, reasoning_summaries, history_record


def monitoring_table_rows(results_path: Optional[Path] = None):
    records = load_reasoning_summary(results_path)
    rows = []
    for item in records:
        rows.append([
            item.get("audio_id"),
            item.get("overall_state"),
            item.get("mean_wheeze_prob"),
            item.get("mean_crackle_prob"),
            item.get("breathing_rate_mean"),
            item.get("comment"),
        ])
    return rows


def search_history_records(query: str, history_path: Optional[Any] = None):
    if not query:
        return []
    query_lower = query.strip().lower()
    records = load_history_records(history_path)
    results = []
    for item in records:
        patient_names = item.get("patient_names") or []
        mix_audio = str(item.get("mix_audio", ""))
        timestamp = str(item.get("timestamp", ""))
        if (
            query_lower in mix_audio.lower()
            or query_lower in timestamp.lower()
            or any(query_lower in str(name).lower() for name in patient_names)
        ):
            results.append([
                item.get("timestamp"),
                mix_audio,
                ", ".join([str(name) for name in patient_names if name]),
                item.get("reasoning_count"),
            ])
    return results


def search_reasoning_records(query: str, results_path: Optional[Any] = None):
    if not query:
        return []
    query_lower = query.strip().lower()
    records = load_reasoning_summary(results_path)
    results = []
    for item in records:
        audio_id = str(item.get("audio_id", ""))
        if query_lower in audio_id.lower():
            results.append([
                item.get("audio_id"),
                item.get("overall_state"),
                item.get("mean_wheeze_prob"),
                item.get("mean_crackle_prob"),
                item.get("breathing_rate_mean"),
                item.get("comment"),
            ])
    return results


def separate_audio(audio_path, patient_names=None, reference_audio_paths=None): # New argument

    import torch
    import torchaudio
    waveform, sr = load_audio(audio_path)
    
    # Load reference waveforms if provided for alignment
    ref_waveforms = []
    if reference_audio_paths:
        for path in reference_audio_paths:
            if path and Path(path).exists():
                ref_wav, _ = load_audio(path)
                ref_waveforms.append(ref_wav.mean(dim=0).cpu().numpy())
            else:
                ref_waveforms.append(None)

    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)

    if sr != 16000:
        waveform = torchaudio.functional.resample(
            waveform,
            sr,
            16000,
        )

    waveform = waveform.squeeze(0)

    with torch.no_grad():
        # Now passing the references to predict_sources for alignment logic
        prediction, _ = predict_sources(waveform.unsqueeze(0), 16000, reference_waveforms=ref_waveforms)

    patient_names = resolve_patient_names(patient_names)
    output_dir = AUDIO_STORAGE_DIR / "separated" / Path(audio_path).stem
    output_audio_paths, output_image_paths = save_separated_sources(
        prediction,
        output_dir,
        Path(audio_path).stem,
        patient_names=patient_names,
    )

    while len(output_audio_paths) < 3:
        output_audio_paths.append(None)
    while len(output_image_paths) < 3:
        output_image_paths.append(None)

    return output_audio_paths + output_image_paths